From 55924ce408cf0369ed2e0e943ae8a03ecd46322a Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 9 Oct 2026 18:34:41 +0200 Subject: [PATCH 01/88] gridpack: in-memory event streaming API (bin/event_stream.py) for fix_max_weight gridpacks New madspace EventStream: a coordinator thread keeps generation jobs running, buffers each channel's unweighted events in launch order and picks the channel of every event from the current cross-section estimates. Reproducible from the seed, nothing written to disk. Shared helpers factored out of EventGenerator (event_utils) and gridpack.py (gridpack_setup.py). --- madgraph/interface/tutorials/mg7.py | 3 +- .../iolibs/template_files/mg7/event_stream.py | 276 +++++++++++ .../iolibs/template_files/mg7/gridpack.py | 285 ++--------- .../template_files/mg7/gridpack_setup.py | 289 +++++++++++ madgraph/iolibs/template_files/mg7/launch.py | 9 +- .../iolibs/template_files/mg7/run_card.toml | 4 +- madspace/CMakeLists.txt | 2 + madspace/include/madspace/driver.hpp | 2 + .../madspace/driver/channel_generator.hpp | 31 +- .../madspace/driver/event_generator.hpp | 7 - .../include/madspace/driver/event_stream.hpp | 167 +++++++ .../include/madspace/driver/event_utils.hpp | 70 +++ madspace/src/driver/channel_generator.cpp | 96 ++-- madspace/src/driver/event_generator.cpp | 194 +------- madspace/src/driver/event_stream.cpp | 466 ++++++++++++++++++ madspace/src/driver/event_utils.cpp | 200 ++++++++ madspace/src/python/madspace.cpp | 136 ++++- madspace/tests/test_event_stream.py | 241 +++++++++ 18 files changed, 2001 insertions(+), 477 deletions(-) create mode 100644 madgraph/iolibs/template_files/mg7/event_stream.py create mode 100644 madgraph/iolibs/template_files/mg7/gridpack_setup.py create mode 100644 madspace/include/madspace/driver/event_stream.hpp create mode 100644 madspace/include/madspace/driver/event_utils.hpp create mode 100644 madspace/src/driver/event_stream.cpp create mode 100644 madspace/src/driver/event_utils.cpp create mode 100644 madspace/tests/test_event_stream.py diff --git a/madgraph/interface/tutorials/mg7.py b/madgraph/interface/tutorials/mg7.py index f5eaf7f384..d26a207d82 100644 --- a/madgraph/interface/tutorials/mg7.py +++ b/madgraph/interface/tutorials/mg7.py @@ -170,7 +170,8 @@ maps the old names onto the new sections. **Gridpacks** work here too, via the `[gridpack]` section and -`bin/gridpack.py`. +`bin/generate_events`. With `run_mode = "fix_max_weight"`, the gridpack also +ships `bin/event_stream.py`, a Python API that streams events in memory. %(see_also)s diff --git a/madgraph/iolibs/template_files/mg7/event_stream.py b/madgraph/iolibs/template_files/mg7/event_stream.py new file mode 100644 index 0000000000..1e5e843ad7 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/event_stream.py @@ -0,0 +1,276 @@ +"""In-memory event generation from a gridpack, as a Python API. + +Only available in gridpacks made with ``[gridpack] run_mode = "fix_max_weight"``: +their channels carry a fixed maximum weight, so every event is final as soon as +it is unweighted and the events can be streamed without a survey or a combine +pass. Nothing is written to disk. + +Usage:: + + import sys + sys.path.insert(0, "/bin") + from event_stream import EventStream + + with EventStream(device="cpu", seed=42) as stream: + for event in stream: # ms.LHEEvent objects, endless + ... + for batch in stream.batches(10000, output_type="torch"): + batch.px # shape (10000, max_particle_count) + +Event weights are in units of the maximum weight: +-1, except for overweight +events. To normalize a sample to the cross section, scale its weights by +``stream.integral / sum(weights)``. The event sequence is reproducible from the seed, +independent of the thread count and of how it is split into batches. +""" + +import dataclasses +import json +import sys +from typing import Any, Iterator, Optional + +import numpy as np + +from gridpack_setup import ( + SOURCE_HASH_MESSAGE, + build_lhe_meta, + gridpack_path, + load_channels, + load_context_data, + load_lhe_completer, + load_madspace_data, + load_run_card, + make_context, + resolve_seed, + source_hash_matches, +) +import madspace as ms + +__all__ = ["EventStream", "EventBatch"] + +# events completed per call into madspace while iterating over single events +_ITER_CHUNK = 1000 + + +@dataclasses.dataclass +class EventBatch: + """A batch of LHE events as a structure of arrays: event fields have shape + (batch_size,), particle fields (batch_size, max_particle_count). Particle + slots beyond ``particle_count`` are padded with zeros, see + ``particle_mask``. numpy arrays or CPU torch tensors, depending on the + ``output_type`` of :meth:`EventStream.batches`.""" + + process_id: Any + weight: Any + scale: Any + alpha_qed: Any + alpha_qcd: Any + particle_count: Any + particle_mask: Any + pdg_id: Any + status_code: Any + mother1: Any + mother2: Any + color: Any + anti_color: Any + px: Any + py: Any + pz: Any + energy: Any + mass: Any + lifetime: Any + spin: Any + + def __len__(self) -> int: + return len(self.weight) + + +class EventStream: + """Endless stream of unweighted events from this gridpack. + + Constructing it starts the generation in background threads, which keep + a buffer of events for every channel. The channel of every event is picked + at random according to the current cross-section estimates, which are + refined as events are generated. + + Arguments left as None take their value from Cards/grid_run_card.toml. + + Args: + device: device to generate on, like "cpu", "cuda" or "cuda:1". + seed: run seed; -1 draws a random one (see :attr:`seed`). + cpu_batch_size: generation batch size on a CPU device. + gpu_batch_size: generation batch size on a GPU device. + thread_pool_size: number of worker threads. + cut_efficiency_threshold: fraction of a batch that has to pass the cuts. + max_cut_repetitions: how often a batch is repeated to pass the cuts. + ignore_source_hash: run even if madspace differs from the version the + gridpack was made with (can lead to errors or incorrect results). + """ + + def __init__( + self, + device: Optional[str] = None, + seed: Optional[int] = None, + cpu_batch_size: Optional[int] = None, + gpu_batch_size: Optional[int] = None, + thread_pool_size: Optional[int] = None, + cut_efficiency_threshold: Optional[float] = None, + max_cut_repetitions: Optional[int] = None, + ignore_source_hash: bool = False, + ): + self._stream = None + run_card = load_run_card() + run_args = run_card["run"] + gen_args = run_card["generation"] + madspace_data = load_madspace_data() + if not source_hash_matches(madspace_data): + if not ignore_source_hash: + raise RuntimeError( + f"{SOURCE_HASH_MESSAGE}. Pass ignore_source_hash=True to run anyway." + ) + print(f"WARNING: {SOURCE_HASH_MESSAGE}", file=sys.stderr) + + with open(gridpack_path("data", "channels.json")) as f: + channels = json.load(f) + if not any("max_weight" in channel for channel in channels.values()): + raise RuntimeError( + "EventStream needs a gridpack made with [gridpack] run_mode = " + '"fix_max_weight"' + ) + + def pick(value, default): + return default if value is None else value + + if device is None: + devices = run_args["device"] + device = devices[0] if isinstance(devices, list) else devices + if thread_pool_size is None: + is_gpu = device.split(":")[0] in ("cuda", "hip") + thread_pool_size = run_args[ + "gpu_thread_pool_size" if is_gpu else "cpu_thread_pool_size" + ] + self._context, backend = make_context( + device, run_args["cpu_mode"], thread_pool_size, thread_pool_size + ) + + config = ms.GeneratorConfig() + config.cpu_batch_size = pick(cpu_batch_size, gen_args["cpu_batch_size"]) + config.gpu_batch_size = pick(gpu_batch_size, gen_args["gpu_batch_size"]) + config.cut_efficiency_threshold = pick( + cut_efficiency_threshold, gen_args["cut_efficiency_threshold"] + ) + config.max_cut_repetitions = pick( + max_cut_repetitions, gen_args["max_cut_repetitions"] + ) + + self.seed = resolve_seed(pick(seed, run_args.get("seed", -1))) + load_context_data([self._context], [backend], madspace_data) + channel_generators = load_channels([self._context], config) + self._stream = ms.EventStream( + self._context, channel_generators, self.seed, load_lhe_completer(), config + ) + + def __iter__(self) -> Iterator["ms.LHEEvent"]: + """Endless iterator over the events as ms.LHEEvent objects. It takes + events from the stream in chunks of 1000; those of the last chunk that + are not iterated over are dropped.""" + while True: + yield from self._stream.next_events(_ITER_CHUNK) + + def next_events(self, count: int) -> list: + """The next ``count`` events as ms.LHEEvent objects.""" + return self._stream.next_events(count) + + def next_batch(self, batch_size: int, output_type: str = "numpy") -> EventBatch: + """The next ``batch_size`` events as an :class:`EventBatch`.""" + if output_type not in ("numpy", "torch"): + raise ValueError('output_type must be "numpy" or "torch"') + arrays = self._stream.next_batch(batch_size) + counts = arrays["particle_count"] + max_count = arrays["pdg_id"].shape[1] + arrays["particle_mask"] = np.arange(max_count)[None, :] < counts[:, None] + if output_type == "torch": + import torch + arrays = {key: torch.from_numpy(value) for key, value in arrays.items()} + return EventBatch(**arrays) + + def batches( + self, + batch_size: int, + output_type: str = "numpy", + max_events: Optional[int] = None, + ) -> Iterator[EventBatch]: + """Iterator over batches of ``batch_size`` events, endless unless + ``max_events`` is given (the last batch may then be smaller).""" + remaining = max_events + while remaining is None or remaining > 0: + size = batch_size if remaining is None else min(batch_size, remaining) + yield self.next_batch(size, output_type) + if remaining is not None: + remaining -= size + + @property + def integral(self) -> float: + """Current estimate of the cross section, in pb.""" + return self._stream.status().mean + + @property + def integral_error(self) -> float: + """Standard error of :attr:`integral`.""" + return self._stream.status().error + + @property + def abs_integral(self) -> float: + """Current estimate of the integral of the absolute weights.""" + return self._stream.status().mean_abs + + @property + def n_weighted(self) -> int: + """Number of weighted events the estimates are based on.""" + return self._stream.status().count + + @property + def n_weighted_after_cuts(self) -> int: + """Number of those weighted events that passed the cuts.""" + return self._stream.status().count_after_cuts + + @property + def n_events(self) -> int: + """Number of events returned so far.""" + return self._stream.event_count() + + @property + def max_particle_count(self) -> int: + """Largest number of particles per event, including resonances.""" + return self._stream.max_particle_count() + + def status(self): + """Combined ms.GeneratorStatus over all channels.""" + return self._stream.status() + + def channel_status(self) -> list: + """ms.GeneratorStatus of every channel; ``count_unweighted`` is the + number of its events returned so far.""" + return self._stream.channel_status() + + def channel_probabilities(self) -> list: + """Current probability of every channel to be picked for an event.""" + return self._stream.channel_probabilities() + + def lhe_meta(self): + """ms.LHEMeta for writing the events to an LHE file, with the current + cross section estimate.""" + return build_lhe_meta(self._stream.status(), self.seed) + + def close(self) -> None: + """Stop the background generation.""" + if self._stream is not None: + self._stream.close() + + def __enter__(self) -> "EventStream": + return self + + def __exit__(self, *exc) -> None: + self.close() + + def __del__(self) -> None: + self.close() diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 848799a7d6..8b6937010c 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -1,37 +1,29 @@ #! /usr/bin/env python3 +import argparse +import glob import gzip +import json import os import shutil -import subprocess import sys -from pathlib import Path - -# search order for madspace package: -# 1. a precompiled install (madspace/install/madspace) -# 2. bundled source that still needs to be built (madspace/install.py) -# 3. otherwise fall back to madspace is available in the environment -_GRIDPACK_DIR = Path(os.path.realpath(__file__)).parent.parent -_LOCAL_MADSPACE_DIR = _GRIDPACK_DIR / "madspace" -_LOCAL_INSTALL_DIR = _LOCAL_MADSPACE_DIR / "install" -if (_LOCAL_INSTALL_DIR / "madspace").is_dir(): - sys.path.insert(0, str(_LOCAL_INSTALL_DIR)) -elif (_LOCAL_MADSPACE_DIR / "install.py").is_file(): - print() - print("You don't have madspace installed for this gridpack") - print("Running interactive madspace installation script") - print() - - _result = subprocess.run([sys.executable, str(_LOCAL_MADSPACE_DIR / "install.py")]) - if _result.returncode != 0: - raise RuntimeError("madspace installation failed — see output above") - sys.path.insert(0, str(_LOCAL_INSTALL_DIR)) +# bin/gridpack_setup.py, next to this script: also locates madspace +from gridpack_setup import ( + SOURCE_HASH_MESSAGE, + build_lhe_meta, + load_channels, + load_context_data, + load_lhe_completer, + load_madspace_data, + load_run_card, + load_systematics, + make_context, + resolve_seed, + source_hash_matches, +) import madspace as ms -import glob -import json -import tomllib -import argparse + def resolve_verbosity(verbosity: str) -> str: """Resolve the run_card "auto" verbosity to "pretty"/"log" depending on @@ -42,43 +34,6 @@ def resolve_verbosity(verbosity: str) -> str: return verbosity -def resolve_seed(seed: int) -> int: - """Resolve the run_card "seed": -1 draws a fresh 64-bit seed via - os.urandom, any other value is used as-is.""" - if seed == -1: - return int.from_bytes(os.urandom(8), "big") - return seed - - -def build_lhe_meta(event_generator, seed: int, systematics=None): - """The LHE header/ metadata of this gridpack run: the cards, beams - and PDF the gridpack was made with (data/lhe_meta.json), with the cross - section and seed of this run. ``systematics`` adds the block, - for the writers that do not inject it themselves (header.lhe); leave it - unset for combine_to_lhe, which does.""" - headers = [] - data = {} - meta_path = os.path.join("data", "lhe_meta.json") - if os.path.exists(meta_path): - with open(meta_path) as f: - data = json.load(f) - headers = [ms.LHEHeader(name=h["name"], content=h["content"], - escape_content=h["escape_content"]) - for h in data.pop("headers")] - headers.append(ms.LHEHeader(name="MG7Seed", content=str(seed))) - if systematics is not None and systematics.weight_ids: - headers.append(ms.LHEHeader( - name="initrwgt", content=systematics.initrwgt(), escape_content=False - )) - status = event_generator.status() - return ms.LHEMeta( - # positional: the pybind arg name for max_weight is non-kwarg-safe - processes=[ms.LHEProcess(status.mean, status.error, status.mean, 1)], - headers=headers, - **data, - ) - - def write_lhe_header(run_path: str, meta) -> None: """Write header.lhe next to events.npy: the
/ blocks the npy formats otherwise drop, with no events.""" @@ -91,22 +46,11 @@ def write_lhe_header(run_path: str, meta) -> None: def main() -> None: - # load run card and metadata. Use the RunCardMG7 representation when the - # madgraph package is importable; gridpacks are meant to be portable, so - # fall back to a plain tomllib parse otherwise (the card is the same TOML). - run_card_path = os.path.join("Cards", "grid_run_card.toml") - try: - from madgraph.various.banner import RunCardMG7 - run_card = RunCardMG7(run_card_path) - except ImportError: - with open(run_card_path, "rb") as f: - run_card = tomllib.load(f) + run_card = load_run_card() run_args = run_card["run"] gen_args = run_card["generation"] grid_args = run_card.get("gridpack", {}) - param_card_path = os.path.join("Cards", "param_card.dat") - with open(os.path.join("data", "data.json")) as f: - madspace_data = json.load(f) + madspace_data = load_madspace_data() # parse command line arguments parser = argparse.ArgumentParser() @@ -159,11 +103,8 @@ def main() -> None: parser.add_argument("--gpu_batch_size", type=int, default=gen_args["gpu_batch_size"]) args = parser.parse_args() - if madspace_data["source_hash"] != ms.SOURCE_HASH: - message = ( - "The madspace version is not identical to the one used to generate " - "the gridpack. This can lead to errors or incorrect results" - ) + if not source_hash_matches(madspace_data): + message = SOURCE_HASH_MESSAGE if not args.ignore_source_hash: sys.exit( f"\033[1m\033[31mERROR\033[39m: {message}.\n" @@ -211,29 +152,17 @@ def resolve_dir(cli_value, card_value): # initialize context device_names = args.device if args.device else run_args["device"] - cpu_mode = run_args["cpu_mode"] contexts = [] backends = [] for device_name in device_names: - if ":" in device_name: - device_type, device_index_str = device_name.split(":") - device_index = int(device_index_str) - else: - device_type = device_name - device_index = 0 - # cpu_mode names the SIMD width of the CPU code, so it applies to the - # 'cpu' devices only: cuda/hip build the backend named after the device. - backends.append(cpu_mode if device_type == "cpu" else device_type) - if device_type == "cuda": - device = ms.cuda_device(device_index) - pool_size = args.gpu_thread_pool_size - elif device_type == "hip": - device = ms.hip_device(device_index) - pool_size = args.gpu_thread_pool_size - else: - device = ms.cpu_device() - pool_size = args.cpu_thread_pool_size - contexts.append(ms.Context(device=device, thread_count=pool_size)) + context, backend = make_context( + device_name, + run_args["cpu_mode"], + args.cpu_thread_pool_size, + args.gpu_thread_pool_size, + ) + contexts.append(context) + backends.append(backend) # set up generator configuration config = ms.GeneratorConfig() @@ -247,31 +176,9 @@ def resolve_dir(cli_value, card_value): config.cut_efficiency_threshold = gen_args["cut_efficiency_threshold"] config.max_cut_repetitions = gen_args["max_cut_repetitions"] - # set up contexts - # run-time matrix-element parameters (bwcutoff) of the run that made the - # gridpack; gridpacks written before they were recorded used the default - me_parameters = madspace_data.get("me_parameters", {}) - globals_path = os.path.join("data", "globals.tar") - for context, backend in zip(contexts, backends): - context.load_globals(globals_path) - for me_path in madspace_data["matrix_elements"]: - context.load_matrix_element( - me_path.format(device=backend), param_card_path, me_parameters - ) - - # set up generators - with open(os.path.join("data", "channels.json")) as f: - channels = json.load(f) - channel_generators = [ - ms.ChannelEventGenerator.load_json( - json.dumps(channel), - contexts, - event_file=os.path.join(temp_dir, f"events.{name}.npy"), - weight_file=os.path.join(temp_dir, f"weights.{name}.npy"), - config=config, - ) - for name, channel in channels.items() - ] + # set up contexts and generators + load_context_data(contexts, backends, madspace_data) + channel_generators = load_channels(contexts, config, temp_dir) event_generator = ms.EventGenerator( contexts=contexts, channels=channel_generators, @@ -281,8 +188,8 @@ def resolve_dir(cli_value, card_value): ) # scale/PDF systematics (as configured when the gridpack was made) - systematics = load_systematics(run_card, backends, param_card_path, - me_parameters) + systematics = load_systematics(run_card, backends, + madspace_data.get("me_parameters", {})) # run generation event_generator.generate() @@ -292,24 +199,24 @@ def resolve_dir(cli_value, card_value): os.path.join(run_path, "events.npy"), systematics ) write_lhe_header(run_path, - build_lhe_meta(event_generator, seed, systematics)) + build_lhe_meta(event_generator.status(), seed, systematics)) # what npy_to_lhe needs to complete these events into LHE later shutil.copy(os.path.join("data", "lhe.json"), os.path.join(run_path, "lhe_completer.json")) elif output_format == "lhe_npy": - lhe_completer = ms.LHECompleter.load(os.path.join("data", "lhe.json")) + lhe_completer = load_lhe_completer() event_generator.combine_to_lhe_npy( os.path.join(run_path, "events.npy"), lhe_completer, systematics ) write_lhe_header(run_path, - build_lhe_meta(event_generator, seed, systematics)) + build_lhe_meta(event_generator.status(), seed, systematics)) elif output_format == "lhe": - lhe_completer = ms.LHECompleter.load(os.path.join("data", "lhe.json")) + lhe_completer = load_lhe_completer() lhe_path = os.path.join(run_path, "events.lhe") # the cross section is known once generate() has converged, which is # before combine_to_lhe writes the block event_generator.combine_to_lhe( - lhe_path, lhe_completer, build_lhe_meta(event_generator, seed), + lhe_path, lhe_completer, build_lhe_meta(event_generator.status(), seed), systematics ) # Ship the LHE compressed, as the launcher that produced this gridpack @@ -332,118 +239,6 @@ def resolve_dir(cli_value, card_value): json.dump(data, f, indent=1) -def _pdf_search_paths(stored_path): - """Directories in which the LHAPDF grids may live: the LHAPDF_DATA_PATH - entries, the paths known to the lhapdf module, then the directory the - gridpack was created with.""" - paths = [] - env = os.environ.get("LHAPDF_DATA_PATH") - if env: - paths += env.split(os.pathsep) - try: - import lhapdf - paths += list(lhapdf.paths()) - except Exception: - pass - if stored_path: - paths.append(stored_path) - return paths - - -def _locate_pdf_file(stored_file): - """Find a PDF grid/info file: the stored absolute path, else the same - / relative to one of the PDF search paths.""" - if os.path.exists(stored_file): - return stored_file - set_dir, name = os.path.split(stored_file) - set_root, set_name = os.path.split(set_dir) - for base in _pdf_search_paths(set_root): - candidate = os.path.join(base, set_name, name) - if os.path.exists(candidate): - return candidate - raise FileNotFoundError( - "PDF file %s not found (set LHAPDF_DATA_PATH to the LHAPDF data directory)" - % stored_file) - - -def load_systematics(run_card, backends=(), param_card_path=None, - me_parameters=None): - """Rebuild the ms.SystematicsCalculator saved with the gridpack - (data/systematics.json) when [systematics] enable is set; None otherwise. - The matrix elements of the mixed-order subprocesses are reloaded into a CPU - context (when a CPU library is available) for the mu_R variations.""" - path = os.path.join("data", "systematics.json") - try: - enabled = bool(run_card["systematics"]["enable"]) - except Exception: - enabled = False - if not enabled or not os.path.exists(path): - return None - with open(path) as f: - data = json.load(f) - config = ms.SystematicsConfig.from_json(json.dumps(data["config"])) - for spec in config.pdf_members: - spec.grid_file = _locate_pdf_file(spec.grid_file) - spec.info_file = _locate_pdf_file(spec.info_file) - subproc_args = [ - ms.SubprocessSystArgs.from_json(json.dumps(a)) for a in data["subproc_args"] - ] - nominal_pdf = None - if config.has_pdf: - nominal_pdf = ms.PdfGrid(_locate_pdf_file(data["nominal_grid_file"])) - # The mu_R variations always need an alpha_s grid, so a gridpack that does - # not name one cannot reweight anything -- drop the systematics rather than - # die here (gridpacks written before the launcher recorded the file for a - # run without parton luminosity). - if not data.get("nominal_info_file"): - print("WARNING systematics: the gridpack records no alpha_s .info file, " - "the scale/PDF weights are dropped") - return None - nominal_alpha_s = ms.AlphaSGrid(_locate_pdf_file(data["nominal_info_file"])) - # PDFs, alpha_s and matrix elements are evaluated on this CPU context - context = ms.Context(device=ms.cpu_device(), thread_count=1) - matrix_elements, flavor_remap = [], [] - need = [i for i, a in enumerate(subproc_args) if a.qcd_power < 0] - # the CPU backend of this run (cpu_mode, resolved when the gridpack was - # saved), else the one the launcher used - backend = data.get("me_backend") - if need and data.get("me_paths"): - cpu = [b for b in backends if not str(b).startswith(("cuda", "hip")) and b != "auto"] - if cpu: - backend = cpu[0] - if need and backend and data.get("me_paths"): - flavor_remap = data.get("flavor_remap", []) - for i, me_path in enumerate(data["me_paths"]): - if i not in need: - matrix_elements.append(None) - continue - lib = me_path.format(device=backend) - if not os.path.exists(lib): - print("WARNING systematics: %s not found, mu_R variations of the " - "mixed-order subprocesses are dropped" % lib) - matrix_elements, flavor_remap = [], [] - break - api = context.load_matrix_element(lib, param_card_path, - me_parameters or {}) - matrix_elements.append(ms.MatrixElement( - api, - [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, - ms.MatrixElement.flavor_in], - [ms.MatrixElement.matrix_element_out], - False, - # the frame the events were generated with, not whatever the - # gridpack's run card says now: it has to match the matrix - # element inside the saved channel generators - data.get("me_frame", []), - data.get("incoming_count", 2), - )) - systematics = ms.SystematicsCalculator( - config, subproc_args, nominal_pdf, nominal_alpha_s, - context, matrix_elements, flavor_remap) - for warning in systematics.warnings: - print("WARNING systematics: %s" % warning) - return systematics - if __name__ == '__main__': os.chdir(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) try: diff --git a/madgraph/iolibs/template_files/mg7/gridpack_setup.py b/madgraph/iolibs/template_files/mg7/gridpack_setup.py new file mode 100644 index 0000000000..b7be349594 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/gridpack_setup.py @@ -0,0 +1,289 @@ +"""Gridpack setup shared by bin/generate_events and bin/event_stream.py: locating +madspace, reading the cards and data files, and building the contexts, channels, +LHE completer and systematics of the gridpack. Paths are resolved relative to +the gridpack directory, so nothing depends on the current directory.""" + +import json +import os +import subprocess +import sys +from pathlib import Path + +GRIDPACK_DIR = Path(os.path.realpath(__file__)).parent.parent + +# search order for madspace package: +# 1. a precompiled install (madspace/install/madspace) +# 2. bundled source that still needs to be built (madspace/install.py) +# 3. otherwise fall back to madspace is available in the environment +_LOCAL_MADSPACE_DIR = GRIDPACK_DIR / "madspace" +_LOCAL_INSTALL_DIR = _LOCAL_MADSPACE_DIR / "install" +if (_LOCAL_INSTALL_DIR / "madspace").is_dir(): + sys.path.insert(0, str(_LOCAL_INSTALL_DIR)) +elif (_LOCAL_MADSPACE_DIR / "install.py").is_file(): + print() + print("You don't have madspace installed for this gridpack") + print("Running interactive madspace installation script") + print() + + _result = subprocess.run([sys.executable, str(_LOCAL_MADSPACE_DIR / "install.py")]) + if _result.returncode != 0: + raise RuntimeError("madspace installation failed — see output above") + sys.path.insert(0, str(_LOCAL_INSTALL_DIR)) + +import madspace as ms + + +def gridpack_path(*parts) -> str: + """A path inside the gridpack directory.""" + return str(GRIDPACK_DIR.joinpath(*parts)) + + +def load_run_card(): + """The gridpack run card (Cards/grid_run_card.toml). Use the RunCardMG7 + representation when the madgraph package is importable; gridpacks are meant + to be portable, so fall back to a plain tomllib parse otherwise (the card is + the same TOML).""" + run_card_path = gridpack_path("Cards", "grid_run_card.toml") + try: + from madgraph.various.banner import RunCardMG7 + return RunCardMG7(run_card_path) + except ImportError: + import tomllib + with open(run_card_path, "rb") as f: + return tomllib.load(f) + + +def load_madspace_data() -> dict: + """data/data.json: the matrix elements, their run-time parameters and the + madspace source hash the gridpack was made with.""" + with open(gridpack_path("data", "data.json")) as f: + return json.load(f) + + +def source_hash_matches(madspace_data: dict) -> bool: + return madspace_data["source_hash"] == ms.SOURCE_HASH + + +SOURCE_HASH_MESSAGE = ( + "The madspace version is not identical to the one used to generate " + "the gridpack. This can lead to errors or incorrect results" +) + + +def resolve_seed(seed: int) -> int: + """Resolve the run_card "seed": -1 draws a fresh 64-bit seed via + os.urandom, any other value is used as-is.""" + if seed == -1: + return int.from_bytes(os.urandom(8), "big") + return seed + + +def make_context(device_name: str, cpu_mode: str, cpu_thread_pool_size: int, + gpu_thread_pool_size: int): + """The context for a device name like "cpu", "cuda:1" or "hip", with the + backend its matrix elements are built for.""" + if ":" in device_name: + device_type, device_index_str = device_name.split(":") + device_index = int(device_index_str) + else: + device_type = device_name + device_index = 0 + # cpu_mode names the SIMD width of the CPU code, so it applies to the + # 'cpu' devices only: cuda/hip build the backend named after the device. + backend = cpu_mode if device_type == "cpu" else device_type + if device_type == "cuda": + device = ms.cuda_device(device_index) + pool_size = gpu_thread_pool_size + elif device_type == "hip": + device = ms.hip_device(device_index) + pool_size = gpu_thread_pool_size + else: + device = ms.cpu_device() + pool_size = cpu_thread_pool_size + return ms.Context(device=device, thread_count=pool_size), backend + + +def load_context_data(contexts, backends, madspace_data: dict) -> None: + """Load the globals and the matrix elements of the gridpack into every + context.""" + # run-time matrix-element parameters (bwcutoff) of the run that made the + # gridpack; gridpacks written before they were recorded used the default + me_parameters = madspace_data.get("me_parameters", {}) + for context, backend in zip(contexts, backends): + context.load_globals(gridpack_path("data", "globals.tar")) + for me_path in madspace_data["matrix_elements"]: + context.load_matrix_element( + _gridpack_relative(me_path.format(device=backend)), + gridpack_path("Cards", "param_card.dat"), + me_parameters, + ) + + +def load_channels(contexts, config, file_dir=None): + """The channel generators stored in data/channels.json. With `file_dir`, + their temporary event files are written there; without, events are only + kept in memory.""" + with open(gridpack_path("data", "channels.json")) as f: + channels = json.load(f) + return [ + ms.ChannelEventGenerator.load_json( + json.dumps(channel), + contexts, + event_file=os.path.join(file_dir, f"events.{name}.npy") if file_dir else "", + weight_file=os.path.join(file_dir, f"weights.{name}.npy") if file_dir else "", + config=config, + ) + for name, channel in channels.items() + ] + + +def load_lhe_completer(): + return ms.LHECompleter.load(gridpack_path("data", "lhe.json")) + + +def build_lhe_meta(status, seed: int, systematics=None): + """The LHE header/ metadata of this gridpack run: the cards, beams + and PDF the gridpack was made with (data/lhe_meta.json), with the cross + section (from the generator `status`) and seed of this run. + ``systematics`` adds the block, for the writers that do not + inject it themselves (header.lhe); leave it unset for combine_to_lhe, + which does.""" + headers = [] + data = {} + meta_path = gridpack_path("data", "lhe_meta.json") + if os.path.exists(meta_path): + with open(meta_path) as f: + data = json.load(f) + headers = [ms.LHEHeader(name=h["name"], content=h["content"], + escape_content=h["escape_content"]) + for h in data.pop("headers")] + headers.append(ms.LHEHeader(name="MG7Seed", content=str(seed))) + if systematics is not None and systematics.weight_ids: + headers.append(ms.LHEHeader( + name="initrwgt", content=systematics.initrwgt(), escape_content=False + )) + return ms.LHEMeta( + # positional: the pybind arg name for max_weight is non-kwarg-safe + processes=[ms.LHEProcess(status.mean, status.error, status.mean, 1)], + headers=headers, + **data, + ) + + +def _gridpack_relative(path: str) -> str: + """Paths stored relative to the gridpack refer to the gridpack directory.""" + return path if os.path.isabs(path) else gridpack_path(path) + + +def _pdf_search_paths(stored_path): + """Directories in which the LHAPDF grids may live: the LHAPDF_DATA_PATH + entries, the paths known to the lhapdf module, then the directory the + gridpack was created with.""" + paths = [] + env = os.environ.get("LHAPDF_DATA_PATH") + if env: + paths += env.split(os.pathsep) + try: + import lhapdf + paths += list(lhapdf.paths()) + except Exception: + pass + if stored_path: + paths.append(stored_path) + return paths + + +def _locate_pdf_file(stored_file): + """Find a PDF grid/info file: the stored path, else the same / + relative to one of the PDF search paths.""" + stored_file = _gridpack_relative(stored_file) + if os.path.exists(stored_file): + return stored_file + set_dir, name = os.path.split(stored_file) + set_root, set_name = os.path.split(set_dir) + for base in _pdf_search_paths(set_root): + candidate = os.path.join(base, set_name, name) + if os.path.exists(candidate): + return candidate + raise FileNotFoundError( + "PDF file %s not found (set LHAPDF_DATA_PATH to the LHAPDF data directory)" + % stored_file) + + +def load_systematics(run_card, backends=(), me_parameters=None): + """Rebuild the ms.SystematicsCalculator saved with the gridpack + (data/systematics.json) when [systematics] enable is set; None otherwise. + The matrix elements of the mixed-order subprocesses are reloaded into a CPU + context (when a CPU library is available) for the mu_R variations.""" + path = gridpack_path("data", "systematics.json") + try: + enabled = bool(run_card["systematics"]["enable"]) + except Exception: + enabled = False + if not enabled or not os.path.exists(path): + return None + with open(path) as f: + data = json.load(f) + config = ms.SystematicsConfig.from_json(json.dumps(data["config"])) + for spec in config.pdf_members: + spec.grid_file = _locate_pdf_file(spec.grid_file) + spec.info_file = _locate_pdf_file(spec.info_file) + subproc_args = [ + ms.SubprocessSystArgs.from_json(json.dumps(a)) for a in data["subproc_args"] + ] + nominal_pdf = None + if config.has_pdf: + nominal_pdf = ms.PdfGrid(_locate_pdf_file(data["nominal_grid_file"])) + # The mu_R variations always need an alpha_s grid, so a gridpack that does + # not name one cannot reweight anything -- drop the systematics rather than + # die here (gridpacks written before the launcher recorded the file for a + # run without parton luminosity). + if not data.get("nominal_info_file"): + print("WARNING systematics: the gridpack records no alpha_s .info file, " + "the scale/PDF weights are dropped") + return None + nominal_alpha_s = ms.AlphaSGrid(_locate_pdf_file(data["nominal_info_file"])) + # PDFs, alpha_s and matrix elements are evaluated on this CPU context + context = ms.Context(device=ms.cpu_device(), thread_count=1) + matrix_elements, flavor_remap = [], [] + need = [i for i, a in enumerate(subproc_args) if a.qcd_power < 0] + # the CPU backend of this run (cpu_mode, resolved when the gridpack was + # saved), else the one the launcher used + backend = data.get("me_backend") + if need and data.get("me_paths"): + cpu = [b for b in backends if not str(b).startswith(("cuda", "hip")) and b != "auto"] + if cpu: + backend = cpu[0] + if need and backend and data.get("me_paths"): + flavor_remap = data.get("flavor_remap", []) + for i, me_path in enumerate(data["me_paths"]): + if i not in need: + matrix_elements.append(None) + continue + lib = _gridpack_relative(me_path.format(device=backend)) + if not os.path.exists(lib): + print("WARNING systematics: %s not found, mu_R variations of the " + "mixed-order subprocesses are dropped" % lib) + matrix_elements, flavor_remap = [], [] + break + api = context.load_matrix_element( + lib, gridpack_path("Cards", "param_card.dat"), me_parameters or {} + ) + matrix_elements.append(ms.MatrixElement( + api, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in], + [ms.MatrixElement.matrix_element_out], + False, + # the frame the events were generated with, not whatever the + # gridpack's run card says now: it has to match the matrix + # element inside the saved channel generators + data.get("me_frame", []), + data.get("incoming_count", 2), + )) + systematics = ms.SystematicsCalculator( + config, subproc_args, nominal_pdf, nominal_alpha_s, + context, matrix_elements, flavor_remap) + for warning in systematics.warnings: + print("WARNING systematics: %s" % warning) + return systematics diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index 63da39cf77..3ea9a82969 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -1968,10 +1968,13 @@ def save_gridpack(self) -> None: bin_path = os.path.join(gridpack_path, "bin") os.mkdir(bin_path) gen_events_file = os.path.join(bin_path, "generate_events") - shutil.copy( - os.path.join(os.path.dirname(__file__), "gridpack.py"), gen_events_file - ) + template_dir = os.path.dirname(__file__) + shutil.copy(os.path.join(template_dir, "gridpack.py"), gen_events_file) os.chmod(gen_events_file, 0o755) + shutil.copy(os.path.join(template_dir, "gridpack_setup.py"), bin_path) + # streaming needs the fixed maximum weights of this mode + if self.gridpack_run_mode == "fix_max_weight": + shutil.copy(os.path.join(template_dir, "event_stream.py"), bin_path) data = { "matrix_elements": matrix_elements, diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 9b4d30f003..c59c842fc5 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -66,7 +66,9 @@ compress = %(gridpack.compress)s # unweighted events are reached; its maximum weight is then kept # fixed in the gridpack run, and its integral is combined with the # one of the gridpack run to share the events between channels -# (can be slow for low-efficiency channels) +# (can be slow for low-efficiency channels); the gridpack +# also gets bin/event_stream.py to stream events in memory +# from Python run_mode = %(gridpack.run_mode)s # gridpack runs only: directories for the temporary npy files and for the final # output files. "" uses the run folder Events/_; if output_dir diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index c7166ea774..cee7f8d4ca 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -209,6 +209,8 @@ add_library( src/driver/context.cpp src/driver/event_generator.cpp src/driver/channel_generator.cpp + src/driver/event_utils.cpp + src/driver/event_stream.cpp src/driver/generator_data.cpp src/driver/status_file.cpp src/driver/thread_pool.cpp diff --git a/madspace/include/madspace/driver.hpp b/madspace/include/madspace/driver.hpp index 948427276c..017da8d979 100644 --- a/madspace/include/madspace/driver.hpp +++ b/madspace/include/madspace/driver.hpp @@ -6,6 +6,8 @@ #include "driver/discrete_optimizer.hpp" #include "driver/event_generator.hpp" #include "driver/event_histograms.hpp" +#include "driver/event_stream.hpp" +#include "driver/event_utils.hpp" #include "driver/format.hpp" #include "driver/io.hpp" #include "driver/lhe_output.hpp" diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 884d72bd2f..b65c1f696b 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -52,8 +52,10 @@ class ChannelEventGenerator { /** * @param contexts One context per device to run on. * @param integrand The channel's compiled integrand. - * @param event_file Path of the weighted event file. - * @param weight_file Path of the max-weight tracking file. + * @param event_file Path of the weighted event file; empty for no + * file (events are only kept in memory). + * @param weight_file Path of the max-weight tracking file; empty + * for no file. * @param config Generator configuration. * @param subprocess_index Index of the owning subprocess. * @param name Human-readable channel name. @@ -86,12 +88,17 @@ class ChannelEventGenerator { std::size_t count_requested_opt() const { return _count_requested_opt; } /// The filled observable histograms, if any were requested. const std::vector& histograms() const { return _histograms; } - /// The weighted event file. - EventFile& event_file() { return _event_file; } - /// The max-weight tracking file. - EventFile& weight_file() { return _weight_file; } + /// The weighted event file. Throws if the channel was created without one. + EventFile& event_file(); + /// The max-weight tracking file. Throws if the channel was created without + /// one. + EventFile& weight_file(); + /// Number of particles per event. + std::size_t particle_count() const { return _particle_count; } /// Current maximum-weight estimate, used for unweighting. double max_weight() const { return _max_weight; } + /// Whether the maximum weight was fixed with @ref set_fixed_max_weight. + bool max_weight_fixed() const { return _max_weight_fixed; } /// Use `max_weight` for unweighting from now on, instead of estimating it /// from the generated events. Must be positive. void set_fixed_max_weight(double max_weight); @@ -162,6 +169,14 @@ class ChannelEventGenerator { /// Append `unweighted_events` to @ref event_file, recording /// `job_max_weight`. void write_events(const TensorVec& unweighted_events, double job_max_weight); + /// Pack `unweighted_events` into `event_buffer` (layout @ref + /// event_file_layout) and their weights into `weight_buffer` (layout + /// `weight_file_layout`), resizing both. + void fill_event_buffers( + const TensorVec& unweighted_events, + EventBuffer& event_buffer, + EventBuffer& weight_buffer + ) const; /// Serialize this generator's state to `file_name`; see @ref load. void save(const std::string& file_name) const; /// The JSON text that @ref save writes; see @ref load_json. With @@ -218,8 +233,8 @@ class ChannelEventGenerator { int _event_layout_extra_flags; int _particle_layout_extra_flags; DataLayout _event_file_layout; - EventFile _event_file; - EventFile _weight_file; + std::optional _event_file; + std::optional _weight_file; std::optional _vegas_optimizer; std::optional _discrete_optimizer; std::size_t _batch_size; diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 53f4e7f10e..20db35a3f2 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -241,13 +241,6 @@ class EventGenerator { double norm_factor, MixMaxRandom& rand_gen ); - void fill_lhe_event( - LHECompleter& lhe_completer, - LHEEvent& lhe_event, - EventBuffer& buffer, - std::size_t event_index, - MixMaxRandom& rand_gen - ); // Layout of the combined events as read from the channel files DataLayout combined_layout() const; // Layout of the written events: the reweighting inputs are dropped unless diff --git a/madspace/include/madspace/driver/event_stream.hpp b/madspace/include/madspace/driver/event_stream.hpp new file mode 100644 index 0000000000..ff0d0eea1a --- /dev/null +++ b/madspace/include/madspace/driver/event_stream.hpp @@ -0,0 +1,167 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "madspace/driver/channel_generator.hpp" +#include "madspace/driver/context.hpp" +#include "madspace/driver/generator_data.hpp" +#include "madspace/driver/io.hpp" +#include "madspace/driver/lhe_output.hpp" +#include "madspace/driver/random.hpp" +#include "madspace/driver/thread_pool.hpp" + +namespace madspace { + +/** + * Endless stream of unweighted events, generated in memory. + * + * Every channel needs a fixed maximum weight (see @ref + * ChannelEventGenerator::set_fixed_max_weight), so that an event is final as + * soon as it is unweighted; channels without one are skipped. A coordinator + * thread keeps generation jobs running on the context's thread pool and + * buffers each channel's unweighted events, in launch order, until they are + * consumed. @ref next_events picks the channel of every event at random, + * proportional to its current cross-section estimate, and completes the + * events with an @ref LHECompleter. + * + * The integral estimates only include the jobs whose events were requested + * so far, so the event sequence is a pure function of the seed: it depends + * neither on thread timing nor on how the requests are split into batches. + * + * Event weights are in units of the maximum weight: +-1, except for + * overweight events. + */ +class EventStream { +public: + /// Install `func`, called periodically while @ref next_events waits for + /// events, to check for a requested abort. + static void set_abort_check_function(std::function func) { + _abort_check_function = func; + } + + /** + * Start the coordinator thread and the first generation jobs. + * + * @param context Context to generate on. + * @param channels The channels, each created with `context` as its + * only context. + * @param seed Run seed every random stream is derived from. + * @param lhe_completer Completes the events into LHE events. + * @param config Generator configuration; the batch sizes and cut + * settings are used. + */ + EventStream( + ContextPtr context, + const std::vector>& channels, + std::uint64_t seed, + const LHECompleter& lhe_completer, + const GeneratorConfig& config + ); + ~EventStream(); + EventStream(const EventStream&) = delete; + EventStream& operator=(const EventStream&) = delete; + + /// The next `count` events of the stream. Blocks until they are available. + std::vector next_events(std::size_t count); + /// Stop the coordinator thread and discard the running jobs. Called by + /// the destructor. + void close(); + /// Combined status over every channel. `count_unweighted` is the number of + /// events returned so far. + GeneratorStatus status() const; + /// Per-channel status; `count_unweighted` is the number of events of the + /// channel returned so far. + std::vector channel_status() const; + /// Current probability of each channel to be picked for the next event. + std::vector channel_probabilities() const; + /// Number of events returned so far. + std::size_t event_count() const; + /// Largest particle count of the returned LHE events. + std::size_t max_particle_count() const { + return _lhe_completer.max_particle_count(); + } + +private: + // A job whose events are buffered, waiting to be consumed. + struct ClosedJob { + GeneratorBatchJob job; + EventBuffer events; + EventBuffer weights; + bool opened = false; + std::size_t next_event = 0; + }; + struct ChannelState { + bool active = false; + // job ids in launch order, waiting to be closed + std::deque launch_order; + std::map finished; + std::deque closed; + std::size_t in_flight = 0; + std::size_t buffered_events = 0; + std::size_t unopened_jobs = 0; + // statistics of the completed and of the opened jobs + std::size_t completed_jobs = 0; + std::size_t completed_events = 0; + std::size_t opened_jobs = 0; + double opened_abs_sum = 0.; + double opened_capped_sum = 0.; + std::size_t returned_events = 0; + }; + // thrown inside the coordinator to stop it + struct StopRequest {}; + + inline static std::function _abort_check_function = [] {}; + + ContextPtr _context; + std::vector> _channels; + std::uint64_t _seed; + LHECompleter _lhe_completer; + GeneratorConfig _config; + DataLayout _combined_layout; + MixMaxRandom _select_rng; + MixMaxRandom _lhe_rng; + + // guards everything below, shared by the consumer and the coordinator + mutable std::mutex _mutex; + std::condition_variable _coordinator_cv; + std::condition_variable _consumer_cv; + std::vector _channel_states; + std::vector _cum_probabilities; + bool _probabilities_dirty = true; + bool _dispatch_requested = false; + bool _stopping = false; + std::exception_ptr _error; + std::size_t _lookahead; + std::size_t _event_count = 0; + bool _started = false; + + // only used by the consumer, which next_events() serializes + std::mutex _consumer_mutex; + // only used by the coordinator thread + std::unordered_map _running_jobs; + std::size_t _results_pending = 0; + std::size_t _job_id = 0; + ResultQueue _result_queue; + std::thread _coordinator; + + void run_coordinator(); + void dispatch_jobs(); + void handle_result(std::size_t job_id); + double expected_job_events(std::size_t channel_index) const; + void update_probabilities(); + std::size_t select_channel(); + ClosedJob& next_job(std::size_t channel_index, std::unique_lock& lock); + void open_job(std::size_t channel_index, ClosedJob& closed_job); + void wait_for_events(std::size_t channel_index, std::unique_lock& lock); +}; + +} // namespace madspace diff --git a/madspace/include/madspace/driver/event_utils.hpp b/madspace/include/madspace/driver/event_utils.hpp new file mode 100644 index 0000000000..2d8e2639b3 --- /dev/null +++ b/madspace/include/madspace/driver/event_utils.hpp @@ -0,0 +1,70 @@ +#pragma once + +#include +#include +#include +#include + +#include "madspace/driver/channel_generator.hpp" +#include "madspace/driver/generator_data.hpp" +#include "madspace/driver/io.hpp" +#include "madspace/driver/lhe_output.hpp" +#include "madspace/driver/random.hpp" + +namespace madspace { + +/// Weight of an unweighted event in units of its channel's maximum weight: +/// +-1, or larger in magnitude for an overweight event. +inline double unweighted_event_weight(double weight, double max_weight) { + return std::copysign(std::max(1., std::abs(weight) / max_weight), weight); +} + +/// Layout of combined events: weight, subprocess index and event data, plus +/// the extra fields (`extra_event_flags`, `extra_particle_flags`) the channels +/// carry. +DataLayout combined_event_layout(int extra_event_flags, int extra_particle_flags); + +/** + * Copy event `in_index` of a channel's event buffer to event `out_index` of a + * combined buffer (layout @ref combined_event_layout), padding the momenta of + * missing particles with zeros. + * + * @param weight Weight to store for the event. + * @param subprocess_index Stored if the channel records no per-event + * subprocess index. + * @param extra_flags The channel's extra `EventRecord` layout flags. + */ +void copy_channel_event( + EventBuffer& in_buffer, + std::size_t in_index, + EventBuffer& out_buffer, + std::size_t out_index, + double weight, + int subprocess_index, + int extra_flags +); + +/// Build `lhe_event` from event `event_index` of a combined buffer and +/// complete it with `lhe_completer`. +void fill_lhe_event( + LHECompleter& lhe_completer, + LHEEvent& lhe_event, + EventBuffer& buffer, + std::size_t event_index, + MixMaxRandom& rand_gen +); + +/// Write the particles of `lhe_event` to event `event_index` of a buffer with +/// LHE particle layout, padding the remaining slots with empty particles. +void write_lhe_particles( + const LHEEvent& lhe_event, EventBuffer& buffer, std::size_t event_index +); + +/// Sum the integrals and event counts of `channels` into `status`. Throws if +/// an integral is not finite. +void sum_channel_status( + const std::vector>& channels, + GeneratorStatus& status +); + +} // namespace madspace diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index b8bb79ee77..c5ffd2d186 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -72,6 +72,18 @@ DerivedSeed generate_phase_seed( ); } +// An empty path means the events are only kept in memory. +std::optional open_event_file( + const std::string& path, const DataLayout& layout, std::size_t particle_count +) { + if (path.empty()) { + return std::nullopt; + } + return std::make_optional( + path, layout, particle_count, EventFile::create, true + ); +} + } // namespace ChannelEventGenerator::ChannelEventGenerator( @@ -114,13 +126,9 @@ ChannelEventGenerator::ChannelEventGenerator( ) ), _event_file( - event_file, - _event_file_layout, - integrand.particle_count(), - EventFile::create, - true + open_event_file(event_file, _event_file_layout, integrand.particle_count()) ), - _weight_file(weight_file, weight_file_layout, 0, EventFile::create, true), + _weight_file(open_event_file(weight_file, weight_file_layout, 0)), _batch_size(config.start_batch_size), _particle_count(integrand.particle_count()), _integrand_channel_function(IntegrandChannelPart(integrand).function()), @@ -241,10 +249,8 @@ ChannelEventGenerator::ChannelEventGenerator( ParticleRecord::f_particle_data | _particle_layout_extra_flags ) ), - _event_file( - event_file, _event_file_layout, particle_count, EventFile::create, true - ), - _weight_file(weight_file, weight_file_layout, 0, EventFile::create, true), + _event_file(open_event_file(event_file, _event_file_layout, particle_count)), + _weight_file(open_event_file(weight_file, weight_file_layout, 0)), _batch_size(config.start_batch_size), _particle_count(particle_count), _integrand_channel_function(integrand_channel_function), @@ -327,14 +333,28 @@ void ChannelEventGenerator::init_field_indices() { _field_indices.rest = _field_indices.random + 1; } +EventFile& ChannelEventGenerator::event_file() { + if (!_event_file) { + throw std::logic_error("Channel was created without an event file"); + } + return *_event_file; +} + +EventFile& ChannelEventGenerator::weight_file() { + if (!_weight_file) { + throw std::logic_error("Channel was created without a weight file"); + } + return *_weight_file; +} + void ChannelEventGenerator::unweight_file(MixMaxRandom& rand_gen) { std::size_t buf_size = 1000000; EventBuffer buffer(0, 0, weight_file_layout); std::size_t accept_count = _unweighted_accept_count; - for (std::size_t i = _unweighted_count; i < _weight_file.event_count(); - i += buf_size) { - _weight_file.seek(i); - _weight_file.read(buffer, buf_size); + auto& w_file = weight_file(); + for (std::size_t i = _unweighted_count; i < w_file.event_count(); i += buf_size) { + w_file.seek(i); + w_file.read(buffer, buf_size); for (std::size_t j = 0; j < buffer.event_count(); ++j) { auto weight = buffer.event(j).weight(); if (std::abs(weight.value()) / _max_weight < rand_gen.generate_double()) { @@ -346,10 +366,10 @@ void ChannelEventGenerator::unweight_file(MixMaxRandom& rand_gen) { ++accept_count; } } - _weight_file.seek(i); - _weight_file.write(buffer); + w_file.seek(i); + w_file.write(buffer); } - _unweighted_count = _weight_file.event_count(); + _unweighted_count = w_file.event_count(); _unweighted_accept_count = accept_count; _status.count_unweighted = accept_count; } @@ -422,10 +442,11 @@ double ChannelEventGenerator::channel_weight_sum(std::size_t event_count) { std::size_t buf_size = 1000000; EventBuffer buffer(0, 0, weight_file_layout); double weight_sum = 0; - _weight_file.seek(0); + auto& w_file = weight_file(); + w_file.seek(0); std::size_t unweighted_count = 0; - for (std::size_t i = 0; i < _weight_file.event_count(); i += buf_size) { - _weight_file.read(buffer, buf_size); + for (std::size_t i = 0; i < w_file.event_count(); i += buf_size) { + w_file.read(buffer, buf_size); bool done = false; for (std::size_t j = 0; j < buffer.event_count(); ++j) { if (unweighted_count == event_count) { @@ -647,8 +668,12 @@ void ChannelEventGenerator::clear_events() { _status.count_opt = 0; _count_requested_opt = 0; _status.count_after_cuts_opt = 0; - _event_file.clear(); - _weight_file.clear(); + if (_event_file) { + _event_file->clear(); + } + if (_weight_file) { + _weight_file->clear(); + } _cross_section.reset(); _abs_cross_section.reset(); _large_weights.clear(); @@ -725,6 +750,21 @@ void ChannelEventGenerator::apply_truncation_budget() { void ChannelEventGenerator::write_events( const std::vector& unweighted_events, double job_max_weight ) { + EventBuffer event_buffer(0, _particle_count, _event_file_layout); + EventBuffer weight_buffer(0, 0, weight_file_layout); + fill_event_buffers(unweighted_events, event_buffer, weight_buffer); + event_file().write(event_buffer); + weight_file().write(weight_buffer); + _status.count_unweighted += event_buffer.event_count() * + (job_max_weight > 0 ? job_max_weight / _max_weight : 1); + _status.done = _status.count_unweighted >= _status.count_target; +} + +void ChannelEventGenerator::fill_event_buffers( + const TensorVec& unweighted_events, + EventBuffer& event_buffer, + EventBuffer& weight_buffer +) const { auto w_view = unweighted_events.at(_field_indices.weight).view(); auto mom_view = unweighted_events.at(_field_indices.momenta).view(); auto colors_view = @@ -739,10 +779,8 @@ void ChannelEventGenerator::write_events( unweighted_events.at(_field_indices.ren_scale).view(); auto alphas_view = unweighted_events.at(_field_indices.alpha_qcd).view(); - EventBuffer event_buffer( - w_view.size(), _event_file.particle_count(), _event_file_layout - ); - EventBuffer weight_buffer(w_view.size(), 0, weight_file_layout); + event_buffer.resize(w_view.size()); + weight_buffer.resize(w_view.size()); for (std::size_t i = 0; i < w_view.size(); ++i) { weight_buffer.event(i).weight() = w_view[i]; auto event = event_buffer.event(i); @@ -803,12 +841,6 @@ void ChannelEventGenerator::write_events( event.subprocess_index() = subproc_view[i]; } } - - _event_file.write(event_buffer); - _weight_file.write(weight_buffer); - _status.count_unweighted += - w_view.size() * (job_max_weight > 0 ? job_max_weight / _max_weight : 1); - _status.done = _status.count_unweighted >= _status.count_target; } void ChannelEventGenerator::save(const std::string& file_name) const { diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index de3a2a83c5..af12146d68 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -8,6 +8,7 @@ #include #include +#include "madspace/driver/event_utils.hpp" #include "madspace/driver/logger.hpp" #include "madspace/driver/random.hpp" #include "madspace/util.hpp" @@ -530,77 +531,7 @@ void EventGenerator::update_integral() { } void EventGenerator::update_integral_status() { - double total_mean = 0., total_var = 0.; - double total_mean_abs = 0., total_var_abs = 0.; - std::size_t total_count = 0, total_count_opt = 0; - std::size_t total_count_after_cuts = 0, total_count_after_cuts_opt = 0; - std::size_t total_integ_count = 0; - std::size_t iterations = 0; - bool optimized = true; - const ChannelEventGenerator* bad_channel = nullptr; - for (auto& channel : _channels) { - auto& status = channel->status(); - auto& cross_section = channel->cross_section(); - auto& abs_cross_section = channel->abs_cross_section(); - // special case for channels with 0 samples, as they have nan variance - if (bad_channel == nullptr && cross_section.count() > 0 && - (!std::isfinite(cross_section.mean()) || - !std::isfinite(cross_section.variance()) || - !std::isfinite(abs_cross_section.mean()) || - !std::isfinite(abs_cross_section.variance()))) { - bad_channel = channel.get(); - } - total_mean += cross_section.mean(); - total_var += cross_section.variance() / cross_section.count(); - total_mean_abs += abs_cross_section.mean(); - total_var_abs += abs_cross_section.variance() / abs_cross_section.count(); - total_count += status.count; - total_count_opt += status.count_opt; - total_count_after_cuts += status.count_after_cuts; - total_count_after_cuts_opt += status.count_after_cuts_opt; - total_integ_count += cross_section.count(); - iterations = std::max(status.iterations, iterations); - if (!channel->status().optimized) { - optimized = false; - } - } - if (bad_channel != nullptr || !std::isfinite(total_mean) || - !std::isfinite(total_mean_abs)) { - std::string where = bad_channel != nullptr - ? std::format( - "channel '{}' after {} samples (mean={}, variance={})", - bad_channel->status().name, - bad_channel->cross_section().count(), - bad_channel->cross_section().mean(), - bad_channel->cross_section().variance() - ) - : std::format("the sum over channels (mean={})", total_mean); - throw std::runtime_error( - std::format( - "non-finite integral in {}. A non-finite weight cannot be recovered " - "from by sampling further, so the integration is aborted. This usually " - "means the matrix element or the phase-space mapping returned nan/inf " - "for " - "some phase-space point -- check the process for a zero or negative " - "width, " - "a parameter point outside the model's validity, or a kinematic " - "configuration at a threshold.", - where - ) - ); - } - - _status.mean = total_mean; - _status.error = std::sqrt(total_var); - _status.mean_abs = total_mean_abs; - _status.error_abs = std::sqrt(total_var_abs); - _status.rel_std_dev = std::sqrt(total_var_abs * total_integ_count) / total_mean_abs; - _status.count = total_count; - _status.count_opt = total_count_opt; - _status.count_after_cuts = total_count_after_cuts; - _status.count_after_cuts_opt = total_count_after_cuts_opt; - _status.iterations = iterations; - _status.optimized = optimized; + sum_channel_status(_channels, _status); } void EventGenerator::update_integral_fractions() { @@ -761,13 +692,7 @@ void EventGenerator::combine_to_lhe_npy( fill_lhe_event(lhe_completer, lhe_event, buffer, i, rand_gen); buffer_out.event(i).from_lhe_event(lhe_event); copy_event_fields(buffer, i, buffer_out, i, syst_weights); - std::size_t j = 0; - for (; j < lhe_event.particles.size(); ++j) { - buffer_out.particle(i, j).from_lhe_particle(lhe_event.particles[j]); - } - for (; j < lhe_completer.max_particle_count(); ++j) { - buffer_out.particle(i, j).from_lhe_particle(LHEParticle{}); - } + write_lhe_particles(lhe_event, buffer_out, i); } event_file.write(buffer_out); if (event_count - last_update_count > 10000) { @@ -902,15 +827,9 @@ void EventGenerator::combine_to_lhe( } DataLayout EventGenerator::combined_layout() const { - return DataLayout( - EventRecord::layout( - EventRecord::f_weight | EventRecord::f_subproc_index | - EventRecord::f_event_data | _channels.at(0)->event_layout_extra_flags() - ), - ParticleRecord::layout( - ParticleRecord::f_particle_data | - _channels.at(0)->particle_layout_extra_flags() - ) + return combined_event_layout( + _channels.at(0)->event_layout_extra_flags(), + _channels.at(0)->particle_layout_extra_flags() ); } @@ -1163,12 +1082,7 @@ void EventGenerator::read_and_combine( std::size_t event_count = std::min(batch_size, channel_data.back().cum_count); buffer.resize(event_count); - bool has_beam1 = _channels.at(0)->event_layout_extra_flags() & EventRecord::f_beam1; - bool has_beam2 = _channels.at(0)->event_layout_extra_flags() & EventRecord::f_beam2; - bool has_partial = - _channels.at(0)->event_layout_extra_flags() & EventRecord::f_partial_weights; - bool has_subproc_index = - _channels.at(0)->event_layout_extra_flags() & EventRecord::f_subproc_index; + int extra_flags = _channels.at(0)->event_layout_extra_flags(); for (std::size_t event_index = 0; event_index < event_count; ++event_index) { std::size_t random_index = rand_gen.generate_int(channel_data.back().cum_count); @@ -1199,93 +1113,17 @@ void EventGenerator::read_and_combine( ++sampled_chan->buffer_index; } - auto event_in = sampled_chan->event_buffer.event(sampled_chan->buffer_index); - auto event_out = buffer.event(event_index); - event_out.weight() = - std::copysign( - std::max(1., std::abs(weight) / channel->max_weight()), weight - ) * - norm_factor; - event_out.subprocess_index() = has_subproc_index - ? event_in.subprocess_index().value() - : static_cast(channel->status().subprocess); - event_out.diagram_index() = event_in.diagram_index(); - event_out.color_index() = event_in.color_index(); - event_out.flavor_index() = event_in.flavor_index(); - event_out.helicity_index() = event_in.helicity_index(); - event_out.ren_scale() = event_in.ren_scale(); - event_out.alpha_qcd() = event_in.alpha_qcd(); - - if (has_beam1) { - event_out.x1() = event_in.x1(); - event_out.fact_scale1() = event_in.fact_scale1(); - } - if (has_beam2) { - event_out.x2() = event_in.x2(); - event_out.fact_scale2() = event_in.fact_scale2(); - } - if (has_partial) { - event_out.partial_weight_product() = event_in.partial_weight_product(); - } - - std::size_t i = 0; - for (; i < sampled_chan->event_buffer.particle_count(); ++i) { - auto particle_in = - sampled_chan->event_buffer.particle(sampled_chan->buffer_index, i); - auto particle_out = buffer.particle(event_index, i); - particle_out.energy() = particle_in.energy(); - particle_out.px() = particle_in.px(); - particle_out.py() = particle_in.py(); - particle_out.pz() = particle_in.pz(); - } - for (; i < buffer.particle_count(); ++i) { - auto particle_out = buffer.particle(event_index, i); - particle_out.energy() = 0.; - particle_out.px() = 0.; - particle_out.py() = 0.; - particle_out.pz() = 0.; - } - ++sampled_chan->buffer_index; - } -} - -void EventGenerator::fill_lhe_event( - LHECompleter& lhe_completer, - LHEEvent& lhe_event, - EventBuffer& buffer, - std::size_t event_index, - MixMaxRandom& rand_gen -) { - EventRecord event_in = buffer.event(event_index); - lhe_event.weight = event_in.weight(); - lhe_event.process_id = 0; - lhe_event.scale = event_in.ren_scale(); - lhe_event.alpha_qed = 0; // TODO: populate this - lhe_event.alpha_qcd = event_in.alpha_qcd(); - lhe_event.particles.clear(); - for (std::size_t i = 0; i < buffer.particle_count(); ++i) { - auto particle_in = buffer.particle(event_index, i); - if (particle_in.energy() == 0.) { - break; - } - lhe_event.particles.push_back( - LHEParticle{ - .px = particle_in.px(), - .py = particle_in.py(), - .pz = particle_in.pz(), - .energy = particle_in.energy(), - } + copy_channel_event( + sampled_chan->event_buffer, + sampled_chan->buffer_index, + buffer, + event_index, + unweighted_event_weight(weight, channel->max_weight()) * norm_factor, + static_cast(channel->status().subprocess), + extra_flags ); + ++sampled_chan->buffer_index; } - lhe_completer.complete_event_data( - lhe_event, - event_in.subprocess_index(), - event_in.diagram_index(), - event_in.color_index(), - event_in.flavor_index(), - event_in.helicity_index(), - rand_gen - ); } void EventGenerator::init_status(const std::string& status) { diff --git a/madspace/src/driver/event_stream.cpp b/madspace/src/driver/event_stream.cpp new file mode 100644 index 0000000000..4798c70421 --- /dev/null +++ b/madspace/src/driver/event_stream.cpp @@ -0,0 +1,466 @@ +#include "madspace/driver/event_stream.hpp" + +#include +#include +#include +#include + +#include "madspace/driver/event_utils.hpp" +#include "madspace/util.hpp" + +using namespace madspace; +using namespace std::chrono_literals; + +namespace { + +// separates the stream's random streams from those of the combine pass +constexpr std::size_t stream_rng_index = 1; + +} // namespace + +EventStream::EventStream( + ContextPtr context, + const std::vector>& channels, + std::uint64_t seed, + const LHECompleter& lhe_completer, + const GeneratorConfig& config +) : + _context(context), + _channels(channels), + _seed(seed), + _lhe_completer(lhe_completer), + _config(config), + _combined_layout([&] { + if (channels.empty()) { + throw std::invalid_argument("EventStream needs at least one channel"); + } + return combined_event_layout( + channels.at(0)->event_layout_extra_flags(), + channels.at(0)->particle_layout_extra_flags() + ); + }()), + _select_rng(DerivedSeed(seed, DerivedSeed::combine_select, 0, 0, stream_rng_index)), + _lhe_rng(DerivedSeed(seed, DerivedSeed::lhe_complete, 0, 0, stream_rng_index)), + _channel_states(channels.size()), + _lookahead( + context->device()->device_type() == DeviceType::cpu + ? config.cpu_batch_size + : config.gpu_batch_size + ) { + bool any_active = false; + for (auto [channel, state] : zip(_channels, _channel_states)) { + state.active = channel->max_weight_fixed(); + any_active |= state.active; + } + if (!any_active) { + throw std::invalid_argument( + "EventStream needs channels with a fixed maximum weight" + ); + } + _coordinator = std::thread([this] { run_coordinator(); }); +} + +EventStream::~EventStream() { close(); } + +void EventStream::close() { + { + std::unique_lock lock(_mutex); + _stopping = true; + } + _coordinator_cv.notify_all(); + _consumer_cv.notify_all(); + if (_coordinator.joinable()) { + _coordinator.join(); + } +} + +void EventStream::run_coordinator() { + try { + // dispatches new jobs when the consumer asked for them, stops on close() + auto poll = [this] { + std::unique_lock lock(_mutex); + if (_stopping) { + throw StopRequest{}; + } + if (_dispatch_requested) { + _dispatch_requested = false; + dispatch_jobs(); + } + }; + while (true) { + { + std::unique_lock lock(_mutex); + if (_stopping) { + break; + } + _dispatch_requested = false; + dispatch_jobs(); + if (_results_pending == 0) { + _coordinator_cv.wait(lock, [this] { + return _stopping || _dispatch_requested; + }); + continue; + } + } + auto result = _result_queue.wait(poll, 10ms); + --_results_pending; + if (result.exception) { + std::rethrow_exception(result.exception); + } + handle_result(result.id); + } + } catch (StopRequest&) { + } catch (...) { + std::unique_lock lock(_mutex); + _error = std::current_exception(); + } + _consumer_cv.notify_all(); + _result_queue.discard(_results_pending, nullptr); + _results_pending = 0; + _running_jobs.clear(); +} + +double EventStream::expected_job_events(std::size_t channel_index) const { + auto& state = _channel_states.at(channel_index); + if (state.completed_jobs > 0) { + return std::max(state.completed_events, 1.) / state.completed_jobs; + } + // before the first job, assume a single event to stay on the safe side + return 1.; +} + +// Called with _mutex held. Which jobs are started depends on timing, but not +// what they contain: that is fixed by the channel's job sequence number. +void EventStream::dispatch_jobs() { + std::size_t max_jobs = 2 * _context->thread_pool().thread_count(); + if (_probabilities_dirty) { + update_probabilities(); + } + while (_running_jobs.size() < max_jobs) { + // the channel furthest below its share of the lookahead + std::size_t best_index = _channels.size(); + double best_ratio = 1.; + double prev_prob = 0.; + for (std::size_t i = 0; i < _channels.size(); ++i) { + double prob = _cum_probabilities.at(i) - prev_prob; + prev_prob = _cum_probabilities.at(i); + auto& state = _channel_states.at(i); + if (!state.active) { + continue; + } + // bounds the memory of channels that rarely yield events + if (state.in_flight + state.unopened_jobs >= max_jobs) { + continue; + } + double available = + state.buffered_events + state.in_flight * expected_job_events(i); + double target = prob * _lookahead; + double ratio; + if (state.buffered_events == 0 && state.in_flight == 0 && + state.unopened_jobs == 0) { + // every channel keeps at least one job buffered or running + ratio = -1.; + } else if (available < target) { + ratio = available / target; + } else { + continue; + } + if (ratio < best_ratio) { + best_ratio = ratio; + best_index = i; + } + } + if (best_index == _channels.size()) { + break; + } + + std::size_t job_id = _job_id++; + auto& job = + _running_jobs + .emplace( + job_id, + GeneratorBatchJob{ + .channel_index = best_index, + .unweight = true, + .batch_event_count = 0, + .split_job_count = 1, + .context_index = 0, + .job_id = job_id, + .is_vegas_batch = false, + } + ) + .first->second; + auto& state = _channel_states.at(best_index); + state.launch_order.push_back(job_id); + ++state.in_flight; + _channels.at(best_index)->start_job(job, _result_queue, _seed, false, 0); + ++_results_pending; + } +} + +void EventStream::handle_result(std::size_t job_id) { + auto& job = _running_jobs.at(job_id); + auto& channel = _channels.at(job.channel_index); + if (job.unweighted_events.empty()) { + // generation stage done, the max weight is fixed so unweight right away + std::unique_lock lock(_mutex); + channel->start_unweight_job(job, _result_queue); + ++_results_pending; + return; + } + + ClosedJob closed_job{ + .job = {}, + .events = + EventBuffer(0, channel->particle_count(), channel->event_file_layout()), + .weights = EventBuffer(0, 0, weight_file_layout), + }; + channel->fill_event_buffers( + job.unweighted_events, closed_job.events, closed_job.weights + ); + // only the weights are needed to integrate the job later + closed_job.job.channel_index = job.channel_index; + closed_job.job.weights = job.weights; + closed_job.job.requested_event_count = job.requested_event_count; + std::size_t channel_index = job.channel_index; + _running_jobs.erase(job_id); + + std::unique_lock lock(_mutex); + auto& state = _channel_states.at(channel_index); + ++state.completed_jobs; + state.completed_events += closed_job.events.event_count(); + state.finished.emplace(job_id, std::move(closed_job)); + // close the jobs in launch order + bool closed_any = false; + while (!state.launch_order.empty()) { + auto node = state.finished.extract(state.launch_order.front()); + if (node.empty()) { + break; + } + state.launch_order.pop_front(); + --state.in_flight; + state.buffered_events += node.mapped().events.event_count(); + ++state.unopened_jobs; + state.closed.push_back(std::move(node.mapped())); + closed_any = true; + } + if (closed_any) { + _consumer_cv.notify_all(); + } +} + +// Called with _mutex held. +void EventStream::update_probabilities() { + _cum_probabilities.resize(_channels.size()); + double total = 0.; + for (std::size_t i = 0; i < _channels.size(); ++i) { + auto& state = _channel_states.at(i); + double prob = 0.; + if (state.active) { + // divided by the mean weight of the unweighted events (> 1 with + // overweight events), so that the weight sums are proportional to + // the cross sections + double mean_weight = state.opened_capped_sum > 0 + ? state.opened_abs_sum / state.opened_capped_sum + : 1.; + prob = _channels.at(i)->abs_integral_estimate() / mean_weight; + if (!std::isfinite(prob) || prob < 0) { + prob = 0.; + } + } + total += prob; + _cum_probabilities.at(i) = total; + } + if (total > 0) { + for (auto& prob : _cum_probabilities) { + prob /= total; + } + } + _probabilities_dirty = false; +} + +std::size_t EventStream::select_channel() { + if (_probabilities_dirty) { + update_probabilities(); + } + if (_cum_probabilities.back() <= 0) { + throw std::runtime_error("EventStream: all channels have a zero cross section"); + } + // the first channel whose cumulative probability exceeds r, so channels + // with zero probability are never picked + double r = _select_rng.generate_double() * _cum_probabilities.back(); + auto it = std::upper_bound(_cum_probabilities.begin(), _cum_probabilities.end(), r); + if (it == _cum_probabilities.end()) { + it = std::lower_bound( + _cum_probabilities.begin(), + _cum_probabilities.end(), + _cum_probabilities.back() + ); + } + return it - _cum_probabilities.begin(); +} + +// Called with _mutex held. +void EventStream::open_job(std::size_t channel_index, ClosedJob& closed_job) { + auto& channel = _channels.at(channel_index); + auto& state = _channel_states.at(channel_index); + channel->integrate(closed_job.job); + double max_weight = channel->max_weight(); + auto w_view = closed_job.job.weights.view(); + for (std::size_t i = 0; i < w_view.size(); ++i) { + double w = std::abs(w_view[i]); + state.opened_abs_sum += w; + state.opened_capped_sum += std::min(w, max_weight); + } + ++state.opened_jobs; + --state.unopened_jobs; + closed_job.opened = true; + _probabilities_dirty = true; +} + +void EventStream::wait_for_events( + std::size_t channel_index, std::unique_lock& lock +) { + _dispatch_requested = true; + _coordinator_cv.notify_all(); + auto& state = _channel_states.at(channel_index); + while (state.closed.empty()) { + if (_error) { + std::rethrow_exception(_error); + } + if (_stopping) { + throw std::runtime_error("EventStream was closed"); + } + if (!_consumer_cv.wait_for(lock, 100ms, [&] { + return !state.closed.empty() || _error || _stopping; + })) { + lock.unlock(); + _abort_check_function(); + lock.lock(); + } + } +} + +// The first job of the channel's buffer that still has events, opened. +// Called with _mutex held. +EventStream::ClosedJob& +EventStream::next_job(std::size_t channel_index, std::unique_lock& lock) { + auto& state = _channel_states.at(channel_index); + while (true) { + if (state.closed.empty()) { + wait_for_events(channel_index, lock); + } + auto& front = state.closed.front(); + if (!front.opened) { + open_job(channel_index, front); + } + if (front.next_event < front.events.event_count()) { + return front; + } + state.closed.pop_front(); + } +} + +std::vector EventStream::next_events(std::size_t count) { + std::unique_lock consumer_lock(_consumer_mutex); + EventBuffer buffer(count, max_particle_count(), _combined_layout); + { + std::unique_lock lock(_mutex); + if (_error) { + std::rethrow_exception(_error); + } + if (_stopping) { + throw std::runtime_error("EventStream was closed"); + } + _lookahead = std::max(_lookahead, 2 * count); + if (!_started) { + // every channel's estimate starts from one job (also those + // without an integral prior), so the first events are already + // drawn with all channels known + for (std::size_t i = 0; i < _channels.size(); ++i) { + if (_channel_states.at(i).active) { + next_job(i, lock); + } + } + _started = true; + } + for (std::size_t i = 0; i < count; ++i) { + std::size_t channel_index = select_channel(); + auto& channel = _channels.at(channel_index); + auto& state = _channel_states.at(channel_index); + auto& job = next_job(channel_index, lock); + std::size_t event_index = job.next_event++; + --state.buffered_events; + ++state.returned_events; + copy_channel_event( + job.events, + event_index, + buffer, + i, + unweighted_event_weight( + job.weights.event(event_index).weight(), channel->max_weight() + ), + static_cast(channel->status().subprocess), + channel->event_layout_extra_flags() + ); + // keep the buffers filled while the batch is assembled + if (state.buffered_events == 0) { + _dispatch_requested = true; + _coordinator_cv.notify_all(); + } + } + _event_count += count; + _dispatch_requested = true; + } + _coordinator_cv.notify_all(); + + std::vector events(count); + for (std::size_t i = 0; i < count; ++i) { + fill_lhe_event(_lhe_completer, events.at(i), buffer, i, _lhe_rng); + } + return events; +} + +GeneratorStatus EventStream::status() const { + std::unique_lock lock(_mutex); + GeneratorStatus status{}; + sum_channel_status(_channels, status); + status.count_unweighted = _event_count; + status.count_target = 0; + status.optimized = true; + status.done = false; + return status; +} + +std::vector EventStream::channel_status() const { + std::unique_lock lock(_mutex); + std::vector ret; + for (auto [channel, state] : zip(_channels, _channel_states)) { + auto status = channel->status(); + status.count_unweighted = state.returned_events; + status.count_target = 0; + ret.push_back(status); + } + return ret; +} + +std::vector EventStream::channel_probabilities() const { + std::unique_lock lock(_mutex); + // const: compute from scratch instead of updating the cache + auto& self = const_cast(*this); + if (self._probabilities_dirty) { + self.update_probabilities(); + } + std::vector probs; + double prev = 0.; + for (double cum_prob : _cum_probabilities) { + probs.push_back(cum_prob - prev); + prev = cum_prob; + } + return probs; +} + +std::size_t EventStream::event_count() const { + std::unique_lock lock(_mutex); + return _event_count; +} diff --git a/madspace/src/driver/event_utils.cpp b/madspace/src/driver/event_utils.cpp new file mode 100644 index 0000000000..7b5567afd9 --- /dev/null +++ b/madspace/src/driver/event_utils.cpp @@ -0,0 +1,200 @@ +#include "madspace/driver/event_utils.hpp" + +#include +#include +#include + +using namespace madspace; + +DataLayout +madspace::combined_event_layout(int extra_event_flags, int extra_particle_flags) { + return DataLayout( + EventRecord::layout( + EventRecord::f_weight | EventRecord::f_subproc_index | + EventRecord::f_event_data | extra_event_flags + ), + ParticleRecord::layout(ParticleRecord::f_particle_data | extra_particle_flags) + ); +} + +void madspace::copy_channel_event( + EventBuffer& in_buffer, + std::size_t in_index, + EventBuffer& out_buffer, + std::size_t out_index, + double weight, + int subprocess_index, + int extra_flags +) { + auto event_in = in_buffer.event(in_index); + auto event_out = out_buffer.event(out_index); + event_out.weight() = weight; + event_out.subprocess_index() = extra_flags & EventRecord::f_subproc_index + ? event_in.subprocess_index().value() + : subprocess_index; + event_out.diagram_index() = event_in.diagram_index(); + event_out.color_index() = event_in.color_index(); + event_out.flavor_index() = event_in.flavor_index(); + event_out.helicity_index() = event_in.helicity_index(); + event_out.ren_scale() = event_in.ren_scale(); + event_out.alpha_qcd() = event_in.alpha_qcd(); + + if (extra_flags & EventRecord::f_beam1) { + event_out.x1() = event_in.x1(); + event_out.fact_scale1() = event_in.fact_scale1(); + } + if (extra_flags & EventRecord::f_beam2) { + event_out.x2() = event_in.x2(); + event_out.fact_scale2() = event_in.fact_scale2(); + } + if (extra_flags & EventRecord::f_partial_weights) { + event_out.partial_weight_product() = event_in.partial_weight_product(); + } + + std::size_t i = 0; + for (; i < in_buffer.particle_count(); ++i) { + auto particle_in = in_buffer.particle(in_index, i); + auto particle_out = out_buffer.particle(out_index, i); + particle_out.energy() = particle_in.energy(); + particle_out.px() = particle_in.px(); + particle_out.py() = particle_in.py(); + particle_out.pz() = particle_in.pz(); + } + for (; i < out_buffer.particle_count(); ++i) { + auto particle_out = out_buffer.particle(out_index, i); + particle_out.energy() = 0.; + particle_out.px() = 0.; + particle_out.py() = 0.; + particle_out.pz() = 0.; + } +} + +void madspace::fill_lhe_event( + LHECompleter& lhe_completer, + LHEEvent& lhe_event, + EventBuffer& buffer, + std::size_t event_index, + MixMaxRandom& rand_gen +) { + EventRecord event_in = buffer.event(event_index); + lhe_event.weight = event_in.weight(); + lhe_event.process_id = 0; + lhe_event.scale = event_in.ren_scale(); + lhe_event.alpha_qed = 0; // TODO: populate this + lhe_event.alpha_qcd = event_in.alpha_qcd(); + lhe_event.particles.clear(); + for (std::size_t i = 0; i < buffer.particle_count(); ++i) { + auto particle_in = buffer.particle(event_index, i); + if (particle_in.energy() == 0.) { + break; + } + lhe_event.particles.push_back( + LHEParticle{ + .px = particle_in.px(), + .py = particle_in.py(), + .pz = particle_in.pz(), + .energy = particle_in.energy(), + } + ); + } + lhe_completer.complete_event_data( + lhe_event, + event_in.subprocess_index(), + event_in.diagram_index(), + event_in.color_index(), + event_in.flavor_index(), + event_in.helicity_index(), + rand_gen + ); +} + +void madspace::write_lhe_particles( + const LHEEvent& lhe_event, EventBuffer& buffer, std::size_t event_index +) { + std::size_t i = 0; + for (; i < lhe_event.particles.size(); ++i) { + buffer.particle(event_index, i).from_lhe_particle(lhe_event.particles[i]); + } + for (; i < buffer.particle_count(); ++i) { + buffer.particle(event_index, i).from_lhe_particle(LHEParticle{}); + } +} + +void madspace::sum_channel_status( + const std::vector>& channels, + GeneratorStatus& status +) { + double total_mean = 0., total_var = 0.; + double total_mean_abs = 0., total_var_abs = 0.; + std::size_t total_count = 0, total_count_opt = 0; + std::size_t total_count_after_cuts = 0, total_count_after_cuts_opt = 0; + std::size_t total_integ_count = 0; + std::size_t iterations = 0; + bool optimized = true; + const ChannelEventGenerator* bad_channel = nullptr; + for (auto& channel : channels) { + auto& chan_status = channel->status(); + auto& cross_section = channel->cross_section(); + auto& abs_cross_section = channel->abs_cross_section(); + // special case for channels with 0 samples, as they have nan variance + if (bad_channel == nullptr && cross_section.count() > 0 && + (!std::isfinite(cross_section.mean()) || + !std::isfinite(cross_section.variance()) || + !std::isfinite(abs_cross_section.mean()) || + !std::isfinite(abs_cross_section.variance()))) { + bad_channel = channel.get(); + } + total_mean += cross_section.mean(); + total_mean_abs += abs_cross_section.mean(); + if (cross_section.count() > 0) { + total_var += cross_section.variance() / cross_section.count(); + total_var_abs += abs_cross_section.variance() / abs_cross_section.count(); + } + total_count += chan_status.count; + total_count_opt += chan_status.count_opt; + total_count_after_cuts += chan_status.count_after_cuts; + total_count_after_cuts_opt += chan_status.count_after_cuts_opt; + total_integ_count += cross_section.count(); + iterations = std::max(chan_status.iterations, iterations); + if (!chan_status.optimized) { + optimized = false; + } + } + if (bad_channel != nullptr || !std::isfinite(total_mean) || + !std::isfinite(total_mean_abs)) { + std::string where = bad_channel != nullptr + ? std::format( + "channel '{}' after {} samples (mean={}, variance={})", + bad_channel->status().name, + bad_channel->cross_section().count(), + bad_channel->cross_section().mean(), + bad_channel->cross_section().variance() + ) + : std::format("the sum over channels (mean={})", total_mean); + throw std::runtime_error( + std::format( + "non-finite integral in {}. A non-finite weight cannot be recovered " + "from by sampling further, so the integration is aborted. This usually " + "means the matrix element or the phase-space mapping returned nan/inf " + "for " + "some phase-space point -- check the process for a zero or negative " + "width, " + "a parameter point outside the model's validity, or a kinematic " + "configuration at a threshold.", + where + ) + ); + } + + status.mean = total_mean; + status.error = std::sqrt(total_var); + status.mean_abs = total_mean_abs; + status.error_abs = std::sqrt(total_var_abs); + status.rel_std_dev = std::sqrt(total_var_abs * total_integ_count) / total_mean_abs; + status.count = total_count; + status.count_opt = total_count_opt; + status.count_after_cuts = total_count_after_cuts; + status.count_after_cuts_opt = total_count_after_cuts_opt; + status.iterations = iterations; + status.optimized = optimized; +} diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 89af1a5b34..3d61b7eca4 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1,4 +1,5 @@ #include +#include #include #include #include @@ -141,6 +142,65 @@ void named_vector_instance(py::module_& m, const char* name) { ); } +// One numpy array per LHE event field (shape (n,)) and per LHE particle field +// (shape (n, max_particle_count), zero padded), plus the particle counts. +py::dict lhe_events_to_numpy( + const std::vector& events, std::size_t max_particle_count +) { + py::ssize_t n = events.size(); + py::ssize_t m = max_particle_count; + auto event_field = [&](auto member) { + using T = std::remove_cvref_t; + py::array_t array(n); + auto view = array.template mutable_unchecked<1>(); + for (py::ssize_t i = 0; i < n; ++i) { + view(i) = events[i].*member; + } + return array; + }; + auto particle_field = [&](auto member) { + using T = std::remove_cvref_t; + py::array_t array({n, m}); + auto view = array.template mutable_unchecked<2>(); + for (py::ssize_t i = 0; i < n; ++i) { + auto& particles = events[i].particles; + for (py::ssize_t j = 0; j < m; ++j) { + view(i, j) = j < static_cast(particles.size()) + ? particles[j].*member + : T{}; + } + } + return array; + }; + py::array_t particle_count(n); + auto count_view = particle_count.mutable_unchecked<1>(); + for (py::ssize_t i = 0; i < n; ++i) { + count_view(i) = events[i].particles.size(); + } + + py::dict ret; + ret["process_id"] = event_field(&LHEEvent::process_id); + ret["weight"] = event_field(&LHEEvent::weight); + ret["scale"] = event_field(&LHEEvent::scale); + ret["alpha_qed"] = event_field(&LHEEvent::alpha_qed); + ret["alpha_qcd"] = event_field(&LHEEvent::alpha_qcd); + ret["particle_count"] = particle_count; + ret["pdg_id"] = particle_field(&LHEParticle::pdg_id); + ret["status_code"] = particle_field(&LHEParticle::status_code); + ret["mother1"] = particle_field(&LHEParticle::mother1); + ret["mother2"] = particle_field(&LHEParticle::mother2); + ret["color"] = particle_field(&LHEParticle::color); + ret["anti_color"] = particle_field(&LHEParticle::anti_color); + ret["px"] = particle_field(&LHEParticle::px); + ret["py"] = particle_field(&LHEParticle::py); + ret["pz"] = particle_field(&LHEParticle::pz); + ret["energy"] = particle_field(&LHEParticle::energy); + ret["mass"] = particle_field(&LHEParticle::mass); + ret["lifetime"] = particle_field(&LHEParticle::lifetime); + ret["spin"] = particle_field(&LHEParticle::spin); + return ret; +} + } // namespace PYBIND11_MODULE(_madspace_py, m) { @@ -3091,8 +3151,8 @@ PYBIND11_MODULE(_madspace_py, m) { &ChannelEventGenerator::load_json, py::arg("channel_json"), py::arg("contexts"), - py::arg("event_file"), - py::arg("weight_file"), + py::arg("event_file") = "", + py::arg("weight_file") = "", py::arg("config"), pydoc::doc("ChannelEventGenerator::load_json") ) @@ -3704,6 +3764,71 @@ PYBIND11_MODULE(_madspace_py, m) { pydoc::doc("EventGenerator::channels") ); + py::classh(m, "EventStream", pydoc::doc("EventStream")) + .def( + py::init< + ContextPtr, + const std::vector>&, + std::uint64_t, + const LHECompleter&, + const GeneratorConfig&>(), + py::arg("context"), + py::arg("channels"), + py::arg("seed"), + py::arg("lhe_completer"), + py::arg("config"), + pydoc::doc("EventStream::EventStream") + ) + .def( + "next_events", + &EventStream::next_events, + py::arg("count"), + py::call_guard(), + pydoc::doc("EventStream::next_events") + ) + .def( + "next_batch", + [](EventStream& stream, std::size_t count) { + std::vector events; + { + py::gil_scoped_release release; + events = stream.next_events(count); + } + return lhe_events_to_numpy(events, stream.max_particle_count()); + }, + py::arg("count"), + "Like next_events, as a dict of numpy arrays: the LHE event fields " + "with shape (count,), the LHE particle fields with shape (count, " + "max_particle_count), padded with zeros, and particle_count." + ) + .def( + "close", + &EventStream::close, + py::call_guard(), + pydoc::doc("EventStream::close") + ) + .def("status", &EventStream::status, pydoc::doc("EventStream::status")) + .def( + "channel_status", + &EventStream::channel_status, + pydoc::doc("EventStream::channel_status") + ) + .def( + "channel_probabilities", + &EventStream::channel_probabilities, + pydoc::doc("EventStream::channel_probabilities") + ) + .def( + "event_count", + &EventStream::event_count, + pydoc::doc("EventStream::event_count") + ) + .def( + "max_particle_count", + &EventStream::max_particle_count, + pydoc::doc("EventStream::max_particle_count") + ); + py::classh logger(m, "Logger", pydoc::doc("Logger")); add_enum( logger, @@ -3780,5 +3905,12 @@ PYBIND11_MODULE(_madspace_py, m) { } }; EventGenerator::set_abort_check_function(abort_check_function); + // EventStream waits with the GIL released + EventStream::set_abort_check_function([] { + py::gil_scoped_acquire gil; + if (PyErr_CheckSignals() != 0) { + throw py::error_already_set(); + } + }); MadnisTraining::set_abort_check_function(abort_check_function); } diff --git a/madspace/tests/test_event_stream.py b/madspace/tests/test_event_stream.py new file mode 100644 index 0000000000..4e75b35bee --- /dev/null +++ b/madspace/tests/test_event_stream.py @@ -0,0 +1,241 @@ +"""EventStream: in-memory event generation from channels with a fixed maximum +weight, as stored in a fix_max_weight gridpack. + +The channels of an e+ e- -> d d~ g process (flat matrix element, so the weights +only vary through the phase-space mapping) are surveyed and their maximum +weights fixed with EventGenerator.fix_max_weights(), then reloaded from their +JSON like a gridpack does. +""" + +import json +import os + +import numpy as np +import pytest + +import madspace as ms + +CM_ENERGY = 1000.0 +PIDS = [11, -11, 1, -1, 21] +SEED = 1234 + + +def build_topology(): + propagators = [ + ms.Propagator( + mass=0.0, width=0.0, integration_order=0, e_min=0.0, e_max=0.0, pdg_id=pid + ) + for pid in (22, -1) + ] + diagram = ms.Diagram( + [0.0, 0.0], + [0.0, 0.0, 0.0], + propagators, + [["i0", "i1", "p0"], ["o0", "p1", "p0"], ["o1", "o2", "p1"]], + ) + return ms.Topology.topologies(diagram)[0] + + +TOPOLOGY = build_topology() +PERMUTATIONS = [[0, 1, 2, 3, 4]] + + +def lhe_completer(): + subproc_args = ms.SubprocArgs( + process_id=0, + topologies=[TOPOLOGY], + permutations=[PERMUTATIONS], + diagram_indices=[[0]], + diagram_color_indices=[[[0]]], + color_flows=[[[0, 0], [0, 0], [501, 0], [0, 502], [502, 501]]], + pdg_color_types={11: 1, -11: 1, 1: 3, -1: -3, 22: 1, 21: 8}, + helicities=[[1, -1, 1, -1, 1]], + pdg_ids=[[PIDS]], + ) + return ms.LHECompleter([subproc_args], bw_cutoff=15.0) + + +def make_config(): + config = ms.GeneratorConfig() + config.cpu_batch_size = 1000 + config.freeze_max_weight_after = 1000 + config.max_overweight_truncation = 0.05 + return config + + +def make_context(run_dir, thread_count): + context = ms.Context(ms.cpu_device(), thread_count) + # RunningCoupling only reads these two keys of an LHAPDF .info file + alphas_file = os.path.join(run_dir, "alphas.info") + with open(alphas_file, "w") as f: + f.write("AlphaS_Qs: [1.0, 10.0, 100.0, 1000.0, 10000.0]\n") + f.write("AlphaS_Vals: [0.3, 0.2, 0.12, 0.1, 0.08]\n") + ms.AlphaSGrid(alphas_file).initialize_globals(context) + return context, alphas_file + + +def make_channel(context, alphas_file, run_dir, name, pt_min): + cuts = ms.Cuts( + [ + ms.CutItem(ms.Observable(PIDS, ms.Observable.obs_pt, [[i]]), min=pt_min) + for i in (2, 3, 4) + ] + ) + mapping = ms.PhaseSpaceMapping( + TOPOLOGY, CM_ENERGY, permutations=PERMUTATIONS, leptonic=True, cuts=cuts + ) + # 0xBADCAFE: the built-in flat test matrix element, no library needed + matrix_element = ms.MatrixElement( + 0xBADCAFE, + len(PIDS), + ms.Integrand.matrix_element_inputs, + ms.Integrand.matrix_element_outputs, + 1, + True, + ) + diff_xs = ms.DifferentialCrossSection( + matrix_element=matrix_element, + cm_energy=CM_ENERGY, + running_coupling=None, + energy_scale=ms.CachedScale(), + ) + integrand = ms.Integrand( + mapping, + [diff_xs], + running_coupling=ms.RunningCoupling(ms.AlphaSGrid(alphas_file)), + energy_scale=ms.EnergyScale(len(PIDS), 91.188), + channel_indices=[0], + ) + return ms.ChannelEventGenerator( + contexts=[context], + integrand=integrand, + event_file=os.path.join(run_dir, f"events.{name}.npy"), + weight_file=os.path.join(run_dir, f"weights.{name}.npy"), + config=make_config(), + subprocess_index=0, + name=name, + histograms=None, + ) + + +@pytest.fixture(scope="module") +def gridpack(tmp_path_factory): + """Channel JSONs with fixed max weights and the globals, like a gridpack.""" + run_dir = str(tmp_path_factory.mktemp("event_stream")) + context, alphas_file = make_context(run_dir, 4) + channels = [ + make_channel(context, alphas_file, run_dir, "a", 20.0), + make_channel(context, alphas_file, run_dir, "b", 60.0), + ] + generator = ms.EventGenerator([context], channels, seed=1, config=make_config()) + generator.survey() + generator.fix_max_weights() + globals_file = os.path.join(run_dir, "globals.tar") + context.save_globals(globals_file) + return { + "run_dir": run_dir, + "channels": [channel.to_json(True) for channel in channels], + "globals": globals_file, + "integral": generator.status().mean, + } + + +def make_stream(gridpack, thread_count=4, seed=SEED): + context = ms.Context(ms.cpu_device(), thread_count) + context.load_globals(gridpack["globals"]) + config = make_config() + channels = [ + ms.ChannelEventGenerator.load_json(channel, [context], config=config) + for channel in gridpack["channels"] + ] + return ms.EventStream(context, channels, seed, lhe_completer(), config) + + +def stream_batches(stream, batch_sizes): + batches = [stream.next_batch(size) for size in batch_sizes] + return {key: np.concatenate([b[key] for b in batches]) for key in batches[0]} + + +def test_reproducible_across_batch_sizes_and_threads(gridpack): + reference = stream_batches(make_stream(gridpack, 4), [3000]) + for thread_count, batch_sizes in [(1, [1000, 1000, 1000]), (3, [1, 999, 2000])]: + events = stream_batches(make_stream(gridpack, thread_count), batch_sizes) + for key, values in reference.items(): + np.testing.assert_array_equal(events[key], values, err_msg=key) + + +def test_seed_changes_events(gridpack): + events1 = make_stream(gridpack, seed=1).next_batch(100) + events2 = make_stream(gridpack, seed=2).next_batch(100) + assert not np.array_equal(events1["px"], events2["px"]) + + +def test_lhe_events_match_batches(gridpack): + lhe_events = make_stream(gridpack).next_events(50) + batch = make_stream(gridpack).next_batch(50) + for i, event in enumerate(lhe_events): + assert event.weight == batch["weight"][i] + assert [p.pdg_id for p in event.particles] == list(batch["pdg_id"][i]) + assert [p.px for p in event.particles] == list(batch["px"][i]) + + +def test_weights_and_particles(gridpack): + stream = make_stream(gridpack) + batch = stream.next_batch(5000) + weights = batch["weight"] + # weights are 1, except for overweight events + assert np.all(weights >= 1.0) + assert np.mean(weights == 1.0) > 0.9 + assert np.all(batch["particle_count"] == len(PIDS)) + assert set(map(tuple, batch["pdg_id"])) == {tuple(PIDS)} + assert batch["px"].shape == (5000, stream.max_particle_count()) + # momentum conservation + outgoing = batch["status_code"] == 1 + assert np.allclose(np.sum(batch["px"] * outgoing, axis=1), 0.0, atol=1e-6) + + +def test_status(gridpack): + stream = make_stream(gridpack) + stream.next_batch(2000) + status = stream.status() + assert stream.event_count() == 2000 + assert status.count_unweighted == 2000 + assert status.count > 0 + assert status.mean == pytest.approx(gridpack["integral"], rel=0.1) + channel_status = stream.channel_status() + assert sum(s.count_unweighted for s in channel_status) == 2000 + assert sum(stream.channel_probabilities()) == pytest.approx(1.0) + + +def test_channel_fractions_follow_probabilities(gridpack): + stream = make_stream(gridpack) + stream.next_batch(20000) + probs = np.array(stream.channel_probabilities()) + counts = np.array([s.count_unweighted for s in stream.channel_status()]) + fractions = counts / counts.sum() + assert np.allclose(fractions, probs, atol=0.02) + + +def test_no_files_written(gridpack, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + stream = make_stream(gridpack) + stream.next_batch(1000) + stream.close() + assert os.listdir(tmp_path) == [] + + +def test_requires_fixed_max_weight(gridpack): + context = ms.Context(ms.cpu_device(), 1) + context.load_globals(gridpack["globals"]) + config = make_config() + channels = [] + for channel in gridpack["channels"]: + data = json.loads(channel) + del data["max_weight"] + channels.append( + ms.ChannelEventGenerator.load_json( + json.dumps(data), [context], config=config + ) + ) + with pytest.raises(ValueError, match="fixed maximum weight"): + ms.EventStream(context, channels, SEED, lhe_completer(), config) From fa1aaaf1b2083ebb30a43096472d78ac7d69bae4 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 12 Oct 2026 11:19:24 +0200 Subject: [PATCH 02/88] fix CI: keyword-only args in ChannelEventGenerator.load_json binding (valid .pyi stub), gridpack section in launch test run card --- madspace/src/python/madspace.cpp | 1 + tests/unit_tests/interface/test_mg7_launch.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 3d61b7eca4..537d4c5844 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -3151,6 +3151,7 @@ PYBIND11_MODULE(_madspace_py, m) { &ChannelEventGenerator::load_json, py::arg("channel_json"), py::arg("contexts"), + py::kw_only(), py::arg("event_file") = "", py::arg("weight_file") = "", py::arg("config"), diff --git a/tests/unit_tests/interface/test_mg7_launch.py b/tests/unit_tests/interface/test_mg7_launch.py index 9eb5e49aaa..8f7fd39788 100644 --- a/tests/unit_tests/interface/test_mg7_launch.py +++ b/tests/unit_tests/interface/test_mg7_launch.py @@ -727,7 +727,10 @@ def test_status_file_path_is_absolute(self): try: process = self.launch.MadgraphProcess.__new__( self.launch.MadgraphProcess) - process.run_card = {"run": {"run_name": "run"}} + process.run_card = { + "run": {"run_name": "run"}, + "gridpack": {"output_dir": "", "temp_output_dir": ""}, + } process.init_event_dir() finally: os.chdir(cwd) From 18e0283b87e202691a47babf69b20445afaa122a Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 12 Oct 2026 11:41:18 +0200 Subject: [PATCH 03/88] docs: event streaming page with a CI-tested p p > e+ e- m_ee example and the EventStream API reference --- .github/workflows/docs_examples.yml | 9 +- docs/source/conf.py | 11 ++ docs/source/event_stream.rst | 151 ++++++++++++++++++ docs/source/gridpacks.rst | 24 +-- docs/source/index.rst | 1 + docs/source/madspace/cpp-api.rst | 1 + docs/source/madspace/python-api.rst | 5 + .../iolibs/template_files/mg7/event_stream.py | 23 +++ .../test_madspace_examples.py | 70 +++++++- 9 files changed, 268 insertions(+), 27 deletions(-) create mode 100644 docs/source/event_stream.rst diff --git a/.github/workflows/docs_examples.yml b/.github/workflows/docs_examples.yml index 3259fe375b..5138407ccf 100644 --- a/.github/workflows/docs_examples.yml +++ b/.github/workflows/docs_examples.yml @@ -3,13 +3,16 @@ name: MadSpace docs examples # The other MG7 acceptance workflows all carry `paths-ignore: ['docs/**']`, so a # docs-only commit currently triggers none of them. This workflow does the # opposite: it triggers ONLY on the paths that can affect the documentation -# examples, so editing docs/source/madspace/examples/*.rst is actually tested. +# examples, so editing docs/source/madspace/examples/*.rst (or the event +# streaming example in docs/source/event_stream.rst) is actually tested. on: push: branches: - '**' paths: - 'docs/source/madspace/examples/**' + - 'docs/source/event_stream.rst' + - 'madgraph/iolibs/template_files/mg7/**' - 'madspace/**' - 'tests/acceptance_tests/test_madspace_examples.py' - '.github/workflows/docs_examples.yml' @@ -18,6 +21,8 @@ on: types: [opened, synchronize, reopened] paths: - 'docs/source/madspace/examples/**' + - 'docs/source/event_stream.rst' + - 'madgraph/iolibs/template_files/mg7/**' - 'madspace/**' - 'tests/acceptance_tests/test_madspace_examples.py' - '.github/workflows/docs_examples.yml' @@ -55,4 +60,4 @@ jobs: - name: run the madspace doc examples run: | cd $GITHUB_WORKSPACE - ./tests/test_manager.py test_madspace_example_simple_mappings test_madspace_example_cuts test_madspace_example_diagram_mapping test_madspace_example_integration_order test_madspace_example_matrix_element test_madspace_example_pdf test_madspace_example_integrator test_madspace_example_integrator_madnis test_madspace_example_flow_training -pA -t0 -l INFO + ./tests/test_manager.py test_madspace_example_simple_mappings test_madspace_example_cuts test_madspace_example_diagram_mapping test_madspace_example_integration_order test_madspace_example_matrix_element test_madspace_example_pdf test_madspace_example_integrator test_madspace_example_integrator_madnis test_madspace_example_flow_training test_gridpack_example_event_stream -pA -t0 -l INFO diff --git a/docs/source/conf.py b/docs/source/conf.py index 541cbea840..38a93e0097 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -14,6 +14,17 @@ if os.path.isdir(os.path.join(_local_install, "madspace")): sys.path.insert(0, _local_install) +# bin/event_stream.py of the gridpacks, for the API reference in event_stream.rst +sys.path.insert( + 0, + os.path.abspath( + os.path.join( + os.path.dirname(__file__), + "..", "..", "madgraph", "iolibs", "template_files", "mg7", + ) + ), +) + # Configuration file for the Sphinx documentation builder. # # For the full list of built-in configuration values, see the documentation: diff --git a/docs/source/event_stream.rst b/docs/source/event_stream.rst new file mode 100644 index 0000000000..306c7b14ca --- /dev/null +++ b/docs/source/event_stream.rst @@ -0,0 +1,151 @@ +Streaming events from Python +============================ + +A gridpack can generate events directly into a Python program, without writing any files. +This is useful when the events are consumed right away, for example by an analysis written +in Python or as training data for a neural network. The ``EventStream`` class in the +gridpack's ``bin/event_stream.py`` returns the events either as LHE event objects or as +batches of NumPy arrays or PyTorch tensors. + +Event streaming needs a gridpack created with ``run_mode = "fix_max_weight"`` (see +:ref:`gridpack-run-mode`). In this mode, the maximum weight of every channel is fixed when +the gridpack is created. An event is final as soon as it is unweighted, so it can be returned +right away without a survey or a combination step. + +Example: the e+ e- invariant mass +--------------------------------- + +This example fills a histogram of the invariant mass of the electron pair in +``p p > e+ e-``. First, create the gridpack at the MadGraph7 prompt: + +.. code-block:: text + :name: event-stream-mg5-commands + + generate p p > e+ e- + output mg7 PROC_ppee + launch PROC_ppee + set save_gridpack true + set run_mode fix_max_weight + done + +The gridpack is written to ``PROC_ppee/Events/run_01/gridpack``. Its ``bin`` directory +contains ``event_stream.py``, so add it to the Python path and import ``EventStream``. The +script can run from any directory. + +.. code-block:: python + + import sys + + import numpy as np + + sys.path.insert(0, "PROC_ppee/Events/run_01/gridpack/bin") + from event_stream import EventStream + +Creating an ``EventStream`` starts the event generation in background threads. Its +arguments, like the device, the seed and the batch sizes, default to the values in the +gridpack's ``Cards/grid_run_card.toml``. The stream is endless, so ``max_events`` limits +the number of events. Each batch is an :py:class:`EventBatch ` +with one array per LHE field. The particle fields have the shape +``(batch_size, max_particle_count)``. + +.. code-block:: python + + bins = np.linspace(60.0, 120.0, 31) + hist = np.zeros(len(bins) - 1) + weight_sum = 0.0 + + with EventStream(seed=42) as stream: + for batch in stream.batches(10000, max_events=100000): + # four-momentum of the e+ e- pair: sum over the outgoing electrons + is_electron = (np.abs(batch.pdg_id) == 11) & (batch.status_code == 1) + e, px, py, pz = ( + np.sum(np.where(is_electron, p, 0.0), axis=1) + for p in (batch.energy, batch.px, batch.py, batch.pz) + ) + m_ee = np.sqrt(np.maximum(e**2 - px**2 - py**2 - pz**2, 0.0)) + hist += np.histogram(m_ee, bins, weights=batch.weight)[0] + weight_sum += np.sum(batch.weight) + cross_section = stream.integral + print(f"cross section: {cross_section:.1f} +- {stream.integral_error:.1f} pb") + +The event weights are ``1``, except for overweight events, so they have to be normalized. +Scaling them to the cross section gives the differential cross section in pb/GeV: + +.. code-block:: python + + dsigma = hist * cross_section / weight_sum / np.diff(bins) + for low, value in zip(bins[:-1], dsigma): + print(f"{low:5.0f} GeV {value:8.3f} pb/GeV " + "#" * int(50 * value / dsigma.max())) + + peak = np.argmax(dsigma) + assert bins[peak] < 91.2 < bins[peak + 1], "the Z peak is missing" + +The output shows the Z peak (shortened):: + + cross section: 880.2 +- 0.6 pb + ... + 84 GeV 7.569 pb/GeV ## + 86 GeV 16.543 pb/GeV ##### + 88 GeV 51.287 pb/GeV ################ + 90 GeV 158.376 pb/GeV ################################################## + 92 GeV 66.743 pb/GeV ##################### + 94 GeV 18.457 pb/GeV ##### + ... + +Output formats +-------------- + +There are three ways to get events from a stream: + +- Iterating over the stream yields one ``madspace.LHEEvent`` at a time. The events are + completed in chunks of 1000 internally. +- :py:meth:`next_events ` returns a list of a given + number of ``LHEEvent`` objects. +- :py:meth:`batches ` and + :py:meth:`next_batch ` return + :py:class:`EventBatch ` objects. With ``output_type="torch"``, + their fields are PyTorch tensors on the CPU instead of NumPy arrays. + +Events have different numbers of particles if the process has several subprocesses or +includes decays. The particle fields are then padded with zeros. ``particle_count`` holds +the number of particles of every event and ``particle_mask`` marks the valid entries. + +Weights and cross section +------------------------- + +Every event weight is ``1``, or ``-1`` for an event with a negative weight. The exception are +overweight events, whose weight exceeded the fixed maximum weight of their channel. They +keep the ratio of their weight to the maximum weight, so their weight is larger than ``1``. +Treating them as ``1`` would bias the distributions in the phase-space regions where they +occur. + +To normalize a sample to the cross section, scale its weights by +``stream.integral / sum(weights)``. The cross section estimate improves while events are +generated: it combines the integrals stored in the gridpack with the weighted events +generated by the stream so far. Read it after generating the sample, as in the example +above. + +How it works +------------ + +Each channel of the gridpack keeps a buffer of unweighted events, which background threads +refill whenever it runs low. The channel of every event is drawn at random, with +probabilities given by the current cross section estimates of the channels. Generating +the phase-space points and evaluating the matrix elements happen in separate steps, as in a +normal gridpack run. + +The events are reproducible: for a given seed, the stream always returns the same sequence of +events. It does not depend on the number of threads or on how the events are split into +batches. + +API reference +------------- + +.. autoclass:: event_stream.EventStream + :members: + :special-members: __iter__ + :member-order: bysource + +.. autoclass:: event_stream.EventBatch + :members: + :member-order: bysource diff --git a/docs/source/gridpacks.rst b/docs/source/gridpacks.rst index 41056de9d6..39c1526857 100644 --- a/docs/source/gridpacks.rst +++ b/docs/source/gridpacks.rst @@ -102,28 +102,8 @@ Streaming events from Python ---------------------------- A gridpack created with ``run_mode = "fix_max_weight"`` also contains -``bin/event_stream.py``. It generates events in memory, without writing any files: - -.. code-block:: python - - import sys - sys.path.insert(0, "/bin") - from event_stream import EventStream - - with EventStream(device="cpu", seed=42) as stream: - for batch in stream.batches(10000, output_type="torch", max_events=100000): - batch.px # shape (10000, max_particle_count) - -The stream is endless. Iterating over ``stream`` yields ``LHEEvent`` objects, and -``batches`` yields ``EventBatch`` dataclasses with numpy arrays or torch tensors, padded to -the largest particle count (see ``particle_mask``). The constructor takes the device, the -seed, the batch sizes and the thread pool size, with defaults from ``grid_run_card.toml``. -Background threads keep a buffer of events for every channel, and every event picks its -channel at random according to the current cross section estimates. ``stream.integral`` -and ``stream.integral_error`` give the current estimate. - -The weights are ``1`` (``-1`` for negative events), except for overweight events. The -events only depend on the seed, not on the number of threads or the batch sizes. +``bin/event_stream.py``. Its ``EventStream`` class generates events in memory and returns +them to a Python program, without writing any files. See :doc:`event_stream`. Run card settings ----------------- diff --git a/docs/source/index.rst b/docs/source/index.rst index 3c597d70cb..f36b5d8b37 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -28,6 +28,7 @@ If you need a stable release, use the first_steps run_card gridpacks + event_stream madevent_to_mg7 .. toctree:: diff --git a/docs/source/madspace/cpp-api.rst b/docs/source/madspace/cpp-api.rst index 8aad6b8723..0beb59ec56 100644 --- a/docs/source/madspace/cpp-api.rst +++ b/docs/source/madspace/cpp-api.rst @@ -125,6 +125,7 @@ The high-level driver plus the tensor and device layer it runs on: integration, cpp/EventHistogramSpec cpp/EventHistograms cpp/EventRecord + cpp/EventStream cpp/FieldLayout cpp/GeneratorBatchJob cpp/GeneratorConfig diff --git a/docs/source/madspace/python-api.rst b/docs/source/madspace/python-api.rst index ec5d0ed19c..e8cb534d40 100644 --- a/docs/source/madspace/python-api.rst +++ b/docs/source/madspace/python-api.rst @@ -40,6 +40,7 @@ Builders that assemble compute-graph functions for channel weights, energy scale .. toctree:: :maxdepth: 1 + python/BufferUnweighter python/CachedPdf python/CachedScale python/ChannelWeightNetwork @@ -122,6 +123,7 @@ The high-level driver plus the tensor and device layer it runs on: integration, python/EventGenerator python/EventHistogramSpec python/EventHistograms + python/EventStream python/GeneratorConfig python/GeneratorStatus python/Histogram @@ -177,3 +179,6 @@ Lower-level helpers and containers. python/NamedTypes python/NamedValues + python/get_stream + python/set_stream + python/stream diff --git a/madgraph/iolibs/template_files/mg7/event_stream.py b/madgraph/iolibs/template_files/mg7/event_stream.py index 1e5e843ad7..3a24ab9c97 100644 --- a/madgraph/iolibs/template_files/mg7/event_stream.py +++ b/madgraph/iolibs/template_files/mg7/event_stream.py @@ -59,25 +59,45 @@ class EventBatch: ``particle_mask``. numpy arrays or CPU torch tensors, depending on the ``output_type`` of :meth:`EventStream.batches`.""" + #: Process ID of the event (IDPRUP). process_id: Any + #: Event weight, +-1 or larger in magnitude for overweight events (XWGTUP). weight: Any + #: Renormalization scale in GeV (SCALUP). scale: Any + #: QED coupling (AQEDUP). alpha_qed: Any + #: QCD coupling (AQCDUP). alpha_qcd: Any + #: Number of particles of the event, including resonances (NUP). particle_count: Any + #: True for the particle slots below ``particle_count``. particle_mask: Any + #: PDG ID of the particle (IDUP). pdg_id: Any + #: -1 for incoming, 1 for outgoing particles, 2 for resonances (ISTUP). status_code: Any + #: Index of the first mother, counting from 1; 0 if there is none (MOTHUP). mother1: Any + #: Index of the second mother, counting from 1; 0 if there is none (MOTHUP). mother2: Any + #: Color line, 0 for a color singlet (ICOLUP). color: Any + #: Anti-color line (ICOLUP). anti_color: Any + #: x component of the momentum in GeV. px: Any + #: y component of the momentum in GeV. py: Any + #: z component of the momentum in GeV. pz: Any + #: Energy in GeV. energy: Any + #: Mass in GeV. mass: Any + #: Proper lifetime c*tau in mm, 0 if not tracked (VTIMUP). lifetime: Any + #: Cosine of the spin angle, or 9 if the spin is not specified (SPINUP). spin: Any def __len__(self) -> int: @@ -106,6 +126,9 @@ class EventStream: gridpack was made with (can lead to errors or incorrect results). """ + #: The seed of the stream; a seed of -1 is replaced by the random seed drawn. + seed: int + def __init__( self, device: Optional[str] = None, diff --git a/tests/acceptance_tests/test_madspace_examples.py b/tests/acceptance_tests/test_madspace_examples.py index 7983ba7b14..ef95ab345b 100644 --- a/tests/acceptance_tests/test_madspace_examples.py +++ b/tests/acceptance_tests/test_madspace_examples.py @@ -55,7 +55,8 @@ # in the subprocess env run_doc_example() builds. os.environ.setdefault('KMP_DUPLICATE_LIB_OK', 'TRUE') -_EXAMPLES_DIR = pjoin(MG5DIR, 'docs', 'source', 'madspace', 'examples') +_DOCS_DIR = pjoin(MG5DIR, 'docs', 'source') +_EXAMPLES_DIR = pjoin(_DOCS_DIR, 'madspace', 'examples') _MADSPACE_INSTALL = pjoin(MG5DIR, 'madspace', 'install') _CODE_BLOCK_RE = re.compile(r'^(?P[ \t]*)\.\.[ \t]+code-block::[ \t]*python[ \t]*$') @@ -106,6 +107,30 @@ def extract_python_blocks(rst_path): return blocks +def extract_named_block(rst_path, name): + """Return the dedented body of the code block whose directive carries + ``:name: ``, or None if there is none.""" + with open(rst_path) as f: + lines = f.read().splitlines() + for i, line in enumerate(lines): + if line.strip() != ':name: %s' % name: + continue + j = i + while not lines[j].lstrip().startswith('.. code-block::'): + j -= 1 + base_indent = len(lines[j]) - len(lines[j].lstrip()) + k = i + 1 + while k < len(lines) and lines[k].strip().startswith(':'): + k += 1 + body = [] + for line in lines[k:]: + if line.strip() and len(line) - len(line.lstrip()) <= base_indent: + break + body.append(line) + return textwrap.dedent('\n'.join(body)).strip() + return None + + def has_madspace(): try: import madspace @@ -147,7 +172,7 @@ def has_madnis(): return False -def run_doc_example(test, page_name, cwd=None): +def run_doc_example(test, page_name, cwd=None, examples_dir=_EXAMPLES_DIR): """Extract the python blocks of ``examples/.rst``, concatenate them into one script, and run it as a subprocess. Fails the test if the script raises or if the page has no python blocks at all (a moved or @@ -157,7 +182,7 @@ def run_doc_example(test, page_name, cwd=None): ``ggttg_process_dir`` below) run with that directory as its working directory, matching what the page's shell commands build. Without it, a fresh empty temporary directory is used and cleaned up afterwards.""" - rst_path = pjoin(_EXAMPLES_DIR, page_name + '.rst') + rst_path = pjoin(examples_dir, page_name + '.rst') if not os.path.isfile(rst_path): test.fail('documentation page not found: %s' % rst_path) @@ -367,5 +392,44 @@ def test_madspace_example_flow_training(self): run_doc_example(self, 'flow-training', cwd=scratch) +class TestGridpackExamples(unittest.TestCase): + + def test_gridpack_example_event_stream(self): + """docs/source/event_stream.rst -- creating a fix_max_weight gridpack + for p p > e+ e- with the MadGraph commands of the page, then filling + an m_ee histogram from its bin/event_stream.py. Needs madspace, numpy, + a C++ compiler and lhapdf/NNPDF40_lo_as_01180.""" + if not has_madspace(): + self.skipTest('madspace unavailable') + if not has_pdf_set('NNPDF40_lo_as_01180'): + self.skipTest('NNPDF40_lo_as_01180 LHAPDF data not found ' + '(set $LHAPDF_DATA_PATH)') + rst_path = pjoin(_DOCS_DIR, 'event_stream.rst') + commands = extract_named_block(rst_path, 'event-stream-mg5-commands') + self.assertTrue(commands, 'MadGraph commands not found in %s' % rst_path) + + with tempfile.TemporaryDirectory(prefix='doc_event_stream_') as scratch: + script = pjoin(scratch, 'mg5_commands.txt') + with open(script, 'w') as f: + f.write('set automatic_html_opening False --no_save\n') + f.write(commands + '\n') + env = dict(os.environ) + if os.path.isdir(_MADSPACE_INSTALL): + env['PYTHONPATH'] = os.pathsep.join( + [_MADSPACE_INSTALL] + + ([env['PYTHONPATH']] if env.get('PYTHONPATH') else []) + ) + result = subprocess.run( + [sys.executable, pjoin(MG5DIR, 'bin', 'madgraph'), '-f', script], + cwd=scratch, env=env, capture_output=True, text=True, + ) + self.assertEqual( + result.returncode, 0, + 'MadGraph commands failed\n--- stdout ---\n%s\n--- stderr ---\n%s' + % (result.stdout[-5000:], result.stderr[-5000:]), + ) + run_doc_example(self, 'event_stream', cwd=scratch, examples_dir=_DOCS_DIR) + + if __name__ == '__main__': unittest.main() From 13bdbe384ab921d36a71b99ba3251501301347e5 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 11:17:34 +0200 Subject: [PATCH 04/88] madspace: optional matrix-element outputs in the Integrand The Integrand looks up the matrix-element outputs by name. Only |M|^2 is required: diagram_amp2 only for diagram-based channel weights, and missing color, helicity and diagram indices are written as -1. With return_matrix_element, DifferentialCrossSection also returns |M|^2 itself, which the Integrand passes on as the matrix_element output. --- .../madspace/phasespace/cross_section.hpp | 10 +- .../include/madspace/phasespace/integrand.hpp | 11 +- madspace/src/phasespace/cross_section.cpp | 21 ++- madspace/src/phasespace/integrand.cpp | 96 +++++++++- madspace/src/python/madspace.cpp | 2 + madspace/tests/test_integrand_me_outputs.py | 175 ++++++++++++++++++ 6 files changed, 297 insertions(+), 18 deletions(-) create mode 100644 madspace/tests/test_integrand_me_outputs.py diff --git a/madspace/include/madspace/phasespace/cross_section.hpp b/madspace/include/madspace/phasespace/cross_section.hpp index 6143b95f30..498e646d9e 100644 --- a/madspace/include/madspace/phasespace/cross_section.hpp +++ b/madspace/include/madspace/phasespace/cross_section.hpp @@ -35,6 +35,8 @@ namespace madspace { * **Returns** * - the return values of @p matrix_element, scaled to the differential cross * section. + * - `raw_matrix_element` – `float`, shape `(batch,)` – the squared matrix + * element without PDFs and flux (with @p return_matrix_element). * * **References** * - [1] T. Heimel, O. Mattelaer, R. Winterhalder, "MadSpace", @@ -62,6 +64,8 @@ class DifferentialCrossSection : public FunctionGenerator { * rather than reconstructed. * @param decay If true, describe the decay of a single * incoming particle instead of a collision. + * @param return_matrix_element If true, also return the squared matrix + * element itself as `raw_matrix_element`. */ DifferentialCrossSection( const MatrixElement& matrix_element, @@ -73,7 +77,8 @@ class DifferentialCrossSection : public FunctionGenerator { const std::variant& pdf1 = std::monostate{}, const std::variant& pdf2 = std::monostate{}, bool input_momentum_fraction = true, - bool decay = false + bool decay = false, + bool return_matrix_element = false ); /// The allowed initial-parton flavour combinations. @@ -86,6 +91,8 @@ class DifferentialCrossSection : public FunctionGenerator { const std::optional& running_coupling() const { return _running_coupling; } + /// Whether the squared matrix element is returned as `raw_matrix_element`. + bool returns_matrix_element() const { return _return_matrix_element; } private: NamedVector build_function_impl( @@ -101,6 +108,7 @@ class DifferentialCrossSection : public FunctionGenerator { std::variant _energy_scale; bool _input_momentum_fraction; bool _decay; + bool _return_matrix_element; }; } // namespace madspace diff --git a/madspace/include/madspace/phasespace/integrand.hpp b/madspace/include/madspace/phasespace/integrand.hpp index 7cb088358a..7bd5dbd1ed 100644 --- a/madspace/include/madspace/phasespace/integrand.hpp +++ b/madspace/include/madspace/phasespace/integrand.hpp @@ -41,14 +41,21 @@ namespace madspace { * - `weight` – `float`, shape `(batch,)` – the event weight. * - `momenta` – `float`, shape `(batch, particle_count, 4)` – the momenta. * - `color_index`, `helicity_index`, `diagram_index`, `flavor_index` – `int`, - * shape `(batch,)` – the sampled discrete quantities. + * shape `(batch,)` – the sampled discrete quantities. Indices the matrix + * element does not return are -1. * - `ren_scale`, `alpha_qcd`, `x1`, `x2`, `fact_scale1`, `fact_scale2` – * `float`, shape `(batch,)` – scales and momentum fractions. * - `random` – `float`, shape `(batch, mapping.random_dim())` – the underlying * random numbers. + * - `matrix_element` – `float`, shape `(batch,)` – the squared matrix element, + * present only if the differential cross sections return it. * - `bias` – `float`, shape `(batch,)` – the event bias, present only with a * custom bias. `weight` then includes the bias factor. * + * Only `matrix_element_out` is required from the matrix element. + * `diagram_amp2_out` is needed for diagram-based channel weights, i.e. with + * several channels and no @p prop_chan_weights. + * * **References** * - [1] T. Heimel, O. Mattelaer, R. Winterhalder, "MadSpace", * https://arxiv.org/abs/2602.06895 @@ -73,7 +80,7 @@ class Integrand : public FunctionGenerator { MatrixElement::random_helicity_in, MatrixElement::random_diagram_in, }; - /// The @ref MatrixElement outputs an `Integrand` requests. + /// The @ref MatrixElement outputs an `Integrand` can use. inline static const std::vector matrix_element_outputs = { MatrixElement::matrix_element_out, diff --git a/madspace/src/phasespace/cross_section.cpp b/madspace/src/phasespace/cross_section.cpp index f5f46b31cd..40b4423b49 100644 --- a/madspace/src/phasespace/cross_section.cpp +++ b/madspace/src/phasespace/cross_section.cpp @@ -11,7 +11,8 @@ DifferentialCrossSection::DifferentialCrossSection( const std::variant& pdf1, const std::variant& pdf2, bool input_momentum_fraction, - bool decay + bool decay, + bool return_matrix_element ) : FunctionGenerator( "DifferentialCrossSection", @@ -50,7 +51,13 @@ DifferentialCrossSection::DifferentialCrossSection( } return arg_types; }(), - matrix_element.return_types() + [&] { + NamedVector ret_types = matrix_element.return_types(); + if (return_matrix_element) { + ret_types.push_back("raw_matrix_element", batch_float); + } + return ret_types; + }() ), _pid_options(pid_options), _matrix_element(matrix_element), @@ -62,7 +69,8 @@ DifferentialCrossSection::DifferentialCrossSection( _e_cm(cm_energy), _energy_scale(energy_scale), _input_momentum_fraction(input_momentum_fraction), - _decay(decay) { + _decay(decay), + _return_matrix_element(return_matrix_element) { if (decay) { if (_has_pdf.at(0) || _has_pdf.at(1)) { throw std::invalid_argument("a decay has no beams, so no PDFs"); @@ -157,14 +165,14 @@ NamedVector DifferentialCrossSection::build_function_impl( throw std::runtime_error("matrix element missing in return values"); } std::size_t me_index = search - _matrix_element.outputs().begin(); + Value raw_matrix_element = me_result.at(me_index); if (_decay) { // Differential decay rate: dGamma = |M|^2 dPhi_n / (2 M), in GeV. The // flux is a compile-time constant here (M is fixed by the decaying // particle), so this is a single multiply and needs no dedicated // instruction. Deliberately no GeV^-2 -> pb conversion: a width is not // a cross section. - me_result.at(me_index) = - fb.mul(me_result.at(me_index), 1. / (2. * _e_cm)); + me_result.at(me_index) = fb.mul(me_result.at(me_index), 1. / (2. * _e_cm)); } else { me_result.at(me_index) = fb.diff_cross_section( x1x2.at(0), @@ -175,5 +183,8 @@ NamedVector DifferentialCrossSection::build_function_impl( _e_cm * _e_cm ); } + if (_return_matrix_element) { + me_result.push_back("raw_matrix_element", raw_matrix_element); + } return me_result; } diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index b05999051c..667bb948f5 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -3,6 +3,7 @@ #include "madspace/util.hpp" #include +#include using namespace madspace; @@ -30,6 +31,16 @@ std::size_t final_channel_count( } } +/// Position of `output` in the return values of a differential cross section +/// built on `matrix_element`, or -1 if it is not returned. +int me_output_index( + const MatrixElement& matrix_element, MatrixElement::MatrixElementOutput output +) { + auto& outputs = matrix_element.outputs(); + auto search = std::find(outputs.begin(), outputs.end(), output); + return search == outputs.end() ? -1 : static_cast(search - outputs.begin()); +} + } // namespace static const BatchSize acc_batch_size("acc_batch_size"); @@ -150,6 +161,9 @@ Integrand::Integrand( ret_types.push_back("partial_weight_product", batch_float); } } + if (diff_xs_first.returns_matrix_element()) { + ret_types.push_back("matrix_element", batch_float); + } // before "random": the outputs after it feed the discrete histograms if (custom_bias) { ret_types.push_back("bias", batch_float); @@ -265,6 +279,41 @@ Integrand::Integrand( _flavor_mirror.push_back(mirror ? 2 : 1); _has_mirror |= mirror; } + auto& me_first = diff_xs.at(0).matrix_element(); + for (auto& dxs : diff_xs) { + if (dxs.matrix_element().outputs() != me_first.outputs() || + dxs.matrix_element().external_inputs() != me_first.external_inputs() || + dxs.returns_matrix_element() != diff_xs.at(0).returns_matrix_element()) { + throw std::invalid_argument( + "all matrix elements of an integrand need the same inputs and outputs" + ); + } + } + for (auto input : me_first.external_inputs()) { + if (input != MatrixElement::momenta_in && input != MatrixElement::flavor_in && + input != MatrixElement::alpha_s_in) { + throw std::invalid_argument( + "the integrand only supplies momenta, flavor and alpha_s to the " + "matrix element, the random inputs have to be sampled internally" + ); + } + } + std::size_t channel_count = final_channel_count( + diff_xs, + first_chan_weight_remap, + first_remapped_chan_count, + second_chan_weight_remap, + second_remapped_chan_count, + prop_chan_weights, + subchan_weights + ); + if (channel_count > 1 && !prop_chan_weights && + me_output_index(me_first, MatrixElement::diagram_amp2_out) == -1) { + throw std::invalid_argument( + "diagram-based channel weights need the diagram_amp2 output of the " + "matrix element" + ); + } _channel_part_ret_types = compute_channel_part_ret_types(); } @@ -724,10 +773,20 @@ NamedVector Integrand::build_common_part( _running_coupling.value().build_function(fb, {args.at("ren_scale")}).at(0); // Evaluate differential cross section - ValueVec xs_args{ - momenta_me, - _flavor_remap.size() > 0 ? fb.gather_int(flavor_id, _flavor_remap) : flavor_id, - }; + auto& me_first = _diff_xs.at(0).matrix_element(); + int amp2_index = me_output_index(me_first, MatrixElement::diagram_amp2_out); + ValueVec xs_args; + for (auto input : me_first.external_inputs()) { + if (input == MatrixElement::momenta_in) { + xs_args.push_back(momenta_me); + } else if (input == MatrixElement::flavor_in) { + xs_args.push_back( + _flavor_remap.size() > 0 + ? fb.gather_int(flavor_id, _flavor_remap) + : flavor_id + ); + } + } xs_args.push_back(x1_acc); xs_args.push_back(x2_acc); xs_args.push_back(flavor_id); @@ -745,7 +804,7 @@ NamedVector Integrand::build_common_part( dxs_vec = _diff_xs.at(0).build_function(fb, xs_args).values(); ps_flavor_id = flavor_id; if (channel_count > 1 && !_prop_chan_weights) { - chan_weights_acc = dxs_vec.at(1); + chan_weights_acc = dxs_vec.at(amp2_index); if (_first_chan_weight_remap.size() > 0) { chan_weights_acc = fb.collect_channel_weights( chan_weights_acc, @@ -774,7 +833,7 @@ NamedVector Integrand::build_common_part( split_out.push_back(indices); } if (channel_count > 1 && !_prop_chan_weights) { - Value split_cw = outputs.at(1); + Value split_cw = outputs.at(amp2_index); if (_first_chan_weight_remap.size() > 0) { split_cw = fb.collect_channel_weights( split_cw, @@ -788,7 +847,7 @@ NamedVector Integrand::build_common_part( ++i; } for (std::size_t i = 0; auto& split_out : split_outputs) { - if (i == 1) { + if (std::cmp_equal(i, amp2_index)) { dxs_vec.push_back(Value()); } else { dxs_vec.push_back(fb.batch_merge_by_index(split_out)); @@ -964,9 +1023,20 @@ NamedVector Integrand::build_common_part( } auto zeros_int = fb.full({static_cast(0), batch_size_val}); auto zeros_float = fb.full({0., batch_size_val}); - outputs.push_back("color_index", scatter_or_drop(zeros_int, dxs_vec.at(2))); - outputs.push_back("helicity_index", scatter_or_drop(zeros_int, dxs_vec.at(3))); - outputs.push_back("diagram_index", scatter_or_drop(zeros_int, dxs_vec.at(4))); + // indices the matrix element does not provide are -1, a later + // reweighting step can fill them in + for (auto [name, me_output] : + {std::pair{"color_index", MatrixElement::color_index_out}, + std::pair{"helicity_index", MatrixElement::helicity_index_out}, + std::pair{"diagram_index", MatrixElement::diagram_index_out}}) { + int index = me_output_index(me_first, me_output); + outputs.push_back( + name, + index == -1 + ? optional_cut(fb.full({static_cast(-1), batch_size_val})) + : scatter_or_drop(zeros_int, dxs_vec.at(index)) + ); + } outputs.push_back("flavor_index", scatter_or_drop(zeros_int, ps_flavor_id)); if (subproc_id) { outputs.push_back( @@ -1005,6 +1075,12 @@ NamedVector Integrand::build_common_part( ); } } + if (_diff_xs.at(0).returns_matrix_element()) { + outputs.push_back( + "matrix_element", + scatter_or_drop(zeros_float, dxs_vec.at(me_first.outputs().size())) + ); + } if (bias_acc) { outputs.push_back( "bias", scatter_or_drop(fb.full({1., batch_size_val}), bias_acc) diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 5df1961e70..4d261fe73c 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -2202,6 +2202,7 @@ PYBIND11_MODULE(_madspace_py, m) { PdfGrid, DifferentialCrossSection::CachedPdf>&, bool, + bool, bool>(), py::arg("matrix_element"), py::arg("cm_energy"), @@ -2212,6 +2213,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("pdf2") = std::monostate{}, py::arg("input_momentum_fraction") = true, py::arg("decay") = false, + py::arg("return_matrix_element") = false, pydoc::doc("DifferentialCrossSection::DifferentialCrossSection") ) .def( diff --git a/madspace/tests/test_integrand_me_outputs.py b/madspace/tests/test_integrand_me_outputs.py new file mode 100644 index 0000000000..c146a9e5a0 --- /dev/null +++ b/madspace/tests/test_integrand_me_outputs.py @@ -0,0 +1,175 @@ +"""Integrand with matrix elements that only return |M|^2: the channel weights, +helicity, color and diagram index are optional outputs of the matrix element. + +Uses the e+ e- -> d d~ g setup of test_event_stream with PythonFunction matrix +elements, as a stand-in for surrogates that only predict |M|^2. +""" + +import os + +import numpy as np +import pytest +from test_event_stream import ( + CM_ENERGY, + PERMUTATIONS, + PIDS, + TOPOLOGY, + make_config, + make_context, +) + +import madspace as ms + + +def python_matrix_element(context, function, outputs=("matrix_element",)): + api = context.add_matrix_element( + ms.PythonFunction( + function, ["momenta"], list(outputs), particle_count=len(PIDS) + ) + ) + output_enums = { + "matrix_element": ms.MatrixElement.matrix_element_out, + "helicity_index": ms.MatrixElement.helicity_index_out, + } + return ms.MatrixElement( + api.index(), + len(PIDS), + [ms.MatrixElement.momenta_in], + [output_enums[output] for output in outputs], + ) + + +def flat_matrix_element(): + return ms.MatrixElement( + 0xBADCAFE, + len(PIDS), + ms.Integrand.matrix_element_inputs, + ms.Integrand.matrix_element_outputs, + 1, + True, + ) + + +def make_integrand(alphas_file, matrix_element, return_matrix_element=False, **kwargs): + mapping = ms.PhaseSpaceMapping( + TOPOLOGY, + CM_ENERGY, + permutations=PERMUTATIONS, + leptonic=True, + cuts=ms.Cuts( + [ + ms.CutItem(ms.Observable(PIDS, ms.Observable.obs_pt, [[i]]), min=20.0) + for i in (2, 3, 4) + ] + ), + ) + diff_xs = ms.DifferentialCrossSection( + matrix_element=matrix_element, + cm_energy=CM_ENERGY, + running_coupling=None, + energy_scale=ms.CachedScale(), + return_matrix_element=return_matrix_element, + ) + return ms.Integrand( + mapping, + [diff_xs], + running_coupling=ms.RunningCoupling(ms.AlphaSGrid(alphas_file)), + energy_scale=ms.EnergyScale(len(PIDS), 91.188), + channel_indices=[0], + **kwargs, + ) + + +def generate(run_dir, context, integrand, name): + config = make_config() + config.target_count = 2000 + channel = ms.ChannelEventGenerator( + contexts=[context], + integrand=integrand, + event_file=os.path.join(run_dir, f"events.{name}.npy"), + weight_file=os.path.join(run_dir, f"weights.{name}.npy"), + config=config, + subprocess_index=0, + name=name, + histograms=None, + ) + generator = ms.EventGenerator([context], [channel], seed=7, config=config) + generator.survey() + generator.generate() + output_file = os.path.join(run_dir, f"combined.{name}.npy") + generator.combine_to_compact_npy(output_file) + return generator.status(), np.load(output_file) + + +def test_only_matrix_element_output(tmp_path): + """Without the index outputs the events carry -1, and |M|^2 = 1 reproduces + the integral of the built-in flat matrix element.""" + context, alphas_file = make_context(str(tmp_path), 2) + matrix_element = python_matrix_element( + context, lambda momenta: np.ones(len(momenta)) + ) + status, events = generate( + str(tmp_path), context, make_integrand(alphas_file, matrix_element), "py" + ) + flat_status, flat_events = generate( + str(tmp_path), + context, + make_integrand(alphas_file, flat_matrix_element()), + "flat", + ) + assert len(events) == len(flat_events) == 2000 + for field in ("helicity_index", "color_index", "diagram_index"): + assert (events[field] == -1).all() + assert (flat_events[field] >= 0).all() + error = np.hypot(status.error, flat_status.error) + assert abs(status.mean - flat_status.mean) < 5 * error + + +def test_provided_index_output(tmp_path): + """An index the matrix element returns is written to the events.""" + context, alphas_file = make_context(str(tmp_path), 1) + matrix_element = python_matrix_element( + context, + lambda momenta: (np.ones(len(momenta)), np.full(len(momenta), 1, np.int32)), + ("matrix_element", "helicity_index"), + ) + _, events = generate( + str(tmp_path), context, make_integrand(alphas_file, matrix_element), "hel" + ) + assert (events["helicity_index"] == 1).all() + assert (events["color_index"] == -1).all() + + +def test_raw_matrix_element_output(tmp_path): + """Returning the raw |M|^2 adds an integrand output without changing the + generated events.""" + context, alphas_file = make_context(str(tmp_path), 1) + matrix_element = python_matrix_element( + context, lambda momenta: 1.0 + momenta[:, 2, 0] / CM_ENERGY + ) + results = [ + generate( + str(tmp_path), + context, + make_integrand(alphas_file, matrix_element, raw), + f"raw{raw}", + ) + for raw in (False, True) + ] + (status_a, events_a), (status_b, events_b) = results + assert status_a.mean == status_b.mean + assert np.array_equal(events_a, events_b) + + +def test_diagram_channel_weights_need_amp2(tmp_path): + context, alphas_file = make_context(str(tmp_path), 1) + matrix_element = python_matrix_element( + context, lambda momenta: np.ones(len(momenta)) + ) + with pytest.raises(ValueError, match="diagram_amp2"): + make_integrand( + alphas_file, + matrix_element, + first_chan_weight_remap=[[0, 1]], + first_remapped_chan_count=2, + ) From 06ff2dc86732c2646519470775e414736013509b Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 11:47:01 +0200 Subject: [PATCH 05/88] madspace: flavor and index of the matrix element as Integrand outputs With return_matrix_element, the Integrand also returns the flavor index passed to the matrix element and, with several differential cross sections, which one was evaluated, so that a later step can evaluate its own matrix element on the same events. --- .../include/madspace/phasespace/integrand.hpp | 5 +++- madspace/src/phasespace/integrand.cpp | 23 ++++++++++++++----- 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/madspace/include/madspace/phasespace/integrand.hpp b/madspace/include/madspace/phasespace/integrand.hpp index 7bd5dbd1ed..3da4e6e9c8 100644 --- a/madspace/include/madspace/phasespace/integrand.hpp +++ b/madspace/include/madspace/phasespace/integrand.hpp @@ -48,7 +48,10 @@ namespace madspace { * - `random` – `float`, shape `(batch, mapping.random_dim())` – the underlying * random numbers. * - `matrix_element` – `float`, shape `(batch,)` – the squared matrix element, - * present only if the differential cross sections return it. + * present only if the differential cross sections return it, followed by + * `matrix_element_flavor` – `int` – the flavor index passed to it, and with + * several differential cross sections `matrix_element_index` – `int` – the + * index of the one evaluated. * - `bias` – `float`, shape `(batch,)` – the event bias, present only with a * custom bias. `weight` then includes the bias factor. * diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index 667bb948f5..5e9893e8ef 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -163,6 +163,10 @@ Integrand::Integrand( } if (diff_xs_first.returns_matrix_element()) { ret_types.push_back("matrix_element", batch_float); + ret_types.push_back("matrix_element_flavor", batch_int); + if (diff_xs.size() > 1) { + ret_types.push_back("matrix_element_index", batch_int); + } } // before "random": the outputs after it feed the discrete histograms if (custom_bias) { @@ -775,16 +779,14 @@ NamedVector Integrand::build_common_part( // Evaluate differential cross section auto& me_first = _diff_xs.at(0).matrix_element(); int amp2_index = me_output_index(me_first, MatrixElement::diagram_amp2_out); + Value me_flavor_id = + _flavor_remap.size() > 0 ? fb.gather_int(flavor_id, _flavor_remap) : flavor_id; ValueVec xs_args; for (auto input : me_first.external_inputs()) { if (input == MatrixElement::momenta_in) { xs_args.push_back(momenta_me); } else if (input == MatrixElement::flavor_in) { - xs_args.push_back( - _flavor_remap.size() > 0 - ? fb.gather_int(flavor_id, _flavor_remap) - : flavor_id - ); + xs_args.push_back(me_flavor_id); } } xs_args.push_back(x1_acc); @@ -800,6 +802,7 @@ NamedVector Integrand::build_common_part( ValueVec dxs_vec; Value ps_flavor_id; Value subproc_id; + Value dxs_index; if (_diff_xs.size() == 1) { dxs_vec = _diff_xs.at(0).build_function(fb, xs_args).values(); ps_flavor_id = flavor_id; @@ -814,7 +817,7 @@ NamedVector Integrand::build_common_part( } } } else { - Value dxs_index = fb.gather_int(flavor_id, _flavor_diff_xs_indices); + dxs_index = fb.gather_int(flavor_id, _flavor_diff_xs_indices); ps_flavor_id = fb.gather_int(flavor_id, _flavor_per_subproc_remap); subproc_id = fb.gather_int(flavor_id, _flavor_subproc_indices); ValueVec split_indices = @@ -1080,6 +1083,14 @@ NamedVector Integrand::build_common_part( "matrix_element", scatter_or_drop(zeros_float, dxs_vec.at(me_first.outputs().size())) ); + outputs.push_back( + "matrix_element_flavor", scatter_or_drop(zeros_int, me_flavor_id) + ); + if (dxs_index) { + outputs.push_back( + "matrix_element_index", scatter_or_drop(zeros_int, dxs_index) + ); + } } if (bias_acc) { outputs.push_back( From 5dfae1e0cebbf98c3e7f0cbb51773707c6a4e753 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 11:47:17 +0200 Subject: [PATCH 06/88] madspace: re- and unweighting steps after the first unweighting Groups of channels can be given a sequence of WeightingSteps, each with its own matrix element, for example a fast surrogate in the integrand followed by the full matrix element. A step multiplies the weights by the ratio of its matrix element to the previous one. An unweighting step then unweights the ratios since the last unweighting step, times the overweight carried from it, with its own maximum weight, truncation and freeze. A reweighting step keeps the ratio in the weight. The events that pass the first unweighting stay on the device and are collected per group and context until a batch is full, or until the end of a round. The decisions are made per channel in job order, with random numbers drawn per event when the events are first unweighted, so the results do not depend on the batching or the number of threads. Integrals, event counts and the sharing of the events between channels refer to the last step, the importance sampling to the first. The color, helicity and diagram indices are taken from the first step that returns them, or replaced with override_outputs. Without groups the generation is unchanged. --- docs/source/madspace/cpp-api.rst | 6 + docs/source/madspace/python-api.rst | 2 + madspace/CMakeLists.txt | 4 + madspace/include/madspace/driver.hpp | 1 + .../madspace/driver/channel_generator.hpp | 113 +++- .../madspace/driver/event_generator.hpp | 22 +- .../madspace/driver/generator_data.hpp | 54 ++ madspace/include/madspace/driver/io.hpp | 11 + madspace/include/madspace/driver/random.hpp | 3 + .../madspace/driver/weighting_pipeline.hpp | 143 +++++ madspace/include/madspace/phasespace.hpp | 1 + .../madspace/phasespace/weighting_step.hpp | 106 ++++ madspace/src/driver/channel_generator.cpp | 462 ++++++++++++++- madspace/src/driver/event_generator.cpp | 148 ++++- madspace/src/driver/event_stream.cpp | 7 + madspace/src/driver/event_utils.cpp | 30 +- madspace/src/driver/generator_data.cpp | 24 + madspace/src/driver/weighting_pipeline.cpp | 553 ++++++++++++++++++ madspace/src/phasespace/weighting_step.cpp | 216 +++++++ madspace/src/python/madspace.cpp | 100 +++- madspace/tests/test_weighting_steps.py | 278 +++++++++ 21 files changed, 2212 insertions(+), 72 deletions(-) create mode 100644 madspace/include/madspace/driver/weighting_pipeline.hpp create mode 100644 madspace/include/madspace/phasespace/weighting_step.hpp create mode 100644 madspace/src/driver/weighting_pipeline.cpp create mode 100644 madspace/src/phasespace/weighting_step.cpp create mode 100644 madspace/tests/test_weighting_steps.py diff --git a/docs/source/madspace/cpp-api.rst b/docs/source/madspace/cpp-api.rst index a8f548477c..1a88ec1966 100644 --- a/docs/source/madspace/cpp-api.rst +++ b/docs/source/madspace/cpp-api.rst @@ -70,6 +70,7 @@ Builders that assemble compute-graph functions for channel weights, energy scale cpp/SubchannelWeights cpp/Unweighter cpp/VegasHistogram + cpp/WeightingStepFunction Phase-space building blocks --------------------------- @@ -85,6 +86,7 @@ Supporting types shared by the mappings: topologies, cuts, observables and densi cpp/PdfGrid cpp/Propagator cpp/Topology + cpp/WeightingStep cpp/initialize_vegas_grid Compute graph @@ -158,6 +160,8 @@ The high-level driver plus the tensor and device layer it runs on: integration, cpp/Runtime cpp/Sizes cpp/StatusFile + cpp/StepChainResult + cpp/StepSettings cpp/SubprocessObservables cpp/SubprocessSystArgs cpp/SystematicsCalculator @@ -172,6 +176,8 @@ The high-level driver plus the tensor and device layer it runs on: integration, cpp/UnalignedRef cpp/Variation cpp/VegasGridOptimizer + cpp/WeightingGroup + cpp/WeightingPipeline cpp/available_backends cpp/compute_generation_batch_event_count cpp/cpu_device diff --git a/docs/source/madspace/python-api.rst b/docs/source/madspace/python-api.rst index 9be3bc155a..e9f61ff485 100644 --- a/docs/source/madspace/python-api.rst +++ b/docs/source/madspace/python-api.rst @@ -86,6 +86,7 @@ Supporting types shared by the mappings: topologies, cuts, observables and densi python/PdfGrid python/Propagator python/Topology + python/WeightingStep python/initialize_vegas_grid Compute graph @@ -156,6 +157,7 @@ The high-level driver plus the tensor and device layer it runs on: integration, python/Variation python/VegasGridOptimizer python/Verbosity + python/WeightingGroup python/available_backends python/compute_generation_batch_event_count python/cpu_device diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index 646acf5685..b69730fa95 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -224,6 +224,7 @@ add_library( src/driver/event_histograms.cpp src/driver/madnis_training.cpp src/driver/random.cpp + src/driver/weighting_pipeline.cpp src/compgraphs/type.cpp src/compgraphs/function.cpp src/compgraphs/instruction.cpp @@ -260,6 +261,7 @@ add_library( src/phasespace/matrix_element.cpp src/phasespace/cross_section.cpp src/phasespace/scale.cpp + src/phasespace/weighting_step.cpp include/madspace/util.hpp include/madspace/driver.hpp include/madspace/driver/adam_optimizer.hpp @@ -279,6 +281,7 @@ add_library( include/madspace/driver/logger.hpp include/madspace/driver/madnis_training.hpp include/madspace/driver/random.hpp + include/madspace/driver/weighting_pipeline.hpp include/madspace/compgraphs.hpp include/madspace/compgraphs/type.hpp include/madspace/compgraphs/function.hpp @@ -317,6 +320,7 @@ add_library( include/madspace/phasespace/matrix_element.hpp include/madspace/phasespace/cross_section.hpp include/madspace/phasespace/scale.hpp + include/madspace/phasespace/weighting_step.hpp include/madspace/mixmax/mixmax.hpp include/madspace/mixmax/mixmax_skip_N17.c ) diff --git a/madspace/include/madspace/driver.hpp b/madspace/include/madspace/driver.hpp index 017da8d979..74142e3ff4 100644 --- a/madspace/include/madspace/driver.hpp +++ b/madspace/include/madspace/driver.hpp @@ -19,3 +19,4 @@ #include "driver/tensor.hpp" #include "driver/thread_pool.hpp" #include "driver/vegas_optimizer.hpp" +#include "driver/weighting_pipeline.hpp" diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 2a561a63dd..71f6dbf394 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -109,8 +110,9 @@ class ChannelEventGenerator { bool has_abs_integral_prior() const { return _abs_integral_prior.has_value(); } /// The integral of the absolute weights: the mean of this run, combined /// with the prior (if any) weighted by the inverse variances. Used to share - /// the events between channels. - double abs_integral_estimate() const; + /// the events between channels. With `final_step`, the integral after the + /// weighting steps, see @ref result_integral, otherwise of the first step. + double abs_integral_estimate(bool final_step = true) const; /// Current generation batch size. std::size_t batch_size() const { return _batch_size; } /// Whether a VEGAS grid or channel-weight optimization pass is still @@ -131,6 +133,87 @@ class ChannelEventGenerator { int particle_layout_extra_flags() const { return _particle_layout_extra_flags; } /// Layout of @ref event_file. const DataLayout& event_file_layout() const { return _event_file_layout; } + /// Layout of @ref weight_file. + const DataLayout& weight_layout() const { + return _step_settings.empty() ? weight_file_layout : step_weight_file_layout; + } + /// Records per event in @ref weight_file, none without steps. + std::size_t weight_record_count() const { + return _step_settings.empty() ? 0 : _step_settings.size() + 1; + } + /// Types of the unweighted events, the integrand outputs before `random`. + const NamedVector& event_types() const { return _event_types; } + /// Whether the matrix element of the integrand returns the color, helicity + /// and diagram index. + const std::array& provided_indices() const { return _provided_indices; } + + /// Integral of the final weights: after the last re- or unweighting step, + /// or the weights of the integrand without steps. + struct IntegralEstimate { + double mean; + double variance; + double mean_abs; + double variance_abs; + std::size_t count; + }; + /// See @ref IntegralEstimate. With steps, but no events through all of + /// them yet, the integral of the first step. + IntegralEstimate result_integral() const; + /// The integral of the first step, see @ref IntegralEstimate. + IntegralEstimate first_step_integral() const; + + /// Set the re- and unweighting steps after the unweighting of this channel, + /// see @ref WeightingStep. Their events are generated by the + /// @ref EventGenerator, before any events of this channel. + void set_weighting_steps(const std::vector& steps); + /// The re- and unweighting steps. + const std::vector& weighting_steps() const { return _step_settings; } + /// Maximum weight of step `step`, counted from 1. + double step_max_weight(std::size_t step) const { + return _step_states.at(step - 1).max_weight; + } + /// Fix the maximum weight of step `step`, like @ref set_fixed_max_weight. + void set_fixed_step_max_weight(std::size_t step, double max_weight); + /// Update the maximum weight of step `step` from the values it unweights, + /// with the truncation budget for `count_target` events. + void update_step_max_weight( + std::size_t step, const std::vector& values, std::size_t count_target + ); + /// Count `count` events accepted by step `step`. + void add_step_accepted(std::size_t step, double count) { + _step_states.at(step - 1).count_accepted += count; + } + /// Count `count` events accepted by the first unweighting with + /// `job_max_weight`, with steps. + void add_first_step_events(std::size_t count, double job_max_weight); + /// Incremented whenever the events are cleared, so that events still in the + /// steps can be discarded. + std::size_t event_epoch() const { return _event_epoch; } + /** + * Write events that passed all steps and add them to the integral. + * + * @param events The events, as returned by the unweighter. + * @param weights The weights of the first step. + * @param max_weight The maximum weight of the first step they were + * unweighted with. + * @param ratios Per step, the ratios of the events. + * @param max_weights Per step, the maximum weight they were unweighted + * with. + * @param biases The biases of the events, empty without. + * @param generated_count Number of points generated for these events. + */ + void write_step_events( + const TensorVec& events, + const std::vector& weights, + double max_weight, + const std::vector>& ratios, + const std::vector& max_weights, + const std::vector& biases, + std::size_t generated_count + ); + /// Signed weight of event `index` of a buffer read from @ref weight_file, + /// in units of the maximum weights, 0 for a rejected event. + double unweighted_weight(EventBuffer& weight_buffer, std::size_t index) const; /// Run a final unweighting pass over @ref weight_file into @ref /// event_file, using `rand_gen`. @@ -187,7 +270,22 @@ class ChannelEventGenerator { std::string to_json(bool include_estimates = false) const; private: + struct StepState { + double max_weight = 0.; + bool max_weight_fixed = false; + std::vector large_weights; + double count_accepted = 0.; + }; + // sums over the generated points of the final weights + struct FinalIntegral { + std::size_t count = 0; + double sum = 0., sum2 = 0., abs_sum = 0., abs_sum2 = 0.; + }; + void apply_truncation_budget(); + void update_status_integrals(); + void unweight_step_file(MixMaxRandom& rand_gen); + StepChainResult step_chain(EventBuffer& weight_buffer, std::size_t index) const; ChannelEventGenerator( const std::vector& contexts, std::size_t particle_count, @@ -215,6 +313,7 @@ class ChannelEventGenerator { RuntimePtr vegas_histogram = nullptr; RuntimePtr discrete_histogram = nullptr; RuntimePtr observable_histograms = nullptr; + RuntimePtr step_random = nullptr; }; struct FieldIndices { @@ -224,6 +323,7 @@ class ChannelEventGenerator { int x1, fact_scale1, x2, fact_scale2, partial_weight_product; int subprocess_index; int bias; + int matrix_element; int random, rest; }; @@ -236,6 +336,7 @@ class ChannelEventGenerator { DataLayout _event_file_layout; std::optional _event_file; std::optional _weight_file; + std::string _weight_file_path; std::optional _vegas_optimizer; std::optional _discrete_optimizer; std::size_t _batch_size; @@ -269,6 +370,14 @@ class ChannelEventGenerator { std::vector _histograms; std::unordered_set _used_globals; FieldIndices _field_indices; + NamedVector _event_types; + std::array _provided_indices{true, true, true}; + std::vector _step_settings; + std::vector _step_states; + std::vector _loaded_step_max_weights; + double _first_step_count = 0.; + FinalIntegral _final_integral; + std::size_t _event_epoch = 0; friend void to_json(nlohmann::json& j, const ChannelEventGenerator& channel); }; diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 20db35a3f2..70b8c6ce57 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -20,6 +20,7 @@ #include "madspace/driver/status_file.hpp" #include "madspace/driver/systematics.hpp" #include "madspace/driver/vegas_optimizer.hpp" +#include "madspace/driver/weighting_pipeline.hpp" #include "madspace/phasespace.hpp" namespace madspace { @@ -54,13 +55,18 @@ class EventGenerator { * @param status_file Optional status file progress is periodically * written to. * @param config Generator configuration. + * @param groups Groups of channels with re- and unweighting steps + * after their first unweighting, see @ref WeightingStep. + * The integrals and event counts are then those after + * the last step. */ EventGenerator( const std::vector& contexts, const std::vector>& channels, std::uint64_t seed, std::shared_ptr status_file = nullptr, - const GeneratorConfig& config = default_config + const GeneratorConfig& config = default_config, + const std::vector& groups = {} ); EventGenerator(EventGenerator&&) = default; EventGenerator& operator=(EventGenerator&&) = default; @@ -184,6 +190,7 @@ class EventGenerator { // dispatch, drained with priority by start_jobs(). std::vector> _context_unweight_queue; ResultQueue _result_queue; + std::unique_ptr _pipeline; // Base seed for reproducible event generation. std::uint64_t _seed; @@ -229,18 +236,25 @@ class EventGenerator { void reset_jobs(); void update_integral(); void update_integral_status(); - void update_integral_fractions(); + // With `final_step`, the integrals after the weighting steps are used. They + // only change in a fixed order once the steps are drained. + void update_integral_fractions(bool final_step = true); void update_counts(); void reset_start_time(); void add_timing_data(const std::string& key); void unweight_all(); - std::tuple, std::size_t, double> init_combine(); + std::tuple, std::size_t, std::vector> + init_combine(); void read_and_combine( std::vector& channel_data, EventBuffer& buffer, - double norm_factor, + const std::vector& norm_factors, MixMaxRandom& rand_gen ); + // run the re- and unweighting steps of all collected events + void drain_pipeline(); + // start the jobs of the steps whose batches are full + std::size_t start_pipeline_jobs(); // Layout of the combined events as read from the channel files DataLayout combined_layout() const; // Layout of the written events: the reweighting inputs are dropped unless diff --git a/madspace/include/madspace/driver/generator_data.hpp b/madspace/include/madspace/driver/generator_data.hpp index c089958006..3e488d4e6f 100644 --- a/madspace/include/madspace/driver/generator_data.hpp +++ b/madspace/include/madspace/driver/generator_data.hpp @@ -112,6 +112,55 @@ struct GeneratorConfig { double batch_overshoot_sigma = 1.0; }; +/// Settings of a re- or unweighting step after the first unweighting, as far +/// as the channels need them; see @ref WeightingStep. +struct StepSettings { + /// Unweighting step if true, reweighting step otherwise. + bool unweight = true; + /// See @ref GeneratorConfig::max_overweight_truncation. + double max_overweight_truncation = 0.01; + /// See @ref GeneratorConfig::freeze_max_weight_after. + std::size_t freeze_max_weight_after = 10000; +}; + +/// Result of @ref evaluate_step_chain. +struct StepChainResult { + /// Probability that the event is accepted by all steps. + double acceptance; + /// Signed weight of the event after all steps, in units of the product of + /// the maximum weights of the unweighting steps. + double weight; +}; + +/** + * Acceptance probability and weight of an event after the first unweighting and + * the re- and unweighting steps that follow. + * + * The first step accepts the event with probability + * @f$\min(1, |w_0| / m_0)@f$ and carries the overweight factor + * @f$o_0 = \max(1, |w_0| / m_0)@f$. Each reweighting step multiplies the weight + * by its ratio @f$r_k@f$. Each unweighting step unweights + * @f$u_k = o \prod r@f$, the overweight factor of the last unweighting step + * times the ratios since, with probability @f$\min(1, u_k / m_k)@f$, and carries + * @f$o_k = \max(1, u_k / m_k)@f$. Raising any maximum weight only lowers the + * acceptance, so the events can be unweighted again with larger maxima. + * + * @param weight0 The stored weight of the first step, @f$\max(|w_0|, m)@f$ + * with the sign of @f$w_0@f$, for a maximum @f$m \le m_0@f$. + * @param max_weight0 The maximum weight @f$m_0@f$ of the first step. + * @param ratios The ratios @f$r_k@f$ of the steps. + * @param max_weights The maximum weights @f$m_k@f$ of the steps, ignored for + * reweighting steps. + * @param steps The steps. + */ +StepChainResult evaluate_step_chain( + double weight0, + double max_weight0, + const double* ratios, + const double* max_weights, + const std::vector& steps +); + /** * Number of events to schedule in the next generation batch. * @@ -257,6 +306,8 @@ struct GeneratorBatchJob { std::size_t context_index; std::size_t job_id; double max_weight; + // target count of the channel when max_weight was taken + std::size_t count_target = 0; // Top-level seed plus job identity, used to derive this job's DerivedSeed(s) // independently for generate vs unweight (see start_job()/submit_unweight_job()). std::optional rng_seed; @@ -271,6 +322,9 @@ struct GeneratorBatchJob { // Nominal batch size start_job() requested; the job generates more raw points // when it repeats batches to get enough of them past the cuts. std::size_t requested_event_count = 0; + // With re- and unweighting steps: unweighted_events stay on the device and + // these are their weights, matrix elements and biases on the CPU + TensorVec step_host; }; void to_json(nlohmann::json& j, const GeneratorStatus& status); diff --git a/madspace/include/madspace/driver/io.hpp b/madspace/include/madspace/driver/io.hpp index 0c150428a9..a9b9b9592a 100644 --- a/madspace/include/madspace/driver/io.hpp +++ b/madspace/include/madspace/driver/io.hpp @@ -124,6 +124,10 @@ class ParticleRecord { // clustering data UnalignedRef cluster_scale() { return &_data[_offsets[0] + 0]; } + // weight files with re- and unweighting steps store one record per step + UnalignedRef step_ratio() { return &_data[0]; } + UnalignedRef step_max_weight() { return &_data[8]; } + void from_lhe_particle(const LHEParticle& particle) { lhe_pdg_id() = particle.pdg_id; lhe_status_code() = particle.status_code; @@ -339,6 +343,13 @@ inline const DataLayout weight_file_layout = DataLayout( ParticleRecord::layout(ParticleRecord::f_none) ); +/// Weight file of a channel with re- and unweighting steps: the weight of the +/// first step and, as "particles", the ratio and the maximum weight of each step. +inline const DataLayout step_weight_file_layout = DataLayout( + EventRecord::layout(EventRecord::f_weight), + {{"ratio", FieldLayout::f64, 0}, {"max_weight", FieldLayout::f64, 0}} +); + class EventBuffer { public: EventBuffer( diff --git a/madspace/include/madspace/driver/random.hpp b/madspace/include/madspace/driver/random.hpp index 821d054a25..e617ecd9f4 100644 --- a/madspace/include/madspace/driver/random.hpp +++ b/madspace/include/madspace/driver/random.hpp @@ -32,6 +32,9 @@ struct DerivedSeed { madnis_unweight, madnis_sample_buffer, global_init, + // per-event random numbers of the re- and unweighting steps; at most + // 16 types fit into the seed + weighting_step_random, }; std::array seed_parts; diff --git a/madspace/include/madspace/driver/weighting_pipeline.hpp b/madspace/include/madspace/driver/weighting_pipeline.hpp new file mode 100644 index 0000000000..2610ddc304 --- /dev/null +++ b/madspace/include/madspace/driver/weighting_pipeline.hpp @@ -0,0 +1,143 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "madspace/driver/channel_generator.hpp" +#include "madspace/driver/context.hpp" +#include "madspace/driver/generator_data.hpp" +#include "madspace/driver/thread_pool.hpp" +#include "madspace/phasespace/weighting_step.hpp" + +namespace madspace { + +/// A group of channels with the same re- and unweighting steps after their +/// first unweighting; see @ref WeightingStep. +struct WeightingGroup { + /// Indices of the channels in the @ref EventGenerator. + std::vector channel_indices; + /// The steps, in order. + std::vector steps; +}; + +/** + * Runs the re- and unweighting steps of the channel groups of an + * @ref EventGenerator. + * + * The events that pass the first unweighting stay on the device. They are + * collected per group and context until a device batch is full, then the + * matrix element of the next step is evaluated on them. The decisions of the + * steps are made on the scheduling thread, in the order of the jobs of each + * channel, and the random numbers they use are drawn per event when the + * events are first unweighted. The events of the channels sharing a batch + * therefore don't influence each other, and the results don't depend on how + * the events are batched. + */ +class WeightingPipeline { +public: + /// Job ids of the pipeline have this bit set. + static constexpr std::size_t job_id_flag = std::size_t(1) << 62; + + WeightingPipeline( + const std::vector& contexts, + const std::vector>& channels, + const std::vector& groups, + const GeneratorConfig& config + ); + /// Whether channel `channel_index` has steps. + bool contains(std::size_t channel_index) const { + return _channel_groups.at(channel_index) != -1; + } + /// Whether `job_id` belongs to the pipeline. + static bool owns_job(std::size_t job_id) { return job_id & job_id_flag; } + /// Take the unweighted events of `job`, in the job order of its channel. + void add_events(GeneratorBatchJob& job); + /// Start the jobs of full batches, or of all collected events with + /// `flush`. Returns the number of started jobs. + std::size_t start_jobs(ResultQueue& result_queue, bool flush); + /// Process the result of the finished job `job_id`. + void commit_job(std::size_t job_id); + /// Whether events are still collected or in jobs. + bool busy() const { return _chunk_count > 0; } + /// Drop all events, after a failed job. + void reset(); + +private: + // the unweighted events of one job of a channel + struct Chunk { + std::size_t channel_index; + std::size_t sequence; + std::size_t epoch; + std::size_t context_index; + std::size_t generated_count; + std::size_t count_target; + // on the device: rows offset to offset + size() of source, before the + // last decision, which keeps the rows in keep + std::shared_ptr source; + std::size_t offset = 0; + std::optional> keep; + TensorVec events; + // per event on the CPU + std::vector weights, matrix_elements, biases; + std::vector overweights, ratio_products; + std::vector> ratios; + double max_weight; + std::vector max_weights; + // results of the last job + std::vector step_matrix_elements, step_random; + TensorVec cpu_events; + + std::size_t size() const { return weights.size(); } + void clear(); + void compact(const std::vector& indices); + }; + using ChunkPtr = std::shared_ptr; + + struct ContextRuntimes { + RuntimePtr gather; + RuntimePtr concat; + std::vector steps; + }; + struct Group { + std::vector settings; + std::vector runtimes; + // [context][stage - 1], stage step_count + 1 copies the events to the CPU + std::vector>> buffers; + std::vector> buffer_sizes; + std::vector multiple_matrix_elements; + int momenta, alpha_s, flavor, me_index, color, helicity, diagram, random; + }; + struct Job { + std::size_t group; + std::size_t stage; + std::size_t context_index; + std::vector chunks; + }; + + void advance(ChunkPtr chunk, std::size_t stage); + void set_ready(ChunkPtr chunk, std::size_t stage); + void commit_stage(ChunkPtr chunk, std::size_t stage); + void decide(Chunk& chunk, std::size_t step); + void run_job(Job& job); + void gather_events(Chunk& chunk, ContextRuntimes& runtimes); + std::size_t batch_size(std::size_t context_index) const; + + std::vector _contexts; + std::vector> _channels; + GeneratorConfig _config; + std::vector _groups; + std::vector _channel_groups; + std::vector _channel_sequences; + // per channel and stage: the next sequence to commit, and the finished + // chunks waiting for their turn + std::vector> _next_sequences; + std::vector>> _ready; + std::unordered_map _jobs; + std::size_t _job_counter = 0; + std::size_t _chunk_count = 0; +}; + +} // namespace madspace diff --git a/madspace/include/madspace/phasespace.hpp b/madspace/include/madspace/phasespace.hpp index d4e3990962..7ffa30a15a 100644 --- a/madspace/include/madspace/phasespace.hpp +++ b/madspace/include/madspace/phasespace.hpp @@ -30,3 +30,4 @@ #include "phasespace/two_particle.hpp" #include "phasespace/unweighter.hpp" #include "phasespace/vegas.hpp" +#include "phasespace/weighting_step.hpp" diff --git a/madspace/include/madspace/phasespace/weighting_step.hpp b/madspace/include/madspace/phasespace/weighting_step.hpp new file mode 100644 index 0000000000..37be0bc14f --- /dev/null +++ b/madspace/include/madspace/phasespace/weighting_step.hpp @@ -0,0 +1,106 @@ +#pragma once + +#include +#include + +#include "madspace/phasespace/base.hpp" +#include "madspace/phasespace/matrix_element.hpp" + +namespace madspace { + +/** + * One re- or unweighting step after the first unweighting of an + * @ref EventGenerator. + * + * The step evaluates its own matrix element on the events that survived the + * previous steps and multiplies their weights by the ratio + * @f$r = |M_\text{step}|^2 / |M_\text{previous}|^2@f$. All other factors of the + * weight (PDFs, flux, Jacobian, channel weights, bias) are those of the first + * step and cancel in the ratio. An unweighting step then unweights the events + * with respect to the product of the ratios since the last unweighting step, + * times the overweight factor carried from it, with its own maximum weight. A + * reweighting step only multiplies the weights. + * + * Typical uses are a fast surrogate of the matrix element in the integrand, + * followed by an unweighting step with the full matrix element, or a + * leading-color integrand reweighted to full color. + */ +struct WeightingStep { + /// What the step does with the ratio of the matrix elements. + enum Mode { + unweight, ///< Unweight with respect to the ratio, with its own maximum. + reweight ///< Multiply the event weight by the ratio. + }; + + /// One matrix element per matrix element of the integrand, called for the + /// events with that `matrix_element_index`. The random inputs are passed + /// in, so they must not be sampled internally. + std::vector matrix_elements; + /// Unweighting or reweighting. + Mode mode = unweight; + /// Like @ref GeneratorConfig::max_overweight_truncation, for this step. + double max_overweight_truncation = 0.01; + /// Like @ref GeneratorConfig::freeze_max_weight_after, for this step. + std::size_t freeze_max_weight_after = 10000; + /// Color, helicity or diagram index outputs that replace the values of the + /// earlier steps. Without, an index is taken from the first step whose + /// matrix element returns it. + std::vector override_outputs; +}; + +/** + * Evaluates the matrix elements of a @ref WeightingStep. + * + * `batch` is the leading batch dimension. + * + * **Arguments** + * - `momenta` – `float`, shape `(batch, particle_count, 4)`. + * - `alpha_s` – `float`, shape `(batch,)`. + * - `flavor` – `int`, shape `(batch,)` – flavor index of the matrix element. + * - `matrix_element_index` – `int`, shape `(batch,)` – which matrix element to + * call, only with several matrix elements. + * - `color_index`, `helicity_index`, `diagram_index` – `int`, shape `(batch,)` + * – the indices of the earlier steps. + * - `random` – `float`, shape `(batch, 4 step_count)` – four random numbers + * per step: for the color, helicity and diagram sampling of the matrix + * element and for the unweighting. + * + * **Returns** + * - `matrix_element` – `float`, shape `(batch,)`. + * - `color_index`, `helicity_index`, `diagram_index` – `int`, shape `(batch,)` + * – the indices the step writes, otherwise the input values. + * - `unweight_random` – `float`, shape `(batch,)` – the random number of this + * step for the unweighting. + */ +class WeightingStepFunction : public FunctionGenerator { +public: + /// Random numbers per event and step: color, helicity, diagram, unweighting. + static constexpr std::size_t random_count = 4; + + /** + * @param step The step. + * @param particle_count Number of particles of the events. + * @param write_indices Whether the color, helicity and diagram indices + * returned by the matrix elements are written. + * @param step_index Position of the step, from 0. + * @param step_count Number of steps, which sets the shape of `random`. + */ + WeightingStepFunction( + const WeightingStep& step, + std::size_t particle_count, + const std::array& write_indices, + std::size_t step_index = 0, + std::size_t step_count = 1 + ); + +private: + NamedVector build_function_impl( + FunctionBuilder& fb, const NamedVector& args + ) const override; + + std::vector _matrix_elements; + std::array _write_indices; + std::size_t _step_index; +}; + +} // namespace madspace diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 57dbac687f..33117b395f 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -1,5 +1,6 @@ #include "madspace/driver/channel_generator.hpp" #include "madspace/driver/random.hpp" +#include "madspace/phasespace/weighting_step.hpp" #include "madspace/util.hpp" using namespace madspace; @@ -132,6 +133,7 @@ ChannelEventGenerator::ChannelEventGenerator( open_event_file(event_file, _event_file_layout, integrand.particle_count()) ), _weight_file(open_event_file(weight_file, weight_file_layout, 0)), + _weight_file_path(weight_file), _batch_size(config.start_batch_size), _particle_count(integrand.particle_count()), _integrand_channel_function(IntegrandChannelPart(integrand).function()), @@ -157,6 +159,17 @@ ChannelEventGenerator::ChannelEventGenerator( init_used_globals(); init_runtimes(); init_field_indices(); + auto& me_outputs = integrand.diff_xs().at(0).matrix_element().outputs(); + for (auto [provided, output] : + zip(_provided_indices, + std::array{ + MatrixElement::color_index_out, + MatrixElement::helicity_index_out, + MatrixElement::diagram_index_out + })) { + provided = + std::find(me_outputs.begin(), me_outputs.end(), output) != me_outputs.end(); + } if (const auto& grid_name = integrand.vegas_grid_name(); grid_name) { _vegas_optimizer = VegasGridOptimizer(contexts, grid_name.value(), config.vegas_damping); @@ -254,6 +267,7 @@ ChannelEventGenerator::ChannelEventGenerator( ), _event_file(open_event_file(event_file, _event_file_layout, particle_count)), _weight_file(open_event_file(weight_file, weight_file_layout, 0)), + _weight_file_path(weight_file), _batch_size(config.start_batch_size), _particle_count(particle_count), _integrand_channel_function(integrand_channel_function), @@ -299,6 +313,13 @@ void ChannelEventGenerator::init_runtimes() { } void ChannelEventGenerator::init_field_indices() { + auto& unweighter_outputs = _unweighter_function.outputs(); + for (auto [key, value] : zip(unweighter_outputs.keys(), unweighter_outputs)) { + // with the common batch size, to compare and combine channels + _event_types.push_back( + key, Type(value.type.dtype, madspace::batch_size, value.type.shape) + ); + } auto index_map = _integrand_common_function.outputs().index_map(); _field_indices.weight = index_map.at("weight"); _field_indices.momenta = index_map.at("momenta"); @@ -333,6 +354,8 @@ void ChannelEventGenerator::init_field_indices() { _field_indices.subprocess_index = -1; } _field_indices.bias = index_map.contains("bias") ? index_map.at("bias") : -1; + _field_indices.matrix_element = + index_map.contains("matrix_element") ? index_map.at("matrix_element") : -1; _field_indices.random = index_map.at("random"); _field_indices.rest = _field_indices.random + 1; } @@ -352,6 +375,10 @@ EventFile& ChannelEventGenerator::weight_file() { } void ChannelEventGenerator::unweight_file(MixMaxRandom& rand_gen) { + if (!_step_settings.empty()) { + unweight_step_file(rand_gen); + return; + } std::size_t buf_size = 1000000; EventBuffer buffer(0, 0, weight_file_layout); std::size_t accept_count = _unweighted_accept_count; @@ -395,11 +422,7 @@ void ChannelEventGenerator::integrate(const GeneratorBatchJob& job) { _cross_section.push(weight); _abs_cross_section.push(std::abs(w_view[i])); } - _status.mean = _cross_section.mean(); - _status.error = _cross_section.error(); - _status.mean_abs = _abs_cross_section.mean(); - _status.error_abs = _abs_cross_section.error(); - _status.rel_std_dev = _abs_cross_section.rel_std_dev(); + update_status_integrals(); _status.count += w_view.size(); _status.count_opt += w_view.size(); _count_requested_opt += job.requested_event_count; @@ -451,7 +474,7 @@ void ChannelEventGenerator::optimize_vegas(const GeneratorBatchJob& job) { double ChannelEventGenerator::channel_weight_sum(std::size_t event_count) { std::size_t buf_size = 1000000; - EventBuffer buffer(0, 0, weight_file_layout); + EventBuffer buffer(0, weight_record_count(), weight_layout()); double weight_sum = 0; auto& w_file = weight_file(); w_file.seek(0); @@ -468,7 +491,9 @@ double ChannelEventGenerator::channel_weight_sum(std::size_t event_count) { if (weight == 0.) { continue; } - weight_sum += std::abs(weight) / _max_weight; + weight_sum += _step_settings.empty() + ? std::abs(weight) / _max_weight + : std::abs(unweighted_weight(buffer, j)); ++unweighted_count; } if (done) { @@ -630,6 +655,7 @@ void ChannelEventGenerator::start_job( void ChannelEventGenerator::prepare_unweight_job(GeneratorBatchJob& job) const { job.max_weight = _max_weight; + job.count_target = _status.count_target; } void ChannelEventGenerator::submit_unweight_job( @@ -652,6 +678,30 @@ void ChannelEventGenerator::submit_unweight_job( ); unweighter_args.push_back(Tensor(job.max_weight, context->device())); TensorVec unw_events = runtimes.unweighter->run(unweighter_args); + if (!_step_settings.empty()) { + // the events stay on the device for the steps, with the random + // numbers of the steps drawn per event + if (job.rng_seed) { + runtimes.step_random->set_seed(DerivedSeed( + job.rng_seed, + DerivedSeed::weighting_step_random, + job.rng_job_index, + job.channel_index, + job.rng_is_survey ? 1 + job.rng_survey_pass : 0 + )); + } + auto random = runtimes.step_random->run({unw_events.at(0)}); + job.step_host.push_back(unw_events.at(_field_indices.weight).cpu()); + job.step_host.push_back( + unw_events.at(_field_indices.matrix_element).cpu() + ); + if (_field_indices.bias != -1) { + job.step_host.push_back(unw_events.at(_field_indices.bias).cpu()); + } + job.unweighted_events = std::move(unw_events); + job.unweighted_events.push_back(random.at(0)); + return; + } for (auto& item : unw_events) { job.unweighted_events.push_back(item.cpu()); } @@ -691,6 +741,16 @@ void ChannelEventGenerator::clear_events() { _cross_section.reset(); _abs_cross_section.reset(); _large_weights.clear(); + _first_step_count = 0.; + _final_integral = {}; + ++_event_epoch; + for (auto& state : _step_states) { + if (!state.max_weight_fixed) { + state.max_weight = 0.; + } + state.large_weights.clear(); + state.count_accepted = 0.; + } for (auto& hist : _histograms) { std::fill(hist.bin_values.begin(), hist.bin_values.end(), 0.); std::fill(hist.bin_errors.begin(), hist.bin_errors.end(), 0.); @@ -698,8 +758,10 @@ void ChannelEventGenerator::clear_events() { } void ChannelEventGenerator::update_max_weight(Tensor weights) { - if (_max_weight_fixed || - _status.count_unweighted > _config.freeze_max_weight_after) { + // with steps, the first step freezes on its own accepted events + double unweighted_count = + _step_settings.empty() ? _status.count_unweighted : _first_step_count; + if (_max_weight_fixed || unweighted_count > _config.freeze_max_weight_after) { return; } @@ -729,8 +791,10 @@ void ChannelEventGenerator::set_target_count(std::size_t target_count) { // (e.g. generate()'s round-one seed) would otherwise never be corrected: a // channel that has met its smaller target is not revisited. Tighten the cap // now; unweight_all() re-unweights the whole file against it at the end. + double unweighted_count = + _step_settings.empty() ? _status.count_unweighted : _first_step_count; if (target_count < old_target && _max_weight != 0 && !_max_weight_fixed && - _status.count_unweighted <= _config.freeze_max_weight_after) { + unweighted_count <= _config.freeze_max_weight_after) { apply_truncation_budget(); } } @@ -751,6 +815,7 @@ void ChannelEventGenerator::apply_truncation_budget() { if (w_sum / w > max_truncation) { if (_max_weight < w) { _status.count_unweighted *= _max_weight / w; + _first_step_count *= _max_weight / w; _max_weight = w; _unweighted_count = 0; _unweighted_accept_count = 0; @@ -874,9 +939,17 @@ std::string ChannelEventGenerator::to_json(bool include_estimates) const { j = *this; if (include_estimates) { j["max_weight"] = _max_weight; - double error = _abs_cross_section.error(); - if (_abs_cross_section.count() > 1 && std::isfinite(error) && error > 0) { - j["abs_integral"] = {{"mean", _abs_cross_section.mean()}, {"error", error}}; + auto integral = result_integral(); + double error = std::sqrt(integral.variance_abs / integral.count); + if (integral.count > 1 && std::isfinite(error) && error > 0) { + j["abs_integral"] = {{"mean", integral.mean_abs}, {"error", error}}; + } + if (!_step_states.empty()) { + json step_max_weights = json::array(); + for (auto& state : _step_states) { + step_max_weights.push_back(state.max_weight); + } + j["step_max_weights"] = step_max_weights; } } return j.dump(); @@ -940,6 +1013,14 @@ ChannelEventGenerator ChannelEventGenerator::load_json( if (channel.contains("max_weight") && channel.at("max_weight").get() > 0) { generator.set_fixed_max_weight(channel.at("max_weight").get()); } + if (channel.contains("provided_indices")) { + generator._provided_indices = + channel.at("provided_indices").get>(); + } + if (channel.contains("step_max_weights")) { + generator._loaded_step_max_weights = + channel.at("step_max_weights").get>(); + } if (channel.contains("abs_integral")) { generator.set_abs_integral_prior( channel.at("abs_integral").at("mean").get(), @@ -966,15 +1047,16 @@ void ChannelEventGenerator::set_abs_integral_prior(double mean, double error) { _abs_integral_prior = {mean, error}; } -double ChannelEventGenerator::abs_integral_estimate() const { - double mean = _abs_cross_section.mean(); +double ChannelEventGenerator::abs_integral_estimate(bool final_step) const { + auto integral = final_step ? result_integral() : first_step_integral(); + double mean = integral.mean_abs; if (!_abs_integral_prior) { return mean; } auto [prior_mean, prior_error] = *_abs_integral_prior; - double error = _abs_cross_section.error(); + double error = std::sqrt(integral.variance_abs / integral.count); // too few samples for a variance - if (_abs_cross_section.count() < 2 || !std::isfinite(error) || !(error > 0)) { + if (integral.count < 2 || !std::isfinite(error) || !(error > 0)) { return prior_mean; } double prior_weight = 1. / (prior_error * prior_error); @@ -982,6 +1064,351 @@ double ChannelEventGenerator::abs_integral_estimate() const { return (prior_mean * prior_weight + mean * weight) / (prior_weight + weight); } +ChannelEventGenerator::IntegralEstimate +ChannelEventGenerator::first_step_integral() const { + return { + .mean = _cross_section.mean(), + .variance = _cross_section.variance(), + .mean_abs = _abs_cross_section.mean(), + .variance_abs = _abs_cross_section.variance(), + .count = _cross_section.count(), + }; +} + +ChannelEventGenerator::IntegralEstimate ChannelEventGenerator::result_integral() const { + if (_step_settings.empty() || _final_integral.count == 0) { + return first_step_integral(); + } + auto& integral = _final_integral; + double count = static_cast(integral.count); + auto variance = [&](double sum2, double mean) { + return integral.count > 1 + ? std::max(0., (sum2 - count * mean * mean) / (count - 1)) + : 0.; + }; + double mean = integral.sum / count; + double mean_abs = integral.abs_sum / count; + return { + .mean = mean, + .variance = variance(integral.sum2, mean), + .mean_abs = mean_abs, + .variance_abs = variance(integral.abs_sum2, mean_abs), + .count = integral.count, + }; +} + +void ChannelEventGenerator::update_status_integrals() { + auto integral = result_integral(); + _status.mean = integral.mean; + _status.error = std::sqrt(integral.variance / integral.count); + _status.mean_abs = integral.mean_abs; + _status.error_abs = std::sqrt(integral.variance_abs / integral.count); + _status.rel_std_dev = std::sqrt(integral.variance_abs) / integral.mean_abs; +} + +void ChannelEventGenerator::set_weighting_steps( + const std::vector& steps +) { + if (steps.empty()) { + throw std::invalid_argument("at least one weighting step is needed"); + } + if (!_step_settings.empty()) { + bool same = + steps.size() == _step_settings.size() && + std::equal( + steps.begin(), + steps.end(), + _step_settings.begin(), + [](auto& a, auto& b) { + return a.unweight == b.unweight && + a.max_overweight_truncation == b.max_overweight_truncation && + a.freeze_max_weight_after == b.freeze_max_weight_after; + } + ); + if (!same) { + throw std::logic_error( + "the weighting steps of a channel cannot be changed" + ); + } + return; + } + if (_field_indices.matrix_element == -1) { + throw std::invalid_argument( + "a channel with weighting steps needs an integrand that returns the " + "matrix element" + ); + } + if (_status.count > 0) { + throw std::logic_error( + "the weighting steps must be set before events are generated" + ); + } + _step_settings = steps; + _step_states.assign(steps.size(), {}); + if (!_loaded_step_max_weights.empty()) { + if (_loaded_step_max_weights.size() != steps.size()) { + throw std::invalid_argument( + "the stored maximum weights do not match the weighting steps" + ); + } + for (std::size_t step = 1; double max_weight : _loaded_step_max_weights) { + if (max_weight > 0) { + set_fixed_step_max_weight(step, max_weight); + } + ++step; + } + } + if (_weight_file) { + _weight_file.reset(); + _weight_file = open_event_file( + _weight_file_path, step_weight_file_layout, steps.size() + 1 + ); + } + auto random_count = + static_cast(WeightingStepFunction::random_count * steps.size()); + FunctionBuilder fb( + NamedVector{{"weight", batch_float}}, + NamedVector{{"random", batch_float_array(random_count)}} + ); + fb.output(0, fb.random(fb.batch_size({fb.input(0)}), random_count)); + Function random_function = fb.function(); + for (auto [context, runtimes] : zip(_contexts, _runtimes)) { + runtimes.step_random = build_runtime(random_function, context, false); + } +} + +void ChannelEventGenerator::set_fixed_step_max_weight( + std::size_t step, double max_weight +) { + if (!(max_weight > 0)) { + throw std::invalid_argument("The fixed maximum weight must be positive"); + } + auto& state = _step_states.at(step - 1); + state.max_weight = max_weight; + state.max_weight_fixed = true; + _unweighted_count = 0; + _unweighted_accept_count = 0; +} + +void ChannelEventGenerator::update_step_max_weight( + std::size_t step, const std::vector& values, std::size_t count_target +) { + auto& state = _step_states.at(step - 1); + auto& settings = _step_settings.at(step - 1); + if (state.max_weight_fixed || + state.count_accepted > settings.freeze_max_weight_after) { + return; + } + double min_nonzero = 0.; + for (double value : values) { + double abs_value = std::abs(value); + if (abs_value != 0 && (min_nonzero == 0 || abs_value < min_nonzero)) { + min_nonzero = abs_value; + } + if (abs_value > state.max_weight) { + state.large_weights.push_back(abs_value); + } + } + if (state.max_weight == 0) { + state.max_weight = min_nonzero; + } + auto& large_weights = state.large_weights; + std::sort(large_weights.begin(), large_weights.end(), std::greater{}); + // same truncation as for the first step, see apply_truncation_budget() + double max_truncation = settings.max_overweight_truncation * + static_cast(std::min(count_target, settings.freeze_max_weight_after)); + double sum = 0; + std::size_t count = 0; + for (double value : large_weights) { + if (value < state.max_weight) { + break; + } + sum += value; + ++count; + if (sum / value > max_truncation) { + if (state.max_weight < value) { + _status.count_unweighted *= state.max_weight / value; + state.count_accepted *= state.max_weight / value; + state.max_weight = value; + _unweighted_count = 0; + _unweighted_accept_count = 0; + } + break; + } + } + large_weights.erase(large_weights.begin() + count, large_weights.end()); +} + +void ChannelEventGenerator::add_first_step_events( + std::size_t count, double job_max_weight +) { + _first_step_count += + count * (job_max_weight > 0 ? job_max_weight / _max_weight : 1); +} + +void ChannelEventGenerator::write_step_events( + const TensorVec& events, + const std::vector& weights, + double max_weight, + const std::vector>& ratios, + const std::vector& max_weights, + const std::vector& biases, + std::size_t generated_count +) { + std::size_t step_count = _step_settings.size(); + std::vector event_ratios(step_count); + // the final weights in units of the cross section, and the fraction of the + // events left by unweighting with the current maxima. Like in write_events(), + // the fraction is rescaled when a maximum rises, so it does not depend on + // the order of the updates. + double weight_unit = max_weight; + double count_factor = max_weight > 0 ? max_weight / _max_weight : 1.; + for (auto [state, settings, job_max_weight] : + zip(_step_states, _step_settings, max_weights)) { + if (settings.unweight) { + weight_unit *= job_max_weight; + count_factor *= job_max_weight > 0 ? job_max_weight / state.max_weight : 1.; + } + } + + _final_integral.count += generated_count; + for (std::size_t i = 0; i < weights.size(); ++i) { + for (std::size_t k = 0; k < step_count; ++k) { + event_ratios.at(k) = ratios.at(k).at(i); + } + auto job_chain = evaluate_step_chain( + weights.at(i), + max_weight, + event_ratios.data(), + max_weights.data(), + _step_settings + ); + double weight = job_chain.weight * weight_unit; + double xs_weight = biases.empty() ? weight : weight / biases.at(i); + _final_integral.sum += xs_weight; + _final_integral.sum2 += xs_weight * xs_weight; + _final_integral.abs_sum += std::abs(weight); + _final_integral.abs_sum2 += weight * weight; + } + + if (weights.size() > 0) { + EventBuffer event_buffer(0, _particle_count, _event_file_layout); + EventBuffer weight_buffer(0, step_count + 1, step_weight_file_layout); + fill_event_buffers(events, event_buffer, weight_buffer); + for (std::size_t i = 0; i < weights.size(); ++i) { + weight_buffer.event(i).weight() = weights.at(i); + auto first = weight_buffer.particle(i, 0); + first.step_ratio() = 1.; + first.step_max_weight() = max_weight; + for (std::size_t k = 0; k < step_count; ++k) { + auto record = weight_buffer.particle(i, k + 1); + record.step_ratio() = ratios.at(k).at(i); + record.step_max_weight() = max_weights.at(k); + } + } + event_file().write(event_buffer); + weight_file().write(weight_buffer); + } + _status.count_unweighted += weights.size() * count_factor; + _status.done = _status.count_unweighted >= _status.count_target; + update_status_integrals(); +} + +StepChainResult +ChannelEventGenerator::step_chain(EventBuffer& weight_buffer, std::size_t index) const { + std::vector ratios, max_weights; + for (std::size_t k = 0; auto& state : _step_states) { + ratios.push_back(weight_buffer.particle(index, k + 1).step_ratio()); + max_weights.push_back(state.max_weight); + ++k; + } + return evaluate_step_chain( + weight_buffer.event(index).weight(), + _max_weight, + ratios.data(), + max_weights.data(), + _step_settings + ); +} + +double ChannelEventGenerator::unweighted_weight( + EventBuffer& weight_buffer, std::size_t index +) const { + double weight = weight_buffer.event(index).weight(); + if (_step_settings.empty()) { + return std::copysign(std::max(1., std::abs(weight) / _max_weight), weight); + } + return weight == 0. ? 0. : step_chain(weight_buffer, index).weight; +} + +void ChannelEventGenerator::unweight_step_file(MixMaxRandom& rand_gen) { + std::size_t buf_size = 100000; + std::size_t step_count = _step_settings.size(); + EventBuffer buffer(0, step_count + 1, step_weight_file_layout); + std::vector ratios(step_count), old_max_weights(step_count), + max_weights(step_count); + for (auto [max_weight, state] : zip(max_weights, _step_states)) { + max_weight = state.max_weight; + } + std::size_t accept_count = _unweighted_accept_count; + auto& w_file = weight_file(); + for (std::size_t i = _unweighted_count; i < w_file.event_count(); i += buf_size) { + w_file.seek(i); + w_file.read(buffer, buf_size); + for (std::size_t j = 0; j < buffer.event_count(); ++j) { + auto weight = buffer.event(j).weight(); + double random = rand_gen.generate_double(); + if (weight.value() == 0.) { + continue; + } + for (std::size_t k = 0; k < step_count; ++k) { + auto record = buffer.particle(j, k + 1); + ratios.at(k) = record.step_ratio(); + old_max_weights.at(k) = record.step_max_weight(); + } + // accepted with the stored maxima: keep with the ratio of the + // acceptances, which only decrease with larger maxima + double old_acceptance = + evaluate_step_chain( + weight.value(), + buffer.particle(j, 0).step_max_weight(), + ratios.data(), + old_max_weights.data(), + _step_settings + ) + .acceptance; + double acceptance = + evaluate_step_chain( + weight.value(), + _max_weight, + ratios.data(), + max_weights.data(), + _step_settings + ) + .acceptance; + if (!(acceptance > random * old_acceptance)) { + weight = 0; + continue; + } + weight = std::copysign( + std::max(std::abs(weight.value()), _max_weight), weight.value() + ); + buffer.particle(j, 0).step_max_weight() = _max_weight; + for (std::size_t k = 0; k < step_count; ++k) { + if (_step_settings.at(k).unweight) { + buffer.particle(j, k + 1).step_max_weight() = max_weights.at(k); + } + } + ++accept_count; + } + w_file.seek(i); + w_file.write(buffer); + } + _unweighted_count = w_file.event_count(); + _unweighted_accept_count = accept_count; + _status.count_unweighted = accept_count; +} + void madspace::to_json(nlohmann::json& j, const ChannelEventGenerator& channel) { json histograms = json::array(); json histogram_function; @@ -1008,5 +1435,6 @@ void madspace::to_json(nlohmann::json& j, const ChannelEventGenerator& channel) {"subprocess_index", channel._status.subprocess}, {"name", channel._status.name}, {"histograms", histograms}, + {"provided_indices", channel._provided_indices}, }; } diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index d78d9f765d..4da3157c8b 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -22,7 +22,8 @@ EventGenerator::EventGenerator( const std::vector>& channels, std::uint64_t seed, std::shared_ptr status_file, - const GeneratorConfig& config + const GeneratorConfig& config, + const std::vector& groups ) : _config(config), _status{ @@ -56,6 +57,11 @@ EventGenerator::EventGenerator( _channel_unweight_order(channels.size()), _channel_unweight_ready(channels.size()), _context_unweight_queue(contexts.size()), + _pipeline( + groups.empty() + ? nullptr + : std::make_unique(contexts, channels, groups, config) + ), _seed(seed), _status_file(status_file) {} @@ -84,7 +90,11 @@ void EventGenerator::commit_generate_job(GeneratorBatchJob& job) { void EventGenerator::commit_unweight_job(GeneratorBatchJob& job) { auto& channel = _channels.at(job.channel_index); - channel->write_events(job.unweighted_events, job.max_weight); + if (_pipeline && _pipeline->contains(job.channel_index)) { + _pipeline->add_events(job); + } else { + channel->write_events(job.unweighted_events, job.max_weight); + } update_counts(); --_channel_job_counts.at(job.channel_index); --_context_job_counts.at(job.context_index); @@ -185,6 +195,16 @@ void EventGenerator::run_generate(GenerateMode mode) { // Wait for the round to fully commit before resyncing cross-channel fractions. while (round_in_flight > 0) { std::size_t job_id = wait_for_result(round_in_flight); + if (WeightingPipeline::owns_job(job_id)) { + _pipeline->commit_job(job_id); + update_counts(); + print_gen_update(false); + std::size_t job_id_refill = _job_id; + std::size_t refill_unweight_dispatched = start_jobs(); + round_in_flight += (_job_id - job_id_refill) + + refill_unweight_dispatched + start_pipeline_jobs(); + continue; + } auto& job = _running_jobs.at(job_id); if (job.unweighted_events.size() == 0) { auto& ready = _channel_ready_gen.at(job.channel_index); @@ -212,8 +232,10 @@ void EventGenerator::run_generate(GenerateMode mode) { } std::size_t job_id_refill = _job_id; std::size_t refill_unweight_dispatched = start_jobs(); - round_in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched; + round_in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched + + start_pipeline_jobs(); } + drain_pipeline(); if (mode == GenerateMode::optimize_only) { if (std::none_of(_channels.begin(), _channels.end(), [](auto& channel) { @@ -291,6 +313,15 @@ void EventGenerator::survey(std::size_t survey_pass) { done = true; while (in_flight > 0) { std::size_t job_id = wait_for_result(in_flight); + if (WeightingPipeline::owns_job(job_id)) { + _pipeline->commit_job(job_id); + update_counts(); + std::size_t job_id_refill = _job_id; + std::size_t refill_unweight_dispatched = start_jobs(); + in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched + + start_pipeline_jobs(); + continue; + } auto& job = _running_jobs.at(job_id); if (job.unweighted_events.size() == 0) { // Generate-stage completion, globally ordered via @@ -339,7 +370,13 @@ void EventGenerator::survey(std::size_t survey_pass) { order.pop_front(); auto& uw_job = _running_jobs.at(commit_id); auto& channel = _channels.at(uw_job.channel_index); - channel->write_events(uw_job.unweighted_events, uw_job.max_weight); + if (_pipeline && _pipeline->contains(uw_job.channel_index)) { + _pipeline->add_events(uw_job); + } else { + channel->write_events( + uw_job.unweighted_events, uw_job.max_weight + ); + } update_counts(); auto& channel_job_count = _channel_job_counts.at(uw_job.channel_index); @@ -358,9 +395,11 @@ void EventGenerator::survey(std::size_t survey_pass) { } std::size_t job_id_refill = _job_id; std::size_t refill_unweight_dispatched = start_jobs(); - in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched; + in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched + + start_pipeline_jobs(); print_survey_update(false, done_event_count, total_event_count, iter); } + drain_pipeline(); } print_survey_update(true, done_event_count, total_event_count, iter - 1); } @@ -392,6 +431,9 @@ std::size_t EventGenerator::wait_for_result(std::size_t& in_flight) { void EventGenerator::reset_jobs() { _running_jobs.clear(); + if (_pipeline) { + _pipeline->reset(); + } _ready_jobs.clear(); _ready_job_rr_cursor = 0; _ready_gen.clear(); @@ -525,30 +567,45 @@ std::size_t EventGenerator::start_jobs() { return unweight_dispatched; } +std::size_t EventGenerator::start_pipeline_jobs() { + return _pipeline ? _pipeline->start_jobs(_result_queue, false) : 0; +} + +// Batches no longer have to be full: the steps don't depend on the batching. +void EventGenerator::drain_pipeline() { + if (!_pipeline) { + return; + } + std::size_t in_flight = 0; + while (_pipeline->busy()) { + in_flight += _pipeline->start_jobs(_result_queue, true); + if (in_flight == 0) { + throw std::logic_error("weighting steps stuck without running jobs"); + } + _pipeline->commit_job(wait_for_result(in_flight)); + update_counts(); + } + update_integral_status(); +} + void EventGenerator::update_integral() { update_integral_status(); - update_integral_fractions(); + update_integral_fractions(false); } void EventGenerator::update_integral_status() { sum_channel_status(_channels, _status); } -void EventGenerator::update_integral_fractions() { - // With an integral known from an earlier run, the sum of the channel estimates - // normalizes the fractions; without, this is _status.mean_abs. - double total_estimate = _status.mean_abs; - if (std::any_of(_channels.begin(), _channels.end(), [](auto& channel) { - return channel->has_abs_integral_prior(); - })) { - total_estimate = 0.; - for (auto& channel : _channels) { - total_estimate += channel->abs_integral_estimate(); - } +void EventGenerator::update_integral_fractions(bool final_step) { + // the sum of the channel estimates normalizes the fractions + double total_estimate = 0.; + for (auto& channel : _channels) { + total_estimate += channel->abs_integral_estimate(final_step); } for (auto [channel, integral_fraction] : zip(_channels, _channel_integral_fractions)) { - integral_fraction = channel->abs_integral_estimate() / total_estimate; + integral_fraction = channel->abs_integral_estimate(final_step) / total_estimate; } // Distribute events between channels, ensure sum is exactly target_count @@ -578,7 +635,12 @@ void EventGenerator::update_integral_fractions() { for (auto [channel, count, fraction] : zip(_channels, counts, _channel_integral_fractions)) { if (_generate_mode == GenerateMode::fix_max_weights) { - count = fraction > 0 ? _config.freeze_max_weight_after : 0; + // the maxima of all steps have to freeze + std::size_t freeze_count = _config.freeze_max_weight_after; + for (auto& step : channel->weighting_steps()) { + freeze_count = std::max(freeze_count, step.freeze_max_weight_after); + } + count = fraction > 0 ? freeze_count : 0; } channel->set_target_count(count); } @@ -616,7 +678,7 @@ void EventGenerator::combine_to_compact_npy( _systematics = systematics; _event_histograms = histograms; MixMaxRandom select_rng(DerivedSeed(_seed, DerivedSeed::combine_select)); - auto [channel_data, particle_count, norm_factor] = init_combine(); + auto [channel_data, particle_count, norm_factors] = init_combine(); DataLayout layout = combined_layout(); DataLayout out_layout = output_layout( EventRecord::f_weight | EventRecord::f_subproc_index | @@ -632,7 +694,7 @@ void EventGenerator::combine_to_compact_npy( print_combine_init(); while (true) { _abort_check_function(); - read_and_combine(channel_data, buffer, norm_factor, select_rng); + read_and_combine(channel_data, buffer, norm_factors, select_rng); if (buffer.event_count() == 0) { break; } @@ -665,7 +727,7 @@ void EventGenerator::combine_to_lhe_npy( std::uint64_t seed = _seed; MixMaxRandom select_rng(DerivedSeed(seed, DerivedSeed::combine_select)); MixMaxRandom rand_gen(DerivedSeed(seed, DerivedSeed::lhe_complete)); - auto [channel_data, particle_count, norm_factor] = init_combine(); + auto [channel_data, particle_count, norm_factors] = init_combine(); DataLayout in_layout = combined_layout(); DataLayout out_layout = output_layout(EventRecord::f_lhe_event, ParticleRecord::f_lhe_particle); @@ -681,7 +743,7 @@ void EventGenerator::combine_to_lhe_npy( print_combine_init(); while (true) { _abort_check_function(); - read_and_combine(channel_data, buffer, norm_factor, select_rng); + read_and_combine(channel_data, buffer, norm_factors, select_rng); if (buffer.event_count() == 0) { break; } @@ -718,7 +780,7 @@ void EventGenerator::combine_to_lhe( std::uint64_t seed = _seed; MixMaxRandom select_rng(DerivedSeed(seed, DerivedSeed::combine_select)); ThreadPool pool(_config.combine_thread_count); - auto [channel_data, particle_count, norm_factor] = init_combine(); + auto [channel_data, particle_count, norm_factors] = init_combine(); struct CombineBuffers { EventBuffer in_buffer; std::string out_buffer; @@ -761,7 +823,7 @@ void EventGenerator::combine_to_lhe( while (idle_buffers.size() > 0 && !done) { std::size_t slot = idle_buffers.back(); auto& [in_buffer, out_buffer, syst_weights] = buffers.at(slot); - read_and_combine(channel_data, in_buffer, norm_factor, select_rng); + read_and_combine(channel_data, in_buffer, norm_factors, select_rng); if (in_buffer.event_count() == 0) { done = true; break; @@ -1045,12 +1107,16 @@ std::vector EventGenerator::histograms() const { return hists; } -std::tuple, std::size_t, double> +std::tuple< + std::vector, + std::size_t, + std::vector> EventGenerator::init_combine() { std::vector channel_data; std::size_t count_sum = 0; std::size_t particle_count = 0; double weight_sum = 0.; + std::vector channel_weight_sums; for (auto& channel : _channels) { particle_count = std::max(particle_count, channel->event_file().particle_count()); @@ -1059,24 +1125,39 @@ EventGenerator::init_combine() { std::size_t count = channel->status().count_target; count_sum += count; channel->event_file().seek(0); - weight_sum += channel->channel_weight_sum(count); + channel_weight_sums.push_back(channel->channel_weight_sum(count)); + weight_sum += channel_weight_sums.back(); channel->weight_file().seek(0); channel_data.push_back({ .cum_count = count_sum, .event_buffer = EventBuffer( 0, channel->event_file().particle_count(), channel->event_file_layout() ), - .weight_buffer = EventBuffer(0, 0, weight_file_layout), + .weight_buffer = EventBuffer( + 0, channel->weight_record_count(), channel->weight_layout() + ), .buffer_index = 0, }); } - return {channel_data, particle_count, _status.mean_abs * count_sum / weight_sum}; + std::vector norm_factors; + for (auto [channel, channel_weight_sum] : zip(_channels, channel_weight_sums)) { + // After reweighting steps the mean event weights differ between + // channels, so the events of each channel are normalized separately + norm_factors.push_back( + _pipeline + ? (channel_weight_sum > 0 + ? channel->status().mean_abs * count_sum / channel_weight_sum + : 0.) + : _status.mean_abs * count_sum / weight_sum + ); + } + return {channel_data, particle_count, norm_factors}; } void EventGenerator::read_and_combine( std::vector& channel_data, EventBuffer& buffer, - double norm_factor, + const std::vector& norm_factors, MixMaxRandom& rand_gen ) { std::size_t batch_size = 1000; @@ -1096,7 +1177,8 @@ void EventGenerator::read_and_combine( std::for_each(sampled_chan, channel_data.end(), [](auto& chan) { --chan.cum_count; }); - auto& channel = _channels.at(sampled_chan - channel_data.begin()); + std::size_t channel_index = sampled_chan - channel_data.begin(); + auto& channel = _channels.at(channel_index); double weight = 0.; while (true) { @@ -1119,7 +1201,9 @@ void EventGenerator::read_and_combine( sampled_chan->buffer_index, buffer, event_index, - unweighted_event_weight(weight, channel->max_weight()) * norm_factor, + channel->unweighted_weight( + sampled_chan->weight_buffer, sampled_chan->buffer_index + ) * norm_factors.at(channel_index), static_cast(channel->status().subprocess), extra_flags ); diff --git a/madspace/src/driver/event_stream.cpp b/madspace/src/driver/event_stream.cpp index 4798c70421..c9974c2e07 100644 --- a/madspace/src/driver/event_stream.cpp +++ b/madspace/src/driver/event_stream.cpp @@ -34,6 +34,13 @@ EventStream::EventStream( if (channels.empty()) { throw std::invalid_argument("EventStream needs at least one channel"); } + for (auto& channel : channels) { + if (!channel->weighting_steps().empty()) { + throw std::invalid_argument( + "EventStream does not support weighting steps yet" + ); + } + } return combined_event_layout( channels.at(0)->event_layout_extra_flags(), channels.at(0)->particle_layout_extra_flags() diff --git a/madspace/src/driver/event_utils.cpp b/madspace/src/driver/event_utils.cpp index 37cc063a82..8817744039 100644 --- a/madspace/src/driver/event_utils.cpp +++ b/madspace/src/driver/event_utils.cpp @@ -136,27 +136,25 @@ void madspace::sum_channel_status( const ChannelEventGenerator* bad_channel = nullptr; for (auto& channel : channels) { auto& chan_status = channel->status(); - auto& cross_section = channel->cross_section(); - auto& abs_cross_section = channel->abs_cross_section(); + auto integral = channel->result_integral(); // special case for channels with 0 samples, as they have nan variance - if (bad_channel == nullptr && cross_section.count() > 0 && - (!std::isfinite(cross_section.mean()) || - !std::isfinite(cross_section.variance()) || - !std::isfinite(abs_cross_section.mean()) || - !std::isfinite(abs_cross_section.variance()))) { + if (bad_channel == nullptr && integral.count > 0 && + (!std::isfinite(integral.mean) || !std::isfinite(integral.variance) || + !std::isfinite(integral.mean_abs) || + !std::isfinite(integral.variance_abs))) { bad_channel = channel.get(); } - total_mean += cross_section.mean(); - total_mean_abs += abs_cross_section.mean(); - if (cross_section.count() > 0) { - total_var += cross_section.variance() / cross_section.count(); - total_var_abs += abs_cross_section.variance() / abs_cross_section.count(); + total_mean += integral.mean; + total_mean_abs += integral.mean_abs; + if (integral.count > 0) { + total_var += integral.variance / integral.count; + total_var_abs += integral.variance_abs / integral.count; } total_count += chan_status.count; total_count_opt += chan_status.count_opt; total_count_after_cuts += chan_status.count_after_cuts; total_count_after_cuts_opt += chan_status.count_after_cuts_opt; - total_integ_count += cross_section.count(); + total_integ_count += integral.count; iterations = std::max(chan_status.iterations, iterations); if (!chan_status.optimized) { optimized = false; @@ -168,9 +166,9 @@ void madspace::sum_channel_status( ? std::format( "channel '{}' after {} samples (mean={}, variance={})", bad_channel->status().name, - bad_channel->cross_section().count(), - bad_channel->cross_section().mean(), - bad_channel->cross_section().variance() + bad_channel->result_integral().count, + bad_channel->result_integral().mean, + bad_channel->result_integral().variance ) : std::format("the sum over channels (mean={})", total_mean); throw std::runtime_error( diff --git a/madspace/src/driver/generator_data.cpp b/madspace/src/driver/generator_data.cpp index 3c1a0da788..1b3679ed61 100644 --- a/madspace/src/driver/generator_data.cpp +++ b/madspace/src/driver/generator_data.cpp @@ -88,3 +88,27 @@ void madspace::to_json(nlohmann::json& j, const Histogram& hist) { {"bin_errors", hist.bin_errors}, }; } + +StepChainResult madspace::evaluate_step_chain( + double weight0, + double max_weight0, + const double* ratios, + const double* max_weights, + const std::vector& steps +) { + double abs_weight0 = std::abs(weight0); + double acceptance = std::min(1., abs_weight0 / max_weight0); + double overweight = std::max(1., abs_weight0 / max_weight0); + double pending = std::copysign(1., weight0); + for (std::size_t i = 0; auto& step : steps) { + pending *= ratios[i]; + if (step.unweight) { + double unweight_value = overweight * std::abs(pending); + acceptance *= std::min(1., unweight_value / max_weights[i]); + overweight = std::max(1., unweight_value / max_weights[i]); + pending = std::copysign(1., pending); + } + ++i; + } + return {acceptance, overweight * pending}; +} diff --git a/madspace/src/driver/weighting_pipeline.cpp b/madspace/src/driver/weighting_pipeline.cpp new file mode 100644 index 0000000000..f1c68c1dcc --- /dev/null +++ b/madspace/src/driver/weighting_pipeline.cpp @@ -0,0 +1,553 @@ +#include "madspace/driver/weighting_pipeline.hpp" + +#include +#include +#include + +#include "madspace/util.hpp" + +using namespace madspace; + +namespace { + +std::vector to_vector(const Tensor& tensor) { + auto view = tensor.view(); + std::vector ret(view.size()); + for (std::size_t i = 0; i < ret.size(); ++i) { + ret[i] = view[i]; + } + return ret; +} + +template +void compact_vector(std::vector& vec, const std::vector& indices) { + if (vec.empty()) { + return; + } + std::vector ret; + ret.reserve(indices.size()); + for (auto index : indices) { + ret.push_back(vec.at(index)); + } + vec = std::move(ret); +} + +Function gather_function(const NamedVector& types) { + NamedVector arg_types{{"indices", batch_int}}; + arg_types.insert_back(types); + FunctionBuilder fb(arg_types, types); + for (std::size_t i = 0; i < types.size(); ++i) { + fb.output(i, fb.batch_gather(fb.input(0), fb.input(i + 1))); + } + return fb.function(); +} + +Function concat_function(const NamedVector& types) { + NamedVector arg_types; + for (std::size_t part = 1; part <= 2; ++part) { + for (auto [key, type] : zip(types.keys(), types)) { + arg_types.push_back(std::format("arg{}_{}", part, key), type); + } + } + FunctionBuilder fb(arg_types, types); + int count = static_cast(types.size()); + for (int i = 0; i < count; ++i) { + auto [cat, sizes] = fb.batch_cat({fb.input(i), fb.input(i + count)}); + fb.output(i, cat); + } + return fb.function(); +} + +} // namespace + +void WeightingPipeline::Chunk::clear() { + source.reset(); + keep.reset(); + events.clear(); + cpu_events.clear(); + weights.clear(); + matrix_elements.clear(); + biases.clear(); + overweights.clear(); + ratio_products.clear(); + for (auto& step_ratios : ratios) { + step_ratios.clear(); + } +} + +void WeightingPipeline::Chunk::compact(const std::vector& indices) { + compact_vector(weights, indices); + compact_vector(matrix_elements, indices); + compact_vector(biases, indices); + compact_vector(overweights, indices); + compact_vector(ratio_products, indices); + for (auto& step_ratios : ratios) { + compact_vector(step_ratios, indices); + } +} + +WeightingPipeline::WeightingPipeline( + const std::vector& contexts, + const std::vector>& channels, + const std::vector& groups, + const GeneratorConfig& config +) : + _contexts(contexts), + _channels(channels), + _config(config), + _channel_groups(channels.size(), -1), + _channel_sequences(channels.size()), + _next_sequences(channels.size()), + _ready(channels.size()) { + const std::array index_outputs{ + MatrixElement::color_index_out, + MatrixElement::helicity_index_out, + MatrixElement::diagram_index_out, + }; + for (std::size_t group_index = 0; auto& group_config : groups) { + if (group_config.steps.empty() || group_config.channel_indices.empty()) { + throw std::invalid_argument( + "a weighting group needs at least one channel and one step" + ); + } + for (auto channel_index : group_config.channel_indices) { + if (channel_index >= channels.size()) { + throw std::invalid_argument( + "channel index of weighting group too large" + ); + } + if (_channel_groups.at(channel_index) != -1) { + throw std::invalid_argument( + "a channel can only belong to one weighting group" + ); + } + _channel_groups.at(channel_index) = static_cast(group_index); + } + + auto& first_channel = channels.at(group_config.channel_indices.at(0)); + NamedVector types = first_channel->event_types(); + for (auto channel_index : group_config.channel_indices) { + auto& channel = channels.at(channel_index); + auto& channel_types = channel->event_types(); + if (channel_types.keys() != types.keys() || + !std::equal( + channel_types.begin(), channel_types.end(), types.begin() + ) || + channel->provided_indices() != first_channel->provided_indices()) { + throw std::invalid_argument( + "the channels of a weighting group need the same events" + ); + } + } + auto& index_map = types.index_map(); + if (!index_map.contains("matrix_element")) { + throw std::invalid_argument( + "a weighting group needs integrands that return the matrix element" + ); + } + std::size_t step_count = group_config.steps.size(); + types.push_back( + "step_random", + batch_float_array(WeightingStepFunction::random_count * step_count) + ); + + Group group{ + .momenta = static_cast(index_map.at("momenta")), + .alpha_s = static_cast(index_map.at("alpha_qcd")), + .flavor = static_cast(index_map.at("matrix_element_flavor")), + .me_index = index_map.contains("matrix_element_index") + ? static_cast(index_map.at("matrix_element_index")) + : -1, + .color = static_cast(index_map.at("color_index")), + .helicity = static_cast(index_map.at("helicity_index")), + .diagram = static_cast(index_map.at("diagram_index")), + .random = static_cast(types.size() - 1), + }; + + std::vector step_functions; + auto provided = first_channel->provided_indices(); + for (std::size_t step_index = 0; auto& step : group_config.steps) { + if (step.matrix_elements.size() > 1 && group.me_index == -1) { + throw std::invalid_argument( + "several matrix elements in a weighting step, but the integrand " + "has only one" + ); + } + group.multiple_matrix_elements.push_back(step.matrix_elements.size() > 1); + group.settings.push_back({ + .unweight = step.mode == WeightingStep::unweight, + .max_overweight_truncation = step.max_overweight_truncation, + .freeze_max_weight_after = step.freeze_max_weight_after, + }); + std::array write_indices{}; + for (std::size_t i = 0; i < index_outputs.size(); ++i) { + auto output = index_outputs.at(i); + auto& outputs = step.matrix_elements.at(0).outputs(); + bool returns = + std::find(outputs.begin(), outputs.end(), output) != outputs.end(); + bool overrides = + std::find( + step.override_outputs.begin(), + step.override_outputs.end(), + output + ) != step.override_outputs.end(); + write_indices.at(i) = returns && (!provided.at(i) || overrides); + provided.at(i) = provided.at(i) || returns; + } + step_functions.push_back( + WeightingStepFunction( + step, + first_channel->particle_count(), + write_indices, + step_index, + step_count + ) + .function() + ); + ++step_index; + } + + Function gather = gather_function(types); + Function concat = concat_function(types); + for (auto& context : contexts) { + ContextRuntimes runtimes{ + .gather = build_runtime(gather, context, false), + .concat = build_runtime(concat, context, false), + }; + for (auto& function : step_functions) { + runtimes.steps.push_back(build_runtime(function, context, false)); + } + group.runtimes.push_back(std::move(runtimes)); + } + group.buffers.assign( + contexts.size(), std::vector>(step_count + 1) + ); + group.buffer_sizes.assign( + contexts.size(), std::vector(step_count + 1) + ); + for (auto channel_index : group_config.channel_indices) { + channels.at(channel_index)->set_weighting_steps(group.settings); + _next_sequences.at(channel_index).assign(step_count + 1, 0); + _ready.at(channel_index).resize(step_count + 1); + } + _groups.push_back(std::move(group)); + ++group_index; + } +} + +std::size_t WeightingPipeline::batch_size(std::size_t context_index) const { + return _contexts.at(context_index)->device()->device_type() == DeviceType::cpu + ? _config.cpu_batch_size + : _config.gpu_batch_size; +} + +void WeightingPipeline::add_events(GeneratorBatchJob& job) { + auto& channel = _channels.at(job.channel_index); + auto chunk = std::make_shared(); + chunk->channel_index = job.channel_index; + chunk->sequence = _channel_sequences.at(job.channel_index)++; + chunk->epoch = channel->event_epoch(); + chunk->context_index = job.context_index; + chunk->generated_count = job.weights.size(0); + chunk->count_target = job.count_target; + chunk->source = std::make_shared(std::move(job.unweighted_events)); + chunk->weights = to_vector(job.step_host.at(0)); + chunk->matrix_elements = to_vector(job.step_host.at(1)); + if (job.step_host.size() > 2) { + chunk->biases = to_vector(job.step_host.at(2)); + } + chunk->max_weight = job.max_weight; + for (double weight : chunk->weights) { + chunk->overweights.push_back(std::max(1., std::abs(weight) / job.max_weight)); + chunk->ratio_products.push_back(std::copysign(1., weight)); + } + channel->add_first_step_events(chunk->size(), job.max_weight); + job.unweighted_events.clear(); + job.step_host.clear(); + ++_chunk_count; + advance(chunk, 1); +} + +void WeightingPipeline::advance(ChunkPtr chunk, std::size_t stage) { + if (chunk->size() == 0) { + set_ready(chunk, stage); + return; + } + auto& group = _groups.at(_channel_groups.at(chunk->channel_index)); + group.buffers.at(chunk->context_index).at(stage - 1).push_back(chunk); + group.buffer_sizes.at(chunk->context_index).at(stage - 1) += chunk->size(); +} + +void WeightingPipeline::set_ready(ChunkPtr chunk, std::size_t stage) { + std::size_t channel_index = chunk->channel_index; + auto& ready = _ready.at(channel_index).at(stage - 1); + auto& next_sequence = _next_sequences.at(channel_index).at(stage - 1); + ready.emplace(chunk->sequence, chunk); + while (!ready.empty() && ready.begin()->first == next_sequence) { + ChunkPtr next_chunk = ready.begin()->second; + ready.erase(ready.begin()); + ++next_sequence; + commit_stage(next_chunk, stage); + } +} + +void WeightingPipeline::commit_stage(ChunkPtr chunk, std::size_t stage) { + auto& channel = _channels.at(chunk->channel_index); + std::size_t step_count = channel->weighting_steps().size(); + // the events of the channel were cleared since + bool stale = chunk->epoch != channel->event_epoch(); + if (stale) { + chunk->clear(); + } + if (stage <= step_count) { + if (!stale && chunk->size() > 0) { + decide(*chunk, stage); + } else { + chunk->ratios.emplace_back(); + chunk->max_weights.push_back(0.); + } + advance(chunk, stage + 1); + return; + } + if (!stale) { + channel->write_step_events( + chunk->cpu_events, + chunk->weights, + chunk->max_weight, + chunk->ratios, + chunk->max_weights, + chunk->biases, + chunk->generated_count + ); + } + --_chunk_count; +} + +void WeightingPipeline::decide(Chunk& chunk, std::size_t step) { + auto& channel = _channels.at(chunk.channel_index); + auto& settings = channel->weighting_steps().at(step - 1); + std::size_t count = chunk.size(); + std::vector ratios(count); + for (std::size_t i = 0; i < count; ++i) { + double previous = chunk.matrix_elements.at(i); + ratios.at(i) = + previous != 0. ? chunk.step_matrix_elements.at(i) / previous : 0.; + chunk.ratio_products.at(i) *= ratios.at(i); + } + std::vector keep; + if (settings.unweight) { + // the overweight of the last unweighting times the ratios since, see + // evaluate_step_chain() + std::vector values(count); + for (std::size_t i = 0; i < count; ++i) { + values.at(i) = + chunk.overweights.at(i) * std::abs(chunk.ratio_products.at(i)); + } + channel->update_step_max_weight(step, values, chunk.count_target); + double max_weight = channel->step_max_weight(step); + for (std::size_t i = 0; i < count; ++i) { + double value = values.at(i); + if (value > 0. && value > chunk.step_random.at(i) * max_weight) { + keep.push_back(static_cast(i)); + } + chunk.overweights.at(i) = std::max(1., value / max_weight); + chunk.ratio_products.at(i) = std::copysign(1., chunk.ratio_products.at(i)); + } + channel->add_step_accepted(step, static_cast(keep.size())); + chunk.max_weights.push_back(max_weight); + } else { + for (std::size_t i = 0; i < count; ++i) { + if (ratios.at(i) != 0.) { + keep.push_back(static_cast(i)); + } + } + chunk.max_weights.push_back(0.); + } + chunk.ratios.push_back(std::move(ratios)); + chunk.matrix_elements = std::move(chunk.step_matrix_elements); + chunk.step_matrix_elements.clear(); + chunk.step_random.clear(); + if (keep.size() < count) { + chunk.compact(keep); + chunk.keep = std::move(keep); + } + if (chunk.size() == 0) { + chunk.source.reset(); + chunk.keep.reset(); + } +} + +std::size_t WeightingPipeline::start_jobs(ResultQueue& result_queue, bool flush) { + std::size_t started = 0; + for (std::size_t group_index = 0; auto& group : _groups) { + for (std::size_t context_index = 0; context_index < _contexts.size(); + ++context_index) { + auto& buffers = group.buffers.at(context_index); + auto& sizes = group.buffer_sizes.at(context_index); + for (std::size_t stage = 1; stage <= buffers.size(); ++stage) { + auto& buffer = buffers.at(stage - 1); + auto& size = sizes.at(stage - 1); + if (buffer.empty() || (!flush && size < batch_size(context_index))) { + continue; + } + std::size_t job_id = job_id_flag | _job_counter++; + auto& job = + _jobs + .emplace( + job_id, + Job{ + .group = group_index, + .stage = stage, + .context_index = context_index, + .chunks = std::move(buffer), + } + ) + .first->second; + buffer.clear(); + size = 0; + result_queue.submit( + _contexts.at(context_index)->thread_pool(), job_id, [this, &job] { + run_job(job); + } + ); + ++started; + } + } + ++group_index; + } + return started; +} + +void WeightingPipeline::gather_events(Chunk& chunk, ContextRuntimes& runtimes) { + auto& source = *chunk.source; + std::size_t count = chunk.size(); + if (!chunk.keep && chunk.offset == 0 && source.at(0).size(0) == count) { + chunk.events = source; + } else { + Tensor indices(DataType::dt_int, {count}); + auto view = indices.view(); + for (std::size_t i = 0; i < count; ++i) { + view[i] = static_cast( + chunk.offset + (chunk.keep ? chunk.keep->at(i) : i) + ); + } + auto device = _contexts.at(chunk.context_index)->device(); + if (device != cpu_device()) { + Tensor device_indices(DataType::dt_int, {count}, device); + device_indices.copy_from(indices); + indices = device_indices; + } + TensorVec args{indices}; + args.insert(args.end(), source.begin(), source.end()); + chunk.events = runtimes.gather->run(args); + } + chunk.source.reset(); + chunk.keep.reset(); + chunk.offset = 0; +} + +void WeightingPipeline::run_job(Job& job) { + auto& group = _groups.at(job.group); + auto& runtimes = group.runtimes.at(job.context_index); + for (auto& chunk : job.chunks) { + gather_events(*chunk, runtimes); + } + if (job.stage > group.settings.size()) { + // after the last step, the events go to the CPU to be written + for (auto& chunk : job.chunks) { + for (auto& tensor : chunk->events) { + chunk->cpu_events.push_back(tensor.cpu()); + } + chunk->events.clear(); + } + return; + } + + std::vector parts; + for (auto& chunk : job.chunks) { + parts.push_back(std::move(chunk->events)); + chunk->events.clear(); + } + // concatenate pairwise, so every event is copied log(chunks) times + while (parts.size() > 1) { + std::vector next_parts; + for (std::size_t i = 0; i < parts.size(); i += 2) { + if (i + 1 == parts.size()) { + next_parts.push_back(std::move(parts.at(i))); + } else { + TensorVec args = parts.at(i); + args.insert(args.end(), parts.at(i + 1).begin(), parts.at(i + 1).end()); + next_parts.push_back(runtimes.concat->run(args)); + } + } + parts = std::move(next_parts); + } + TensorVec events = std::move(parts.at(0)); + + TensorVec step_args{ + events.at(group.momenta), events.at(group.alpha_s), events.at(group.flavor) + }; + if (group.multiple_matrix_elements.at(job.stage - 1)) { + step_args.push_back(events.at(group.me_index)); + } + step_args.push_back(events.at(group.color)); + step_args.push_back(events.at(group.helicity)); + step_args.push_back(events.at(group.diagram)); + step_args.push_back(events.at(group.random)); + TensorVec outputs = runtimes.steps.at(job.stage - 1)->run(step_args); + auto matrix_elements = to_vector(outputs.at(0).cpu()); + auto random = to_vector(outputs.at(4).cpu()); + events.at(group.color) = outputs.at(1); + events.at(group.helicity) = outputs.at(2); + events.at(group.diagram) = outputs.at(3); + + auto shared_events = std::make_shared(std::move(events)); + std::size_t offset = 0; + for (auto& chunk : job.chunks) { + std::size_t count = chunk->size(); + chunk->source = shared_events; + chunk->offset = offset; + chunk->step_matrix_elements.assign( + matrix_elements.begin() + offset, matrix_elements.begin() + offset + count + ); + chunk->step_random.assign( + random.begin() + offset, random.begin() + offset + count + ); + offset += count; + } +} + +void WeightingPipeline::commit_job(std::size_t job_id) { + auto node = _jobs.extract(job_id); + auto& job = node.mapped(); + for (auto& chunk : job.chunks) { + set_ready(chunk, job.stage); + } +} + +void WeightingPipeline::reset() { + _jobs.clear(); + for (auto& group : _groups) { + for (auto& buffers : group.buffers) { + for (auto& buffer : buffers) { + buffer.clear(); + } + } + for (auto& sizes : group.buffer_sizes) { + std::fill(sizes.begin(), sizes.end(), 0); + } + } + for (auto& channel_ready : _ready) { + for (auto& ready : channel_ready) { + ready.clear(); + } + } + // the sequences continue, as the dropped chunks are never committed + for (std::size_t channel_index = 0; channel_index < _channels.size(); + ++channel_index) { + for (auto& next_sequence : _next_sequences.at(channel_index)) { + next_sequence = _channel_sequences.at(channel_index); + } + } + _chunk_count = 0; +} diff --git a/madspace/src/phasespace/weighting_step.cpp b/madspace/src/phasespace/weighting_step.cpp new file mode 100644 index 0000000000..f8f5820ecd --- /dev/null +++ b/madspace/src/phasespace/weighting_step.cpp @@ -0,0 +1,216 @@ +#include "madspace/phasespace/weighting_step.hpp" + +#include "madspace/util.hpp" + +using namespace madspace; + +namespace { + +const std::array index_outputs{ + MatrixElement::color_index_out, + MatrixElement::helicity_index_out, + MatrixElement::diagram_index_out, +}; + +int output_position( + const MatrixElement& matrix_element, MatrixElement::MatrixElementOutput output +) { + auto& outputs = matrix_element.outputs(); + auto search = std::find(outputs.begin(), outputs.end(), output); + return search == outputs.end() ? -1 : static_cast(search - outputs.begin()); +} + +} // namespace + +WeightingStepFunction::WeightingStepFunction( + const WeightingStep& step, + std::size_t particle_count, + const std::array& write_indices, + std::size_t step_index, + std::size_t step_count +) : + FunctionGenerator( + "WeightingStepFunction", + [&] { + NamedVector arg_types{ + {"momenta", batch_four_vec_array(particle_count)}, + {"alpha_s", batch_float}, + {"flavor", batch_int}, + }; + if (step.matrix_elements.size() > 1) { + arg_types.push_back("matrix_element_index", batch_int); + } + arg_types.push_back("color_index", batch_int); + arg_types.push_back("helicity_index", batch_int); + arg_types.push_back("diagram_index", batch_int); + arg_types.push_back("random", batch_float_array(random_count * step_count)); + return arg_types; + }(), + { + {"matrix_element", batch_float}, + {"color_index", batch_int}, + {"helicity_index", batch_int}, + {"diagram_index", batch_int}, + {"unweight_random", batch_float}, + } + ), + _matrix_elements(step.matrix_elements), + _write_indices(write_indices), + _step_index(step_index) { + if (step_index >= step_count) { + throw std::invalid_argument("step index out of range"); + } + if (_matrix_elements.empty()) { + throw std::invalid_argument("a weighting step needs a matrix element"); + } + for (auto& matrix_element : _matrix_elements) { + if (matrix_element.outputs() != _matrix_elements.at(0).outputs()) { + throw std::invalid_argument( + "the matrix elements of a weighting step need the same outputs" + ); + } + if (matrix_element.particle_count() != particle_count) { + throw std::invalid_argument( + "the matrix elements of a weighting step need the particle count " + "of the events" + ); + } + if (output_position(matrix_element, MatrixElement::matrix_element_out) == -1) { + throw std::invalid_argument( + "the matrix elements of a weighting step must return the matrix " + "element" + ); + } + for (auto input : matrix_element.external_inputs()) { + switch (input) { + case MatrixElement::momenta_in: + case MatrixElement::alpha_s_in: + case MatrixElement::flavor_in: + case MatrixElement::random_color_in: + case MatrixElement::random_helicity_in: + case MatrixElement::random_diagram_in: + case MatrixElement::helicity_in: + case MatrixElement::diagram_in: + break; + default: + throw std::invalid_argument( + "unsupported input of a weighting step matrix element" + ); + } + } + if (matrix_element.external_inputs().size() != matrix_element.inputs().size()) { + throw std::invalid_argument( + "the random inputs of a weighting step matrix element are passed in " + "and must not be sampled internally" + ); + } + } +} + +NamedVector WeightingStepFunction::build_function_impl( + FunctionBuilder& fb, const NamedVector& args +) const { + ValueVec all_random = fb.unstack(args.at("random")); + ValueVec random( + all_random.begin() + random_count * _step_index, + all_random.begin() + random_count * (_step_index + 1) + ); + ValueVec indices{ + args.at("color_index"), args.at("helicity_index"), args.at("diagram_index") + }; + auto me_args = [&](const MatrixElement& matrix_element, auto&& arg) { + ValueVec ret; + for (auto input : matrix_element.inputs()) { + switch (input) { + case MatrixElement::momenta_in: + ret.push_back(arg(args.at("momenta"))); + break; + case MatrixElement::alpha_s_in: + ret.push_back(arg(args.at("alpha_s"))); + break; + case MatrixElement::flavor_in: + ret.push_back(arg(args.at("flavor"))); + break; + case MatrixElement::random_color_in: + ret.push_back(arg(random.at(0))); + break; + case MatrixElement::random_helicity_in: + ret.push_back(arg(random.at(1))); + break; + case MatrixElement::random_diagram_in: + ret.push_back(arg(random.at(2))); + break; + case MatrixElement::helicity_in: + ret.push_back(arg(indices.at(1))); + break; + case MatrixElement::diagram_in: + ret.push_back(arg(indices.at(2))); + break; + default: + break; + } + } + return ret; + }; + // matrix element plus the index outputs, missing ones are empty + auto select_outputs = + [&](const MatrixElement& matrix_element, const ValueVec& outputs) { + ValueVec ret{outputs.at( + output_position(matrix_element, MatrixElement::matrix_element_out) + )}; + for (auto output : index_outputs) { + int position = output_position(matrix_element, output); + ret.push_back(position == -1 ? Value() : outputs.at(position)); + } + return ret; + }; + + ValueVec results; + if (_matrix_elements.size() == 1) { + auto& matrix_element = _matrix_elements.at(0); + results = select_outputs( + matrix_element, + matrix_element + .build_function(fb, me_args(matrix_element, [](Value v) { return v; })) + .values() + ); + } else { + ValueVec split_indices = fb.batch_split_by_index( + args.at("matrix_element_index"), + static_cast(_matrix_elements.size()) + ); + std::vector split_results(1 + index_outputs.size()); + for (auto [matrix_element, split_index] : + zip(_matrix_elements, split_indices)) { + ValueVec outputs = select_outputs( + matrix_element, + matrix_element + .build_function( + fb, + me_args( + matrix_element, + [&](Value v) { return fb.batch_gather(split_index, v); } + ) + ) + .values() + ); + for (auto [output, split] : zip(outputs, split_results)) { + if (output) { + split.push_back(output); + split.push_back(split_index); + } + } + } + for (auto& split : split_results) { + results.push_back(split.empty() ? Value() : fb.batch_merge_by_index(split)); + } + } + + ValueVec outputs{results.at(0)}; + for (std::size_t i = 0; i < index_outputs.size(); ++i) { + Value result = results.at(i + 1); + outputs.push_back(result && _write_indices.at(i) ? result : indices.at(i)); + } + outputs.push_back(random.at(3)); + return {return_types().keys(), outputs}; +} diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 4d261fe73c..6b9c9e50fa 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -3361,6 +3361,102 @@ PYBIND11_MODULE(_madspace_py, m) { &ChannelEventGenerator::set_fixed_max_weight, py::arg("max_weight"), pydoc::doc("ChannelEventGenerator::set_fixed_max_weight") + ) + .def( + "max_weight", + &ChannelEventGenerator::max_weight, + pydoc::doc("ChannelEventGenerator::max_weight") + ) + .def( + "step_max_weight", + &ChannelEventGenerator::step_max_weight, + py::arg("step"), + pydoc::doc("ChannelEventGenerator::step_max_weight") + ) + .def( + "set_fixed_step_max_weight", + &ChannelEventGenerator::set_fixed_step_max_weight, + py::arg("step"), + py::arg("max_weight"), + pydoc::doc("ChannelEventGenerator::set_fixed_step_max_weight") + ); + + py::classh weighting_step( + m, "WeightingStep", pydoc::doc("WeightingStep") + ); + add_enum( + weighting_step, + "Mode", + { + {"unweight", WeightingStep::unweight}, + {"reweight", WeightingStep::reweight}, + } + ); + weighting_step + .def( + py::init([](const std::vector& matrix_elements, + WeightingStep::Mode mode, + double max_overweight_truncation, + std::size_t freeze_max_weight_after, + const std::vector& + override_outputs) { + return WeightingStep{ + .matrix_elements = matrix_elements, + .mode = mode, + .max_overweight_truncation = max_overweight_truncation, + .freeze_max_weight_after = freeze_max_weight_after, + .override_outputs = override_outputs, + }; + }), + py::arg("matrix_elements"), + py::arg("mode") = WeightingStep::unweight, + py::arg("max_overweight_truncation") = 0.01, + py::arg("freeze_max_weight_after") = 10000, + py::arg("override_outputs") = + std::vector{}, + pydoc::doc("WeightingStep") + ) + .def_readwrite( + "matrix_elements", + &WeightingStep::matrix_elements, + pydoc::doc("WeightingStep::matrix_elements") + ) + .def_readwrite("mode", &WeightingStep::mode, pydoc::doc("WeightingStep::mode")) + .def_readwrite( + "max_overweight_truncation", + &WeightingStep::max_overweight_truncation, + pydoc::doc("WeightingStep::max_overweight_truncation") + ) + .def_readwrite( + "freeze_max_weight_after", + &WeightingStep::freeze_max_weight_after, + pydoc::doc("WeightingStep::freeze_max_weight_after") + ) + .def_readwrite( + "override_outputs", + &WeightingStep::override_outputs, + pydoc::doc("WeightingStep::override_outputs") + ); + + py::classh(m, "WeightingGroup", pydoc::doc("WeightingGroup")) + .def( + py::init([](const std::vector& channel_indices, + const std::vector& steps) { + return WeightingGroup{ + .channel_indices = channel_indices, .steps = steps + }; + }), + py::arg("channel_indices"), + py::arg("steps"), + pydoc::doc("WeightingGroup") + ) + .def_readwrite( + "channel_indices", + &WeightingGroup::channel_indices, + pydoc::doc("WeightingGroup::channel_indices") + ) + .def_readwrite( + "steps", &WeightingGroup::steps, pydoc::doc("WeightingGroup::steps") ); py::classh(m, "PdfMemberSpec", pydoc::doc("PdfMemberSpec")) @@ -3845,7 +3941,8 @@ PYBIND11_MODULE(_madspace_py, m) { const std::vector>&, std::uint64_t, std::shared_ptr, - const GeneratorConfig&>(), + const GeneratorConfig&, + const std::vector&>(), py::arg("contexts"), py::arg("channels"), py::arg("seed"), @@ -3855,6 +3952,7 @@ PYBIND11_MODULE(_madspace_py, m) { EventGenerator::default_config, "EventGenerator.default_config" ), + py::arg("groups") = std::vector{}, pydoc::doc("EventGenerator::EventGenerator") ) .def( diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py new file mode 100644 index 0000000000..73dc0a87a0 --- /dev/null +++ b/madspace/tests/test_weighting_steps.py @@ -0,0 +1,278 @@ +"""Re- and unweighting steps after the first unweighting (WeightingStep). + +The integrand uses a deformed "surrogate" matrix element, and a later step +re- or unweights the events with the "true" one. The results have to agree with +a run that uses the true matrix element in the integrand, and must not depend +on the number of threads. Uses the e+ e- -> d d~ g setup of test_event_stream +with PythonFunction matrix elements. +""" + +import os + +import numpy as np +import pytest +from test_event_stream import ( + CM_ENERGY, + PERMUTATIONS, + PIDS, + TOPOLOGY, + make_config, + make_context, +) +from test_integrand_me_outputs import python_matrix_element + +import madspace as ms + + +def true_me(momenta): + return 1.0 + 3.0 * (momenta[:, 2, 0] / CM_ENERGY) ** 2 + + +def surrogate_me(momenta): + return true_me(momenta) * (1.0 + 0.3 * np.sin(momenta[:, 3, 1] / 50.0)) + + +def make_integrand(alphas_file, matrix_element, pt_min=20.0): + mapping = ms.PhaseSpaceMapping( + TOPOLOGY, + CM_ENERGY, + permutations=PERMUTATIONS, + leptonic=True, + cuts=ms.Cuts( + [ + ms.CutItem(ms.Observable(PIDS, ms.Observable.obs_pt, [[i]]), min=pt_min) + for i in (2, 3, 4) + ] + ), + ) + diff_xs = ms.DifferentialCrossSection( + matrix_element=matrix_element, + cm_energy=CM_ENERGY, + running_coupling=None, + energy_scale=ms.CachedScale(), + return_matrix_element=True, + ) + return ms.Integrand( + mapping, + [diff_xs], + running_coupling=ms.RunningCoupling(ms.AlphaSGrid(alphas_file)), + energy_scale=ms.EnergyScale(len(PIDS), 91.188), + channel_indices=[0], + ) + + +def generate( + run_dir, + integrand_me, + step_modes=(), + step_me=true_me, + target_count=5000, + threads=2, + pt_mins=(20.0,), + grouped=(0,), + seed=7, + step_outputs=("matrix_element",), + override_outputs=(), + integrand_outputs=("matrix_element",), +): + """Survey and generate, returns the status, the combined events and the + channels.""" + os.makedirs(run_dir, exist_ok=True) + context, alphas_file = make_context(run_dir, threads) + me0 = python_matrix_element(context, integrand_me, integrand_outputs) + me1 = python_matrix_element(context, step_me, step_outputs) + config = make_config() + config.target_count = target_count + config.cpu_batch_size = 500 + channels = [ + ms.ChannelEventGenerator( + contexts=[context], + integrand=make_integrand(alphas_file, me0, pt_min), + event_file=os.path.join(run_dir, f"events.{i}.npy"), + weight_file=os.path.join(run_dir, f"weights.{i}.npy"), + config=config, + subprocess_index=0, + name=str(i), + histograms=None, + ) + for i, pt_min in enumerate(pt_mins) + ] + steps = [ + ms.WeightingStep([me1], mode=mode, override_outputs=list(override_outputs)) + for mode in step_modes + ] + groups = [ms.WeightingGroup(list(grouped), steps)] if steps else [] + generator = ms.EventGenerator( + [context], channels, seed=seed, config=config, groups=groups + ) + generator.survey() + generator.generate() + output_file = os.path.join(run_dir, "combined.npy") + generator.combine_to_compact_npy(output_file) + return generator.status(), np.load(output_file), channels + + +@pytest.fixture(scope="module") +def reference(tmp_path_factory): + """The true matrix element in the integrand.""" + return generate( + str(tmp_path_factory.mktemp("reference")), true_me, target_count=20000 + ) + + +def check_compatible(status, events, reference): + ref_status, ref_events, _ = reference + error = np.hypot(status.error, ref_status.error) + assert abs(status.mean - ref_status.mean) < 4 * error + # the events are normalized to the cross section + assert np.mean(events["weight"]) == pytest.approx(status.mean_abs, rel=1e-9) + # mean energy of the first quark, weighted + energies = [ + np.sum(ev["weight"] * ev["part3_energy"]) / np.sum(ev["weight"]) + for ev in (events, ref_events) + ] + assert energies[0] == pytest.approx(energies[1], rel=0.01) + + +def test_unweight_step(tmp_path, reference): + status, events, channels = generate( + str(tmp_path), surrogate_me, ["unweight"], target_count=20000 + ) + assert len(events) == 20000 + check_compatible(status, events, reference) + # unweighted: apart from the overweights, all events have the same weight + weights = events["weight"] + assert np.mean(np.isclose(weights, np.min(weights))) > 0.98 + assert channels[0].step_max_weight(1) > 0 + + +def test_reweight_step(tmp_path, reference): + status, events, _ = generate( + str(tmp_path), surrogate_me, ["reweight"], target_count=20000 + ) + check_compatible(status, events, reference) + # the ratios of the matrix elements stay in the weights + weights = events["weight"] + assert np.std(weights) / np.mean(weights) > 0.1 + + +def test_reweight_then_unweight(tmp_path, reference): + """The unweighting step absorbs the ratio of the reweighting step: the + second step has the same matrix element, so it only unweights.""" + status, events, _ = generate( + str(tmp_path), surrogate_me, ["reweight", "unweight"], target_count=20000 + ) + check_compatible(status, events, reference) + weights = events["weight"] + assert np.mean(np.isclose(weights, np.min(weights))) > 0.98 + + +@pytest.mark.parametrize("step_modes", [["unweight"], ["reweight", "unweight"]]) +def test_reproducible(tmp_path, step_modes): + """Two grouped channels and one without steps, which share the batches + differently with different thread counts.""" + results = [ + generate( + str(tmp_path / str(threads)), + surrogate_me, + step_modes, + threads=threads, + pt_mins=(20.0, 40.0, 60.0), + grouped=(0, 2), + ) + for threads in (1, 4) + ] + (status_a, events_a, _), (status_b, events_b, _) = results + assert len(events_a) == 5000 + assert np.array_equal(events_a, events_b) + assert status_a.mean == status_b.mean + + +def test_index_outputs(tmp_path): + """A step fills in indices the integrand matrix element does not return, + and replaces the others only with override_outputs.""" + + def with_helicity(momenta): + return true_me(momenta), np.full(len(momenta), 1, np.int32) + + outputs = ("matrix_element", "helicity_index") + _, events, _ = generate( + str(tmp_path / "fill"), + surrogate_me, + ["unweight"], + step_me=with_helicity, + step_outputs=outputs, + ) + assert (events["helicity_index"] == 1).all() + assert (events["color_index"] == -1).all() + + def surrogate_with_helicity(momenta): + return surrogate_me(momenta), np.zeros(len(momenta), np.int32) + + for override, helicity in [((), 0), ((ms.MatrixElement.helicity_index_out,), 1)]: + _, events, _ = generate( + str(tmp_path / f"override{helicity}"), + surrogate_with_helicity, + ["unweight"], + step_me=with_helicity, + step_outputs=outputs, + integrand_outputs=outputs, + override_outputs=override, + ) + assert (events["helicity_index"] == helicity).all() + + +def test_invalid_groups(tmp_path): + context, alphas_file = make_context(str(tmp_path), 1) + me = python_matrix_element(context, true_me) + config = make_config() + + def channel(return_matrix_element=True): + integrand = make_integrand(alphas_file, me) + if not return_matrix_element: + diff_xs = ms.DifferentialCrossSection( + matrix_element=me, + cm_energy=CM_ENERGY, + running_coupling=None, + energy_scale=ms.CachedScale(), + ) + integrand = ms.Integrand( + ms.PhaseSpaceMapping( + TOPOLOGY, CM_ENERGY, permutations=PERMUTATIONS, leptonic=True + ), + [diff_xs], + running_coupling=ms.RunningCoupling(ms.AlphaSGrid(alphas_file)), + energy_scale=ms.EnergyScale(len(PIDS), 91.188), + channel_indices=[0], + ) + return ms.ChannelEventGenerator( + [context], integrand, "", "", config, 0, "c", None + ) + + step = ms.WeightingStep([me]) + with pytest.raises(ValueError, match="matrix element"): + ms.EventGenerator( + [context], [channel(False)], seed=1, groups=[ms.WeightingGroup([0], [step])] + ) + with pytest.raises(ValueError, match="one weighting group"): + ms.EventGenerator( + [context], + [channel()], + seed=1, + groups=[ms.WeightingGroup([0], [step]), ms.WeightingGroup([0], [step])], + ) + with pytest.raises(ValueError, match="sampled internally"): + sampled = ms.MatrixElement( + me.matrix_element_index(), + len(PIDS), + [ms.MatrixElement.momenta_in, ms.MatrixElement.random_helicity_in], + [ms.MatrixElement.matrix_element_out], + 1, + True, + ) + ms.EventGenerator( + [context], + [channel()], + seed=1, + groups=[ms.WeightingGroup([0], [ms.WeightingStep([sampled])])], + ) From 5d667106e38d4dcfcf9455a2c3674a9358c02a5a Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 12:01:37 +0200 Subject: [PATCH 07/88] madspace: release the GIL in EventGenerator.optimize and fix_max_weights Like survey and generate: with matrix elements or custom functions in Python, the worker threads otherwise wait for the GIL held by the calling thread. --- madspace/src/python/madspace.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 6b9c9e50fa..8ebda73b88 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -3971,11 +3971,13 @@ PYBIND11_MODULE(_madspace_py, m) { .def( "optimize", &EventGenerator::optimize, + py::call_guard(), pydoc::doc("EventGenerator::optimize") ) .def( "fix_max_weights", &EventGenerator::fix_max_weights, + py::call_guard(), pydoc::doc("EventGenerator::fix_max_weights") ) .def( From 900ec212e40108f55988f6554aed4dab4f6050af Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 12:04:28 +0200 Subject: [PATCH 08/88] madspace: weighting groups in JSON, fixed maxima of the steps weighting_groups_to_json/from_json store the groups with their matrix elements by index, like the channel functions store theirs, so that a gridpack can restore them. With include_estimates, the channel JSON also holds the maximum weights of the steps, which stay fixed when the channel is loaded again. Tested with fix_max_weights() and a reload into a fresh context. --- docs/source/madspace/cpp-api.rst | 2 + docs/source/madspace/python-api.rst | 2 + .../madspace/driver/weighting_pipeline.hpp | 7 + .../madspace/phasespace/matrix_element.hpp | 2 + madspace/src/driver/weighting_pipeline.cpp | 121 ++++++++++++++++++ madspace/src/python/madspace.cpp | 12 ++ madspace/tests/test_weighting_steps.py | 76 +++++++++++ 7 files changed, 222 insertions(+) diff --git a/docs/source/madspace/cpp-api.rst b/docs/source/madspace/cpp-api.rst index 1a88ec1966..dcc79f73bb 100644 --- a/docs/source/madspace/cpp-api.rst +++ b/docs/source/madspace/cpp-api.rst @@ -192,6 +192,8 @@ The high-level driver plus the tensor and device layer it runs on: integration, cpp/select_combine_channel_index cpp/set_lib_path cpp/set_simd_vector_size + cpp/weighting_groups_from_json + cpp/weighting_groups_to_json Utilities --------- diff --git a/docs/source/madspace/python-api.rst b/docs/source/madspace/python-api.rst index e9f61ff485..95caadfb9c 100644 --- a/docs/source/madspace/python-api.rst +++ b/docs/source/madspace/python-api.rst @@ -172,6 +172,8 @@ The high-level driver plus the tensor and device layer it runs on: integration, python/select_combine_channel_index python/set_lib_path python/set_simd_vector_size + python/weighting_groups_from_json + python/weighting_groups_to_json Utilities --------- diff --git a/madspace/include/madspace/driver/weighting_pipeline.hpp b/madspace/include/madspace/driver/weighting_pipeline.hpp index 2610ddc304..5f18e93342 100644 --- a/madspace/include/madspace/driver/weighting_pipeline.hpp +++ b/madspace/include/madspace/driver/weighting_pipeline.hpp @@ -23,6 +23,13 @@ struct WeightingGroup { std::vector steps; }; +/// JSON text of `groups`, for example to store them in a gridpack. The matrix +/// elements are stored by their index in the context. +std::string weighting_groups_to_json(const std::vector& groups); +/// The groups stored with @ref weighting_groups_to_json. The matrix elements +/// must be loaded into the contexts in the same order as before. +std::vector weighting_groups_from_json(const std::string& text); + /** * Runs the re- and unweighting steps of the channel groups of an * @ref EventGenerator. diff --git a/madspace/include/madspace/phasespace/matrix_element.hpp b/madspace/include/madspace/phasespace/matrix_element.hpp index f677765d78..a74f67a4a7 100644 --- a/madspace/include/madspace/phasespace/matrix_element.hpp +++ b/madspace/include/madspace/phasespace/matrix_element.hpp @@ -142,6 +142,8 @@ class MatrixElement : public FunctionGenerator { const std::vector& inputs() const { return _inputs; } /// The configured output quantities. const std::vector& outputs() const { return _outputs; } + /// Whether the `random_*` inputs are drawn internally. + bool sample_random_inputs() const { return _sample_random_inputs; } /// The inputs that must be supplied by the caller (excludes the ones drawn /// internally when @p sample_random_inputs is set). std::vector external_inputs() const; diff --git a/madspace/src/driver/weighting_pipeline.cpp b/madspace/src/driver/weighting_pipeline.cpp index f1c68c1dcc..47143c9e6d 100644 --- a/madspace/src/driver/weighting_pipeline.cpp +++ b/madspace/src/driver/weighting_pipeline.cpp @@ -1,11 +1,14 @@ #include "madspace/driver/weighting_pipeline.hpp" +#include #include #include #include #include "madspace/util.hpp" +#include + using namespace madspace; namespace { @@ -58,8 +61,126 @@ Function concat_function(const NamedVector& types) { return fb.function(); } +using json = nlohmann::json; + +json matrix_element_to_json(const MatrixElement& matrix_element) { + std::vector me_frame; + for (std::size_t i = 0; auto flag : matrix_element.frame_mask()) { + if (flag) { + me_frame.push_back(static_cast(i + 1)); + } + ++i; + } + // see the MatrixElement constructor: an empty reference mask with a frame + // means that the frame is the one of the incoming particles + auto& reference_mask = matrix_element.reference_mask(); + auto& incoming_mask = + reference_mask.empty() ? matrix_element.frame_mask() : reference_mask; + std::size_t incoming_count = incoming_mask.empty() + ? 2 + : std::count(incoming_mask.begin(), incoming_mask.end(), 1); + json inputs = json::array(), outputs = json::array(); + for (auto input : matrix_element.inputs()) { + inputs.push_back(static_cast(input)); + } + for (auto output : matrix_element.outputs()) { + outputs.push_back(static_cast(output)); + } + return { + {"index", matrix_element.matrix_element_index()}, + {"particle_count", matrix_element.particle_count()}, + {"inputs", inputs}, + {"outputs", outputs}, + {"diagram_count", matrix_element.diagram_count()}, + {"sample_random_inputs", matrix_element.sample_random_inputs()}, + {"me_frame", me_frame}, + {"incoming_count", incoming_count}, + }; +} + +MatrixElement matrix_element_from_json(const json& j) { + std::vector inputs; + std::vector outputs; + for (int input : j.at("inputs")) { + inputs.push_back(static_cast(input)); + } + for (int output : j.at("outputs")) { + outputs.push_back(static_cast(output)); + } + return MatrixElement( + j.at("index").get(), + j.at("particle_count").get(), + inputs, + outputs, + j.at("diagram_count").get(), + j.at("sample_random_inputs").get(), + j.at("me_frame").get>(), + j.at("incoming_count").get() + ); +} + } // namespace +std::string +madspace::weighting_groups_to_json(const std::vector& groups) { + json j = json::array(); + for (auto& group : groups) { + json steps = json::array(); + for (auto& step : group.steps) { + json matrix_elements = json::array(), override_outputs = json::array(); + for (auto& matrix_element : step.matrix_elements) { + matrix_elements.push_back(matrix_element_to_json(matrix_element)); + } + for (auto output : step.override_outputs) { + override_outputs.push_back(static_cast(output)); + } + steps.push_back({ + {"matrix_elements", matrix_elements}, + {"mode", + step.mode == WeightingStep::unweight ? "unweight" : "reweight"}, + {"max_overweight_truncation", step.max_overweight_truncation}, + {"freeze_max_weight_after", step.freeze_max_weight_after}, + {"override_outputs", override_outputs}, + }); + } + j.push_back({{"channel_indices", group.channel_indices}, {"steps", steps}}); + } + return j.dump(); +} + +std::vector +madspace::weighting_groups_from_json(const std::string& text) { + std::vector groups; + for (auto& group_json : json::parse(text)) { + WeightingGroup group{ + .channel_indices = + group_json.at("channel_indices").get>() + }; + for (auto& step_json : group_json.at("steps")) { + WeightingStep step{ + .mode = step_json.at("mode").get() == "unweight" + ? WeightingStep::unweight + : WeightingStep::reweight, + .max_overweight_truncation = + step_json.at("max_overweight_truncation").get(), + .freeze_max_weight_after = + step_json.at("freeze_max_weight_after").get(), + }; + for (auto& me_json : step_json.at("matrix_elements")) { + step.matrix_elements.push_back(matrix_element_from_json(me_json)); + } + for (int output : step_json.at("override_outputs")) { + step.override_outputs.push_back( + static_cast(output) + ); + } + group.steps.push_back(std::move(step)); + } + groups.push_back(std::move(group)); + } + return groups; +} + void WeightingPipeline::Chunk::clear() { source.reset(); keep.reset(); diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 8ebda73b88..acbd265a22 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -3458,6 +3458,18 @@ PYBIND11_MODULE(_madspace_py, m) { .def_readwrite( "steps", &WeightingGroup::steps, pydoc::doc("WeightingGroup::steps") ); + m.def( + "weighting_groups_to_json", + &weighting_groups_to_json, + py::arg("groups"), + pydoc::doc("weighting_groups_to_json") + ); + m.def( + "weighting_groups_from_json", + &weighting_groups_from_json, + py::arg("text"), + pydoc::doc("weighting_groups_from_json") + ); py::classh(m, "PdfMemberSpec", pydoc::doc("PdfMemberSpec")) .def( diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py index 73dc0a87a0..28f35fdb29 100644 --- a/madspace/tests/test_weighting_steps.py +++ b/madspace/tests/test_weighting_steps.py @@ -7,6 +7,7 @@ with PythonFunction matrix elements. """ +import json import os import numpy as np @@ -276,3 +277,78 @@ def channel(return_matrix_element=True): seed=1, groups=[ms.WeightingGroup([0], [ms.WeightingStep([sampled])])], ) + + +def test_fixed_max_weights(tmp_path): + """fix_max_weights() freezes the maxima of all steps. Reloaded like in a + gridpack, with the groups from JSON, the channels keep them.""" + config = make_config() + run_dir = str(tmp_path / "fix") + os.makedirs(run_dir) + context, alphas_file = make_context(run_dir, 2) + surrogate = python_matrix_element(context, surrogate_me) + full = python_matrix_element(context, true_me) + step = ms.WeightingStep( + [full], freeze_max_weight_after=1000, max_overweight_truncation=0.05 + ) + groups = [ms.WeightingGroup([0, 1], [step])] + channels = [ + ms.ChannelEventGenerator( + contexts=[context], + integrand=make_integrand(alphas_file, surrogate, pt_min), + event_file=os.path.join(run_dir, f"events.{i}.npy"), + weight_file=os.path.join(run_dir, f"weights.{i}.npy"), + config=config, + subprocess_index=0, + name=str(i), + histograms=None, + ) + for i, pt_min in enumerate((20.0, 60.0)) + ] + generator = ms.EventGenerator( + [context], channels, seed=1, config=config, groups=groups + ) + generator.survey() + generator.fix_max_weights() + channel_jsons = [channel.to_json(True) for channel in channels] + groups_json = ms.weighting_groups_to_json(groups) + globals_file = os.path.join(run_dir, "globals.tar") + context.save_globals(globals_file) + max_weights = [ + (channel.max_weight(), channel.step_max_weight(1)) for channel in channels + ] + assert all(m0 > 0 and m1 > 0 for m0, m1 in max_weights) + for channel_json, (_, m1) in zip(channel_jsons, max_weights): + assert json.loads(channel_json)["step_max_weights"] == [m1] + + reload_dir = str(tmp_path / "reload") + os.makedirs(reload_dir) + new_context = ms.Context(ms.cpu_device(), 2) + new_context.load_globals(globals_file) + # same order as before, so the stored indices match + python_matrix_element(new_context, surrogate_me) + python_matrix_element(new_context, true_me) + new_channels = [ + ms.ChannelEventGenerator.load_json( + channel_json, + [new_context], + event_file=os.path.join(reload_dir, f"events.{i}.npy"), + weight_file=os.path.join(reload_dir, f"weights.{i}.npy"), + config=config, + ) + for i, channel_json in enumerate(channel_jsons) + ] + new_groups = ms.weighting_groups_from_json(groups_json) + assert ms.weighting_groups_to_json(new_groups) == groups_json + new_generator = ms.EventGenerator( + [new_context], new_channels, seed=2, config=config, groups=new_groups + ) + new_generator.generate() + assert [ + (channel.max_weight(), channel.step_max_weight(1)) for channel in new_channels + ] == max_weights + status, new_status = generator.status(), new_generator.status() + assert new_status.count_unweighted == config.target_count + assert abs(status.mean - new_status.mean) < 4 * np.hypot( + status.error, new_status.error + ) From 384d84527941dc63024a90b220eb43a4bb31b48a Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 12:22:31 +0200 Subject: [PATCH 09/88] madspace: release the GIL when an EventStream is destroyed Destroying the stream waits for its running jobs, which can need the GIL, for example with custom functions or matrix elements in Python. Destroyed from Python with the GIL held, this deadlocked. --- madspace/src/python/madspace.cpp | 27 +++++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index acbd265a22..0cd53de689 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -4046,12 +4046,27 @@ PYBIND11_MODULE(_madspace_py, m) { py::classh(m, "EventStream", pydoc::doc("EventStream")) .def( - py::init< - ContextPtr, - const std::vector>&, - std::uint64_t, - const LHECompleter&, - const GeneratorConfig&>(), + py::init( + [](ContextPtr context, + const std::vector>& channels, + std::uint64_t seed, + const LHECompleter& lhe_completer, + const GeneratorConfig& config) { + // Destroying the stream waits for its running jobs, which can + // need the GIL, e.g. for matrix elements in Python + return std::shared_ptr( + new EventStream(context, channels, seed, lhe_completer, config), + [](EventStream* stream) { + if (PyGILState_Check()) { + py::gil_scoped_release release; + delete stream; + } else { + delete stream; + } + } + ); + } + ), py::arg("context"), py::arg("channels"), py::arg("seed"), From 4fd1d75970ce801d17efa68a3eb0d49dcefaa49f Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 12:22:42 +0200 Subject: [PATCH 10/88] madspace: re- and unweighting steps in the EventStream The EventStream takes the weighting groups like the EventGenerator and runs the events of their channels through the same pipeline, which now hands the events after the last step to a sink: the EventGenerator writes them to the channel files, the stream buffers them as finished jobs. All maxima have to be fixed, so every decision only depends on the event itself. As for the first step, the integral after the steps only includes the jobs whose events were requested, so the stream stays a pure function of the seed. Also passes the biases of the jobs on to their integration, which the stream dropped before. --- docs/source/madspace/cpp-api.rst | 2 + .../madspace/driver/channel_generator.hpp | 80 +++++++---- .../include/madspace/driver/event_stream.hpp | 23 ++- .../madspace/driver/weighting_pipeline.hpp | 12 +- madspace/src/driver/channel_generator.cpp | 122 +++++++++------- madspace/src/driver/event_stream.cpp | 93 +++++++++--- madspace/src/driver/weighting_pipeline.cpp | 35 +++-- madspace/src/python/madspace.cpp | 8 +- madspace/tests/test_integrand_me_outputs.py | 2 + madspace/tests/test_weighting_steps.py | 132 +++++++++++++----- 10 files changed, 358 insertions(+), 151 deletions(-) diff --git a/docs/source/madspace/cpp-api.rst b/docs/source/madspace/cpp-api.rst index dcc79f73bb..783ce2b3d7 100644 --- a/docs/source/madspace/cpp-api.rst +++ b/docs/source/madspace/cpp-api.rst @@ -161,6 +161,8 @@ The high-level driver plus the tensor and device layer it runs on: integration, cpp/Sizes cpp/StatusFile cpp/StepChainResult + cpp/StepEvents + cpp/StepIntegral cpp/StepSettings cpp/SubprocessObservables cpp/SubprocessSystArgs diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 71f6dbf394..846f20f457 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -20,6 +20,35 @@ namespace madspace { +/// Events of one job that passed all re- and unweighting steps, see +/// @ref WeightingPipeline. +struct StepEvents { + /// Index of the channel. + std::size_t channel_index; + /// Id of the generation job of the events. + std::size_t job_id; + /// The events on the CPU, as returned by the unweighter. + TensorVec events; + /// The weights of the first step. + std::vector weights; + /// The maximum weight of the first step they were unweighted with. + double max_weight; + /// Per step, the ratios of the matrix elements. + std::vector> ratios; + /// Per step, the maximum weight they were unweighted with. + std::vector max_weights; + /// The biases of the events, empty without. + std::vector biases; + /// Number of points generated for these events. + std::size_t generated_count; +}; + +/// Sums of the final weights over the generated points. +struct StepIntegral { + std::size_t count = 0; + double sum = 0., sum2 = 0., abs_sum = 0., abs_sum2 = 0.; +}; + /** * Event generation and integration for one integration channel. * @@ -172,6 +201,15 @@ class ChannelEventGenerator { double step_max_weight(std::size_t step) const { return _step_states.at(step - 1).max_weight; } + /// Whether the maximum weights of all unweighting steps are fixed. + bool step_max_weights_fixed() const { + for (auto [state, settings] : zip(_step_states, _step_settings)) { + if (settings.unweight && !state.max_weight_fixed) { + return false; + } + } + return true; + } /// Fix the maximum weight of step `step`, like @ref set_fixed_max_weight. void set_fixed_step_max_weight(std::size_t step, double max_weight); /// Update the maximum weight of step `step` from the values it unweights, @@ -189,28 +227,19 @@ class ChannelEventGenerator { /// Incremented whenever the events are cleared, so that events still in the /// steps can be discarded. std::size_t event_epoch() const { return _event_epoch; } - /** - * Write events that passed all steps and add them to the integral. - * - * @param events The events, as returned by the unweighter. - * @param weights The weights of the first step. - * @param max_weight The maximum weight of the first step they were - * unweighted with. - * @param ratios Per step, the ratios of the events. - * @param max_weights Per step, the maximum weight they were unweighted - * with. - * @param biases The biases of the events, empty without. - * @param generated_count Number of points generated for these events. - */ - void write_step_events( - const TensorVec& events, - const std::vector& weights, - double max_weight, - const std::vector>& ratios, - const std::vector& max_weights, - const std::vector& biases, - std::size_t generated_count - ); + /// Write events that passed all steps and add them to the integral. + void write_step_events(const StepEvents& step_events); + /// The contribution of `step_events` to the integral of the final weights. + StepIntegral step_integral(const StepEvents& step_events) const; + /// Add to the integral of the final weights. + void add_step_integral(const StepIntegral& integral); + /// Pack `step_events` into an event buffer and a weight buffer (layouts + /// @ref event_file_layout and @ref weight_layout). + void fill_step_buffers( + const StepEvents& step_events, + EventBuffer& event_buffer, + EventBuffer& weight_buffer + ) const; /// Signed weight of event `index` of a buffer read from @ref weight_file, /// in units of the maximum weights, 0 for a rejected event. double unweighted_weight(EventBuffer& weight_buffer, std::size_t index) const; @@ -276,11 +305,6 @@ class ChannelEventGenerator { std::vector large_weights; double count_accepted = 0.; }; - // sums over the generated points of the final weights - struct FinalIntegral { - std::size_t count = 0; - double sum = 0., sum2 = 0., abs_sum = 0., abs_sum2 = 0.; - }; void apply_truncation_budget(); void update_status_integrals(); @@ -376,7 +400,7 @@ class ChannelEventGenerator { std::vector _step_states; std::vector _loaded_step_max_weights; double _first_step_count = 0.; - FinalIntegral _final_integral; + StepIntegral _final_integral; std::size_t _event_epoch = 0; friend void to_json(nlohmann::json& j, const ChannelEventGenerator& channel); diff --git a/madspace/include/madspace/driver/event_stream.hpp b/madspace/include/madspace/driver/event_stream.hpp index ff0d0eea1a..7839f886df 100644 --- a/madspace/include/madspace/driver/event_stream.hpp +++ b/madspace/include/madspace/driver/event_stream.hpp @@ -18,6 +18,7 @@ #include "madspace/driver/lhe_output.hpp" #include "madspace/driver/random.hpp" #include "madspace/driver/thread_pool.hpp" +#include "madspace/driver/weighting_pipeline.hpp" namespace madspace { @@ -26,7 +27,9 @@ namespace madspace { * * Every channel needs a fixed maximum weight (see @ref * ChannelEventGenerator::set_fixed_max_weight), so that an event is final as - * soon as it is unweighted; channels without one are skipped. A coordinator + * soon as it is unweighted; channels without one are skipped. With re- and + * unweighting steps (see @ref WeightingStep), the maximum weights of the steps + * have to be fixed as well, and an event is final after the last step. A coordinator * thread keeps generation jobs running on the context's thread pool and * buffers each channel's unweighted events, in launch order, until they are * consumed. @ref next_events picks the channel of every event at random, @@ -38,7 +41,8 @@ namespace madspace { * neither on thread timing nor on how the requests are split into batches. * * Event weights are in units of the maximum weight: +-1, except for - * overweight events. + * overweight events. After reweighting steps, they also contain the ratios of + * the reweighting steps. */ class EventStream { public: @@ -58,13 +62,16 @@ class EventStream { * @param lhe_completer Completes the events into LHE events. * @param config Generator configuration; the batch sizes and cut * settings are used. + * @param groups Groups of channels with re- and unweighting steps, + * see @ref EventGenerator. */ EventStream( ContextPtr context, const std::vector>& channels, std::uint64_t seed, const LHECompleter& lhe_completer, - const GeneratorConfig& config + const GeneratorConfig& config, + const std::vector& groups = {} ); ~EventStream(); EventStream(const EventStream&) = delete; @@ -98,6 +105,10 @@ class EventStream { EventBuffer weights; bool opened = false; std::size_t next_event = 0; + // with steps: the integral after the steps and the sum of the + // absolute event weights + StepIntegral step_integral; + double abs_weight_sum = 0.; }; struct ChannelState { bool active = false; @@ -151,11 +162,17 @@ class EventStream { std::size_t _results_pending = 0; std::size_t _job_id = 0; ResultQueue _result_queue; + std::unique_ptr _pipeline; + // jobs whose events are in the steps + std::unordered_map _step_jobs; std::thread _coordinator; void run_coordinator(); void dispatch_jobs(); void handle_result(std::size_t job_id); + void + close_job(std::size_t job_id, std::size_t channel_index, ClosedJob&& closed_job); + void close_step_events(StepEvents& step_events); double expected_job_events(std::size_t channel_index) const; void update_probabilities(); std::size_t select_channel(); diff --git a/madspace/include/madspace/driver/weighting_pipeline.hpp b/madspace/include/madspace/driver/weighting_pipeline.hpp index 5f18e93342..a61a6c6a47 100644 --- a/madspace/include/madspace/driver/weighting_pipeline.hpp +++ b/madspace/include/madspace/driver/weighting_pipeline.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -48,11 +49,18 @@ class WeightingPipeline { /// Job ids of the pipeline have this bit set. static constexpr std::size_t job_id_flag = std::size_t(1) << 62; + /// Receives the events after the last step, in the job order of each + /// channel. + using Sink = std::function; + + /// Without `sink`, the events are written to the channel files with + /// @ref ChannelEventGenerator::write_step_events. WeightingPipeline( const std::vector& contexts, const std::vector>& channels, const std::vector& groups, - const GeneratorConfig& config + const GeneratorConfig& config, + Sink sink = {} ); /// Whether channel `channel_index` has steps. bool contains(std::size_t channel_index) const { @@ -76,6 +84,7 @@ class WeightingPipeline { // the unweighted events of one job of a channel struct Chunk { std::size_t channel_index; + std::size_t job_id; std::size_t sequence; std::size_t epoch; std::size_t context_index; @@ -145,6 +154,7 @@ class WeightingPipeline { std::unordered_map _jobs; std::size_t _job_counter = 0; std::size_t _chunk_count = 0; + Sink _sink; }; } // namespace madspace diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 33117b395f..16097d9556 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -1246,72 +1246,90 @@ void ChannelEventGenerator::add_first_step_events( count * (job_max_weight > 0 ? job_max_weight / _max_weight : 1); } -void ChannelEventGenerator::write_step_events( - const TensorVec& events, - const std::vector& weights, - double max_weight, - const std::vector>& ratios, - const std::vector& max_weights, - const std::vector& biases, - std::size_t generated_count -) { +StepIntegral ChannelEventGenerator::step_integral(const StepEvents& step_events) const { std::size_t step_count = _step_settings.size(); - std::vector event_ratios(step_count); - // the final weights in units of the cross section, and the fraction of the - // events left by unweighting with the current maxima. Like in write_events(), - // the fraction is rescaled when a maximum rises, so it does not depend on - // the order of the updates. - double weight_unit = max_weight; - double count_factor = max_weight > 0 ? max_weight / _max_weight : 1.; - for (auto [state, settings, job_max_weight] : - zip(_step_states, _step_settings, max_weights)) { + // the final weights in units of the cross section + double weight_unit = step_events.max_weight; + for (auto [settings, job_max_weight] : + zip(_step_settings, step_events.max_weights)) { if (settings.unweight) { weight_unit *= job_max_weight; - count_factor *= job_max_weight > 0 ? job_max_weight / state.max_weight : 1.; } } - - _final_integral.count += generated_count; - for (std::size_t i = 0; i < weights.size(); ++i) { + StepIntegral integral{.count = step_events.generated_count}; + std::vector ratios(step_count); + for (std::size_t i = 0; i < step_events.weights.size(); ++i) { for (std::size_t k = 0; k < step_count; ++k) { - event_ratios.at(k) = ratios.at(k).at(i); + ratios.at(k) = step_events.ratios.at(k).at(i); + } + double weight = weight_unit * + evaluate_step_chain( + step_events.weights.at(i), + step_events.max_weight, + ratios.data(), + step_events.max_weights.data(), + _step_settings + ) + .weight; + double xs_weight = + step_events.biases.empty() ? weight : weight / step_events.biases.at(i); + integral.sum += xs_weight; + integral.sum2 += xs_weight * xs_weight; + integral.abs_sum += std::abs(weight); + integral.abs_sum2 += weight * weight; + } + return integral; +} + +void ChannelEventGenerator::add_step_integral(const StepIntegral& integral) { + _final_integral.count += integral.count; + _final_integral.sum += integral.sum; + _final_integral.sum2 += integral.sum2; + _final_integral.abs_sum += integral.abs_sum; + _final_integral.abs_sum2 += integral.abs_sum2; + update_status_integrals(); +} + +void ChannelEventGenerator::fill_step_buffers( + const StepEvents& step_events, EventBuffer& event_buffer, EventBuffer& weight_buffer +) const { + fill_event_buffers(step_events.events, event_buffer, weight_buffer); + for (std::size_t i = 0; i < step_events.weights.size(); ++i) { + weight_buffer.event(i).weight() = step_events.weights.at(i); + auto first = weight_buffer.particle(i, 0); + first.step_ratio() = 1.; + first.step_max_weight() = step_events.max_weight; + for (std::size_t k = 0; k < _step_settings.size(); ++k) { + auto record = weight_buffer.particle(i, k + 1); + record.step_ratio() = step_events.ratios.at(k).at(i); + record.step_max_weight() = step_events.max_weights.at(k); } - auto job_chain = evaluate_step_chain( - weights.at(i), - max_weight, - event_ratios.data(), - max_weights.data(), - _step_settings - ); - double weight = job_chain.weight * weight_unit; - double xs_weight = biases.empty() ? weight : weight / biases.at(i); - _final_integral.sum += xs_weight; - _final_integral.sum2 += xs_weight * xs_weight; - _final_integral.abs_sum += std::abs(weight); - _final_integral.abs_sum2 += weight * weight; } +} - if (weights.size() > 0) { +void ChannelEventGenerator::write_step_events(const StepEvents& step_events) { + add_step_integral(step_integral(step_events)); + std::size_t count = step_events.weights.size(); + if (count > 0) { EventBuffer event_buffer(0, _particle_count, _event_file_layout); - EventBuffer weight_buffer(0, step_count + 1, step_weight_file_layout); - fill_event_buffers(events, event_buffer, weight_buffer); - for (std::size_t i = 0; i < weights.size(); ++i) { - weight_buffer.event(i).weight() = weights.at(i); - auto first = weight_buffer.particle(i, 0); - first.step_ratio() = 1.; - first.step_max_weight() = max_weight; - for (std::size_t k = 0; k < step_count; ++k) { - auto record = weight_buffer.particle(i, k + 1); - record.step_ratio() = ratios.at(k).at(i); - record.step_max_weight() = max_weights.at(k); - } - } + EventBuffer weight_buffer(0, weight_record_count(), step_weight_file_layout); + fill_step_buffers(step_events, event_buffer, weight_buffer); event_file().write(event_buffer); weight_file().write(weight_buffer); } - _status.count_unweighted += weights.size() * count_factor; + // the fraction left by unweighting with the current maxima. Like in + // write_events(), it is rescaled when a maximum rises, so it does not + // depend on the order of the updates. + double count_factor = + step_events.max_weight > 0 ? step_events.max_weight / _max_weight : 1.; + for (auto [state, settings, job_max_weight] : + zip(_step_states, _step_settings, step_events.max_weights)) { + if (settings.unweight && job_max_weight > 0) { + count_factor *= job_max_weight / state.max_weight; + } + } + _status.count_unweighted += count * count_factor; _status.done = _status.count_unweighted >= _status.count_target; - update_status_integrals(); } StepChainResult diff --git a/madspace/src/driver/event_stream.cpp b/madspace/src/driver/event_stream.cpp index c9974c2e07..8779a10a91 100644 --- a/madspace/src/driver/event_stream.cpp +++ b/madspace/src/driver/event_stream.cpp @@ -23,7 +23,8 @@ EventStream::EventStream( const std::vector>& channels, std::uint64_t seed, const LHECompleter& lhe_completer, - const GeneratorConfig& config + const GeneratorConfig& config, + const std::vector& groups ) : _context(context), _channels(channels), @@ -34,13 +35,6 @@ EventStream::EventStream( if (channels.empty()) { throw std::invalid_argument("EventStream needs at least one channel"); } - for (auto& channel : channels) { - if (!channel->weighting_steps().empty()) { - throw std::invalid_argument( - "EventStream does not support weighting steps yet" - ); - } - } return combined_event_layout( channels.at(0)->event_layout_extra_flags(), channels.at(0)->particle_layout_extra_flags() @@ -54,9 +48,18 @@ EventStream::EventStream( ? config.cpu_batch_size : config.gpu_batch_size ) { + if (!groups.empty()) { + _pipeline = std::make_unique( + std::vector{context}, + channels, + groups, + config, + [this](StepEvents& step_events) { close_step_events(step_events); } + ); + } bool any_active = false; for (auto [channel, state] : zip(_channels, _channel_states)) { - state.active = channel->max_weight_fixed(); + state.active = channel->max_weight_fixed() && channel->step_max_weights_fixed(); any_active |= state.active; } if (!any_active) { @@ -102,6 +105,11 @@ void EventStream::run_coordinator() { } _dispatch_requested = false; dispatch_jobs(); + if (_pipeline) { + // partial batches only once no events are coming anymore + _results_pending += + _pipeline->start_jobs(_result_queue, _running_jobs.empty()); + } if (_results_pending == 0) { _coordinator_cv.wait(lock, [this] { return _stopping || _dispatch_requested; @@ -125,6 +133,10 @@ void EventStream::run_coordinator() { _result_queue.discard(_results_pending, nullptr); _results_pending = 0; _running_jobs.clear(); + if (_pipeline) { + _pipeline->reset(); + } + _step_jobs.clear(); } double EventStream::expected_job_events(std::size_t channel_index) const { @@ -206,6 +218,10 @@ void EventStream::dispatch_jobs() { } void EventStream::handle_result(std::size_t job_id) { + if (WeightingPipeline::owns_job(job_id)) { + _pipeline->commit_job(job_id); + return; + } auto& job = _running_jobs.at(job_id); auto& channel = _channels.at(job.channel_index); if (job.unweighted_events.empty()) { @@ -215,6 +231,17 @@ void EventStream::handle_result(std::size_t job_id) { ++_results_pending; return; } + if (_pipeline && _pipeline->contains(job.channel_index)) { + // keep what open_job() needs until the events are through the steps + auto& step_job = _step_jobs[job_id]; + step_job.channel_index = job.channel_index; + step_job.weights = job.weights; + step_job.biases = job.biases; + step_job.requested_event_count = job.requested_event_count; + _pipeline->add_events(job); + _running_jobs.erase(job_id); + return; + } ClosedJob closed_job{ .job = {}, @@ -228,10 +255,35 @@ void EventStream::handle_result(std::size_t job_id) { // only the weights are needed to integrate the job later closed_job.job.channel_index = job.channel_index; closed_job.job.weights = job.weights; + closed_job.job.biases = job.biases; closed_job.job.requested_event_count = job.requested_event_count; std::size_t channel_index = job.channel_index; _running_jobs.erase(job_id); + close_job(job_id, channel_index, std::move(closed_job)); +} + +void EventStream::close_step_events(StepEvents& step_events) { + auto& channel = _channels.at(step_events.channel_index); + auto node = _step_jobs.extract(step_events.job_id); + ClosedJob closed_job{ + .job = std::move(node.mapped()), + .events = + EventBuffer(0, channel->particle_count(), channel->event_file_layout()), + .weights = + EventBuffer(0, channel->weight_record_count(), channel->weight_layout()), + .step_integral = channel->step_integral(step_events), + }; + channel->fill_step_buffers(step_events, closed_job.events, closed_job.weights); + for (std::size_t i = 0; i < closed_job.weights.event_count(); ++i) { + closed_job.abs_weight_sum += + std::abs(channel->unweighted_weight(closed_job.weights, i)); + } + close_job(step_events.job_id, step_events.channel_index, std::move(closed_job)); +} +void EventStream::close_job( + std::size_t job_id, std::size_t channel_index, ClosedJob&& closed_job +) { std::unique_lock lock(_mutex); auto& state = _channel_states.at(channel_index); ++state.completed_jobs; @@ -312,12 +364,19 @@ void EventStream::open_job(std::size_t channel_index, ClosedJob& closed_job) { auto& channel = _channels.at(channel_index); auto& state = _channel_states.at(channel_index); channel->integrate(closed_job.job); - double max_weight = channel->max_weight(); - auto w_view = closed_job.job.weights.view(); - for (std::size_t i = 0; i < w_view.size(); ++i) { - double w = std::abs(w_view[i]); - state.opened_abs_sum += w; - state.opened_capped_sum += std::min(w, max_weight); + if (!channel->weighting_steps().empty()) { + channel->add_step_integral(closed_job.step_integral); + // the mean weight of the events after the steps + state.opened_abs_sum += closed_job.abs_weight_sum; + state.opened_capped_sum += closed_job.events.event_count(); + } else { + double max_weight = channel->max_weight(); + auto w_view = closed_job.job.weights.view(); + for (std::size_t i = 0; i < w_view.size(); ++i) { + double w = std::abs(w_view[i]); + state.opened_abs_sum += w; + state.opened_capped_sum += std::min(w, max_weight); + } } ++state.opened_jobs; --state.unopened_jobs; @@ -404,9 +463,7 @@ std::vector EventStream::next_events(std::size_t count) { event_index, buffer, i, - unweighted_event_weight( - job.weights.event(event_index).weight(), channel->max_weight() - ), + channel->unweighted_weight(job.weights, event_index), static_cast(channel->status().subprocess), channel->event_layout_extra_flags() ); diff --git a/madspace/src/driver/weighting_pipeline.cpp b/madspace/src/driver/weighting_pipeline.cpp index 47143c9e6d..68b34a2e05 100644 --- a/madspace/src/driver/weighting_pipeline.cpp +++ b/madspace/src/driver/weighting_pipeline.cpp @@ -211,7 +211,8 @@ WeightingPipeline::WeightingPipeline( const std::vector& contexts, const std::vector>& channels, const std::vector& groups, - const GeneratorConfig& config + const GeneratorConfig& config, + Sink sink ) : _contexts(contexts), _channels(channels), @@ -219,7 +220,13 @@ WeightingPipeline::WeightingPipeline( _channel_groups(channels.size(), -1), _channel_sequences(channels.size()), _next_sequences(channels.size()), - _ready(channels.size()) { + _ready(channels.size()), + _sink(sink) { + if (!_sink) { + _sink = [this](StepEvents& step_events) { + _channels.at(step_events.channel_index)->write_step_events(step_events); + }; + } const std::array index_outputs{ MatrixElement::color_index_out, MatrixElement::helicity_index_out, @@ -366,6 +373,7 @@ void WeightingPipeline::add_events(GeneratorBatchJob& job) { auto& channel = _channels.at(job.channel_index); auto chunk = std::make_shared(); chunk->channel_index = job.channel_index; + chunk->job_id = job.job_id; chunk->sequence = _channel_sequences.at(job.channel_index)++; chunk->epoch = channel->event_epoch(); chunk->context_index = job.context_index; @@ -430,18 +438,21 @@ void WeightingPipeline::commit_stage(ChunkPtr chunk, std::size_t stage) { advance(chunk, stage + 1); return; } + --_chunk_count; if (!stale) { - channel->write_step_events( - chunk->cpu_events, - chunk->weights, - chunk->max_weight, - chunk->ratios, - chunk->max_weights, - chunk->biases, - chunk->generated_count - ); + StepEvents step_events{ + .channel_index = chunk->channel_index, + .job_id = chunk->job_id, + .events = std::move(chunk->cpu_events), + .weights = std::move(chunk->weights), + .max_weight = chunk->max_weight, + .ratios = std::move(chunk->ratios), + .max_weights = std::move(chunk->max_weights), + .biases = std::move(chunk->biases), + .generated_count = chunk->generated_count, + }; + _sink(step_events); } - --_chunk_count; } void WeightingPipeline::decide(Chunk& chunk, std::size_t step) { diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 0cd53de689..38aa898a4b 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -4051,11 +4051,14 @@ PYBIND11_MODULE(_madspace_py, m) { const std::vector>& channels, std::uint64_t seed, const LHECompleter& lhe_completer, - const GeneratorConfig& config) { + const GeneratorConfig& config, + const std::vector& groups) { // Destroying the stream waits for its running jobs, which can // need the GIL, e.g. for matrix elements in Python return std::shared_ptr( - new EventStream(context, channels, seed, lhe_completer, config), + new EventStream( + context, channels, seed, lhe_completer, config, groups + ), [](EventStream* stream) { if (PyGILState_Check()) { py::gil_scoped_release release; @@ -4072,6 +4075,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("seed"), py::arg("lhe_completer"), py::arg("config"), + py::arg("groups") = std::vector{}, pydoc::doc("EventStream::EventStream") ) .def( diff --git a/madspace/tests/test_integrand_me_outputs.py b/madspace/tests/test_integrand_me_outputs.py index c146a9e5a0..b85e6cad15 100644 --- a/madspace/tests/test_integrand_me_outputs.py +++ b/madspace/tests/test_integrand_me_outputs.py @@ -29,7 +29,9 @@ def python_matrix_element(context, function, outputs=("matrix_element",)): ) output_enums = { "matrix_element": ms.MatrixElement.matrix_element_out, + "color_index": ms.MatrixElement.color_index_out, "helicity_index": ms.MatrixElement.helicity_index_out, + "diagram_index": ms.MatrixElement.diagram_index_out, } return ms.MatrixElement( api.index(), diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py index 28f35fdb29..c5d3b1083a 100644 --- a/madspace/tests/test_weighting_steps.py +++ b/madspace/tests/test_weighting_steps.py @@ -17,6 +17,7 @@ PERMUTATIONS, PIDS, TOPOLOGY, + lhe_completer, make_config, make_context, ) @@ -279,15 +280,30 @@ def channel(return_matrix_element=True): ) -def test_fixed_max_weights(tmp_path): - """fix_max_weights() freezes the maxima of all steps. Reloaded like in a - gridpack, with the groups from JSON, the channels keep them.""" +def full_me_with_indices(momenta): + zeros = np.zeros(len(momenta), np.int32) + return true_me(momenta), zeros, zeros, zeros + + +FULL_ME_OUTPUTS = ("matrix_element", "color_index", "helicity_index", "diagram_index") + + +def load_matrix_elements(context): + """The matrix elements, always in the same order, so that their indices + match the ones stored in the channels and groups.""" + surrogate = python_matrix_element(context, surrogate_me) + full = python_matrix_element(context, full_me_with_indices, FULL_ME_OUTPUTS) + return surrogate, full + + +@pytest.fixture(scope="module") +def fixed_gridpack(tmp_path_factory): + """Two grouped channels with all maxima fixed by fix_max_weights(), stored + like in a gridpack.""" config = make_config() - run_dir = str(tmp_path / "fix") - os.makedirs(run_dir) + run_dir = str(tmp_path_factory.mktemp("fixed")) context, alphas_file = make_context(run_dir, 2) - surrogate = python_matrix_element(context, surrogate_me) - full = python_matrix_element(context, true_me) + surrogate, full = load_matrix_elements(context) step = ms.WeightingStep( [full], freeze_max_weight_after=1000, max_overweight_truncation=0.05 ) @@ -310,45 +326,91 @@ def test_fixed_max_weights(tmp_path): ) generator.survey() generator.fix_max_weights() - channel_jsons = [channel.to_json(True) for channel in channels] - groups_json = ms.weighting_groups_to_json(groups) globals_file = os.path.join(run_dir, "globals.tar") context.save_globals(globals_file) - max_weights = [ - (channel.max_weight(), channel.step_max_weight(1)) for channel in channels - ] - assert all(m0 > 0 and m1 > 0 for m0, m1 in max_weights) - for channel_json, (_, m1) in zip(channel_jsons, max_weights): - assert json.loads(channel_json)["step_max_weights"] == [m1] - - reload_dir = str(tmp_path / "reload") - os.makedirs(reload_dir) - new_context = ms.Context(ms.cpu_device(), 2) - new_context.load_globals(globals_file) - # same order as before, so the stored indices match - python_matrix_element(new_context, surrogate_me) - python_matrix_element(new_context, true_me) - new_channels = [ + return { + "channels": [channel.to_json(True) for channel in channels], + "groups": ms.weighting_groups_to_json(groups), + "globals": globals_file, + "max_weights": [ + (channel.max_weight(), channel.step_max_weight(1)) for channel in channels + ], + "status": generator.status(), + } + + +def load_gridpack(gridpack, thread_count, run_dir=None): + context = ms.Context(ms.cpu_device(), thread_count) + context.load_globals(gridpack["globals"]) + load_matrix_elements(context) + config = make_config() + channels = [ ms.ChannelEventGenerator.load_json( channel_json, - [new_context], - event_file=os.path.join(reload_dir, f"events.{i}.npy"), - weight_file=os.path.join(reload_dir, f"weights.{i}.npy"), + [context], + event_file=os.path.join(run_dir, f"events.{i}.npy") if run_dir else "", + weight_file=os.path.join(run_dir, f"weights.{i}.npy") if run_dir else "", config=config, ) - for i, channel_json in enumerate(channel_jsons) + for i, channel_json in enumerate(gridpack["channels"]) ] - new_groups = ms.weighting_groups_from_json(groups_json) - assert ms.weighting_groups_to_json(new_groups) == groups_json - new_generator = ms.EventGenerator( - [new_context], new_channels, seed=2, config=config, groups=new_groups + groups = ms.weighting_groups_from_json(gridpack["groups"]) + return context, channels, groups, config + + +def test_fixed_max_weights(fixed_gridpack, tmp_path): + """fix_max_weights() freezes the maxima of all steps. Reloaded like in a + gridpack, with the groups from JSON, the channels keep them.""" + max_weights = fixed_gridpack["max_weights"] + assert all(m0 > 0 and m1 > 0 for m0, m1 in max_weights) + for channel_json, (_, m1) in zip(fixed_gridpack["channels"], max_weights): + assert json.loads(channel_json)["step_max_weights"] == [m1] + + context, channels, groups, config = load_gridpack(fixed_gridpack, 2, str(tmp_path)) + assert ms.weighting_groups_to_json(groups) == fixed_gridpack["groups"] + generator = ms.EventGenerator( + [context], channels, seed=2, config=config, groups=groups ) - new_generator.generate() + generator.generate() assert [ - (channel.max_weight(), channel.step_max_weight(1)) for channel in new_channels + (channel.max_weight(), channel.step_max_weight(1)) for channel in channels ] == max_weights - status, new_status = generator.status(), new_generator.status() + status, new_status = fixed_gridpack["status"], generator.status() assert new_status.count_unweighted == config.target_count assert abs(status.mean - new_status.mean) < 4 * np.hypot( status.error, new_status.error ) + + +def make_stream(fixed_gridpack, thread_count=4, seed=5): + context, channels, groups, config = load_gridpack(fixed_gridpack, thread_count) + return ms.EventStream( + context, channels, seed, lhe_completer(), config, groups=groups + ) + + +def stream_batches(stream, batch_sizes): + batches = [stream.next_batch(size) for size in batch_sizes] + return {key: np.concatenate([b[key] for b in batches]) for key in batches[0]} + + +def test_event_stream_reproducible(fixed_gridpack): + reference = stream_batches(make_stream(fixed_gridpack, 4), [3000]) + for thread_count, batch_sizes in [(1, [1000, 2000]), (3, [1, 999, 2000])]: + events = stream_batches(make_stream(fixed_gridpack, thread_count), batch_sizes) + for key, values in reference.items(): + np.testing.assert_array_equal(events[key], values, err_msg=key) + + +def test_event_stream_cross_section(fixed_gridpack): + """The streamed events have unit weights apart from overweights, and the + integral after the step matches the generator's.""" + stream = make_stream(fixed_gridpack) + batch = stream.next_batch(5000) + weights = batch["weight"] + assert np.all(weights >= 1.0) + assert np.mean(weights == 1.0) > 0.9 + status, stream_status = fixed_gridpack["status"], stream.status() + assert abs(status.mean - stream_status.mean) < 4 * np.hypot( + status.error, stream_status.error + ) From 3bcbdac97070a9cc7a45b5e15f8c9e4bd5b410c0 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 13:34:13 +0200 Subject: [PATCH 11/88] madspace: enum location of WeightingStep.Mode for the .pyi stubs The default mode of the WeightingStep constructor made pybind11-stubgen fail the wheel build. --- madspace/generate_pyi.py | 1 + 1 file changed, 1 insertion(+) diff --git a/madspace/generate_pyi.py b/madspace/generate_pyi.py index 2663d9363a..03633a792a 100644 --- a/madspace/generate_pyi.py +++ b/madspace/generate_pyi.py @@ -57,6 +57,7 @@ def main(): "--enum-class-locations=Activation:MLP.Activation", "--enum-class-locations=CutMode:Cuts.CutMode", "--enum-class-locations=LRSchedule:AdamOptimizer.LRSchedule", + "--enum-class-locations=Mode:WeightingStep.Mode", # Fail the build on unresolvable names/expressions (e.g. an enum # default value needing its own --enum-class-locations entry) # instead of silently emitting a stub with holes in it. From 3a044ae9ee2b62d1c7fa50dd2c46ffbe6e0841d3 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 14:33:16 +0200 Subject: [PATCH 12/88] mg7: [surrogate] run card section A surrogate of the matrix element replaces it in the integrand, and the full matrix element unweights or reweights the events in a second pass, as one weighting step per subprocess. The surrogate is a .so library with the inputs of the madmatrix libraries, loaded per subprocess with {subprocess} and {device} in its path, or a Python file defining custom_surrogate(momenta, pdg_ids, alpha_s). In unweight mode, the first_pass_* settings apply to the unweighting with the surrogate and those of [generation] to the second pass. Needs the denominators sde_strategy; no gridpacks yet. --- docs/source/run_card.rst | 43 ++++++ .../mg7/custom_function_loader.py | 84 +++++++++-- madgraph/iolibs/template_files/mg7/launch.py | 132 +++++++++++++++++- .../iolibs/template_files/mg7/run_card.toml | 12 ++ madgraph/various/banner.py | 14 ++ 5 files changed, 271 insertions(+), 14 deletions(-) diff --git a/docs/source/run_card.rst b/docs/source/run_card.rst index b36341f67d..0a2c95d62a 100644 --- a/docs/source/run_card.rst +++ b/docs/source/run_card.rst @@ -355,6 +355,49 @@ An empty string disables the function. See :doc:`custom_functions` for details. Custom event bias. Events are generated with the weight times the bias and written with the weight divided by the bias. +.. _run-card-surrogate: + +[surrogate] +----------- + +A fast surrogate of the matrix element, for example a neural network. The integrand and +its optimization use the surrogate. The events are then reweighted or unweighted a second +time with the full matrix element, which is only evaluated for the events that pass the +first unweighting. The cross section and the events are those of the full matrix element. +The surrogate does not return channel weights, so it needs ``sde_strategy = +"denominators"``. + +``function`` (``""``) + The surrogate. An empty string disables it. + + - A ``.so`` library implementing the UMAMI interface with the inputs of the madmatrix + libraries (momenta, :math:`\alpha_s`, flavor index and the random numbers). It is + loaded once per subprocess: ``{subprocess}`` and ``{device}`` in the path are replaced + by the subprocess name (e.g. ``P1_gg_ttx``) and the device, like in + ``lib/libmadmatrix_{subprocess}_{device}.so``. + - A Python file defining ``custom_surrogate(momenta, pdg_ids, alpha_s)``, which returns + :math:`|M|^2` for a batch of events. ``momenta`` has the shape ``(batch, n, 4)``, + ``pdg_ids`` the shape ``(batch, n)`` and identifies subprocess and flavor, and + ``alpha_s`` the shape ``(batch,)``. ``FRAMEWORK = "torch"`` passes torch tensors, + which also run on GPUs, otherwise numpy arrays are passed. See + :doc:`custom_functions`. + +``mode`` (``"unweight"``) + ``"unweight"`` unweights the events a second time with respect to the ratio of the full + matrix element to the surrogate. The events then have unit weights apart from + overweights. ``"reweight"`` keeps the ratio in the event weights. + +``first_pass_max_overweight_truncation`` (``0.01``), ``first_pass_freeze_max_weight_after`` (``10000``) + ``max_overweight_truncation`` and ``freeze_max_weight_after`` of the unweighting with + the surrogate, in unweight mode. The ones of ``[generation]`` are used for the second + pass, which gives the final events. In reweight mode, ``[generation]`` sets the + unweighting with the surrogate. + +``param_card`` (``""``) + The param card a ``.so`` surrogate is initialized with. An empty string uses + ``Cards/param_card.dat``. A madmatrix library with a different param card, for example, + is a surrogate for testing. + .. _run-card-multiparticles: [multiparticles] diff --git a/madgraph/iolibs/template_files/mg7/custom_function_loader.py b/madgraph/iolibs/template_files/mg7/custom_function_loader.py index b6b1da919c..b06f34150e 100644 --- a/madgraph/iolibs/template_files/mg7/custom_function_loader.py +++ b/madgraph/iolibs/template_files/mg7/custom_function_loader.py @@ -35,18 +35,10 @@ def _import_module(path): return _modules[path] -def load_custom_function(ms, contexts, device_types, kind, path): - """Load the custom ``kind`` function at ``path`` into each context, in the - same order everywhere. Returns the ``MatrixElementApi`` of each context.""" - if path.endswith(".so"): - return [context.load_matrix_element(path, "") for context in contexts] - if not path.endswith(".py"): - raise RuntimeError( - f"[custom_functions] {kind} = {path!r}: expected a .so library or a .py file" - ) - +def _python_function(path, name, device_types): + """The function ``name`` of the Python module at ``path`` and its + FRAMEWORK, checked against the devices of the run.""" module = _import_module(path) - name = f"custom_{kind}" function = getattr(module, name, None) if not callable(function): raise RuntimeError(f"{path} does not define a function {name}") @@ -66,7 +58,77 @@ def load_custom_function(ms, contexts, device_types, kind, path): logger.info("Limiting torch to one thread per madspace thread") torch.set_num_threads(1) + return function, framework + + +def load_custom_function(ms, contexts, device_types, kind, path): + """Load the custom ``kind`` function at ``path`` into each context, in the + same order everywhere. Returns the ``MatrixElementApi`` of each context.""" + if path.endswith(".so"): + return [context.load_matrix_element(path, "") for context in contexts] + if not path.endswith(".py"): + raise RuntimeError( + f"[custom_functions] {kind} = {path!r}: expected a .so library or a .py file" + ) + + name = f"custom_{kind}" + function, framework = _python_function(path, name, device_types) python_function = ms.PythonFunction( function, PYTHON_INPUTS[kind], [kind], framework=framework, name=name ) return [context.add_matrix_element(python_function) for context in contexts] + + +def load_surrogate( + ms, contexts, backends, device_types, path, subprocess, pdg_ids, param_card +): + """Load the [surrogate] function for one subprocess into each context, and + return the ``MatrixElementApi`` of each context. + + A .so path is a UMAMI library with the inputs of the madmatrix libraries, + where ``{subprocess}`` and ``{device}`` are replaced, initialized with + ``param_card``. A .py file defines + ``custom_surrogate(momenta, pdg_ids, alpha_s)``, which returns |M|^2; + ``pdg_ids[flavor_index]`` are the PDG ids of the subprocess for each + flavor index of its matrix element.""" + if path.endswith(".so"): + return [ + context.load_matrix_element( + path.format(subprocess=subprocess, device=backend), param_card + ) + for context, backend in zip(contexts, backends) + ] + if not path.endswith(".py"): + raise RuntimeError( + f"[surrogate] function = {path!r}: expected a .so library or a .py file" + ) + function, framework = _python_function(path, "custom_surrogate", device_types) + particle_count = len(pdg_ids[0]) + if framework == "torch": + import torch + + tables = {} + + def surrogate(momenta, flavor_index, alpha_s): + device = momenta.device + if device not in tables: + tables[device] = torch.tensor(pdg_ids, dtype=torch.int32, device=device) + return function(momenta, tables[device][flavor_index.long()], alpha_s) + + else: + import numpy as np + + table = np.array(pdg_ids, dtype=np.int32) + + def surrogate(momenta, flavor_index, alpha_s): + return function(momenta, table[flavor_index], alpha_s) + + python_function = ms.PythonFunction( + surrogate, + ["momenta", "flavor_index", "alpha_s"], + ["matrix_element"], + framework=framework, + particle_count=particle_count, + name=f"custom_surrogate ({subprocess})", + ) + return [context.add_matrix_element(python_function) for context in contexts] diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index ce315138e9..b871e59039 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -38,7 +38,10 @@ from madgraph.iolibs.template_files.mg7 import plots from madgraph.iolibs.template_files.mg7 import systematics_summary from madgraph.iolibs.template_files.mg7 import npy_to_lhe -from madgraph.iolibs.template_files.mg7.custom_function_loader import load_custom_function +from madgraph.iolibs.template_files.mg7.custom_function_loader import ( + load_custom_function, + load_surrogate, +) from madgraph.various.banner import RunCardMG7 from madgraph.various import misc from madgraph.interface.extended_cmd import Cmd @@ -289,6 +292,7 @@ def __init__(self): self.init_event_dir() self.init_context() self.init_custom_functions() + self.init_surrogate() self.init_cuts() self.init_histograms() self.init_generator_config() @@ -557,6 +561,55 @@ def init_custom_functions(self) -> None: self.custom_function_paths.append([kind, path]) logger.info("Using the custom %s function %s", kind, path) + def init_surrogate(self) -> None: + """Check the [surrogate] settings. The surrogates are loaded per + subprocess, by MadgraphSubprocess.""" + settings = self.run_card["surrogate"] + self.surrogate = None + if not settings["function"]: + return + path = settings["function"] + if path.endswith(".py") and not os.path.isfile(path): + raise RuntimeError(f"[surrogate] function = {path!r}: file not found") + if self.run_card["phasespace"]["sde_strategy"] != "denominators": + raise RuntimeError( + "[surrogate] needs [phasespace] sde_strategy = \"denominators\", " + "as the surrogate does not return channel weights" + ) + if self.run_card["run"]["dummy_matrix_element"]: + raise RuntimeError("[surrogate] cannot be used with dummy_matrix_element") + if self.run_card["gridpack"]["save_gridpack"]: + raise RuntimeError("[surrogate] does not support gridpacks yet") + self.surrogate = dict(settings) + if not self.surrogate["param_card"]: + self.surrogate["param_card"] = self.param_card_path + logger.info( + "Using the surrogate %s, %s with the full matrix element", + path, + "unweighting" if settings["mode"] == "unweight" else "reweighting", + ) + + def weighting_groups(self, subprocess_channels: list[list[int]]) -> list: + """One ms.WeightingGroup per subprocess with a [surrogate]: the second + pass with the full matrix element over its channels, given by their + indices in the event generator.""" + if self.surrogate is None: + return [] + gen_args = self.run_card["generation"] + return [ + ms.WeightingGroup( + channel_indices, + [ms.WeightingStep( + subproc.step_matrix_elements, + mode=self.surrogate["mode"], + max_overweight_truncation=gen_args["max_overweight_truncation"], + freeze_max_weight_after=gen_args["freeze_max_weight_after"], + )], + ) + for subproc, channel_indices in zip(self.subprocesses, subprocess_channels) + if channel_indices + ] + @staticmethod def build_custom_function(build_dir: str) -> None: """Run make in a custom_ directory, logging to compile.log.""" @@ -1240,6 +1293,14 @@ def init_generator_config(self) -> None: cfg.vegas_damping = vegas_args["damping"] cfg.max_overweight_truncation = gen_args["max_overweight_truncation"] cfg.freeze_max_weight_after = gen_args["freeze_max_weight_after"] + if self.surrogate is not None and self.surrogate["mode"] == "unweight": + # [generation] applies to the second pass, which gives the events + cfg.max_overweight_truncation = ( + self.surrogate["first_pass_max_overweight_truncation"] + ) + cfg.freeze_max_weight_after = ( + self.surrogate["first_pass_freeze_max_weight_after"] + ) cfg.start_batch_size = vegas_args["start_batch_size"] cfg.max_batch_size = vegas_args["max_batch_size"] cfg.survey_min_iters = gen_args["survey_min_iters"] @@ -1370,6 +1431,7 @@ def make_subprocesses( def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenerator: channel_generators = [] + subprocess_channels = [[] for _ in self.subprocesses] for i, (subproc, phasespace) in enumerate(zip(self.subprocesses, phasespaces)): for integrand, channel in zip( subproc.build_integrands(phasespace), phasespace.channels @@ -1385,6 +1447,7 @@ def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenera name=f"{i}.{channel.name}", histograms=subproc.histograms ) + subprocess_channels[i].append(len(channel_generators)) channel_generators.append(channel.event_generator) event_generator = ms.EventGenerator( @@ -1393,6 +1456,7 @@ def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenera status_file=self.status_file, config=self.event_generator_config, seed=self.run_seed, + groups=self.weighting_groups(subprocess_channels), ) unused_globals = ( set(self.contexts[0].global_names()) - event_generator.used_globals() @@ -2499,11 +2563,14 @@ def __init__( self.unmerged_meta = None if unmerged_meta is None: api_path_formats = [self.meta["me_path"]] + member_metas = [self.meta] else: api_path_formats = [] + member_metas = [] for subproc in self.meta["subprocesses"]: submeta = unmerged_meta[subproc] api_path_formats.append(submeta["me_path"]) + member_metas.append(submeta) if len(api_path_formats) == 1: self.meta = submeta else: @@ -2606,6 +2673,35 @@ def cut_pids(observable_kwargs): me_parameters(self.process.run_card) ) self.matrix_elements.append(mat) + self.surrogates = [] + if self.process.surrogate is not None: + for path_format, member_meta in zip(api_path_formats, member_metas): + self.surrogates.append(self.load_surrogate(path_format, member_meta)) + + def load_surrogate(self, me_path_format: str, member_meta: dict): + """The [surrogate] of one matrix element of the subprocess.""" + name = re.sub( + r"^libmadmatrix_(.*)_\{device\}\.so$", r"\1", os.path.basename(me_path_format) + ) + # PDG ids for each flavor index of the matrix element + flavors = member_meta["flavors"] + pdg_ids = [flavors[0]["options"][0]] * (max(f["index"] for f in flavors) + 1) + for flavor in flavors: + pdg_ids[flavor["index"]] = flavor["options"][0] + settings = self.process.surrogate + apis = load_surrogate( + ms, + self.process.contexts, + self.process.backends, + self.process.device_types, + settings["function"], + name, + pdg_ids, + settings["param_card"], + ) + if len({api.index() for api in apis}) != 1: + raise RuntimeError("surrogate loaded at different indices") + return apis[0] def flavor_pids(self) -> list[list[int]]: """PDG ids of the external particles for every flavor of the @@ -3250,8 +3346,37 @@ def build_integrands( flavor_mirror.append(flav["mirror"]) cross_sections = [] - for matrix_element in self.matrix_elements: - if matrix_element: + self.step_matrix_elements = [] + for me_index, matrix_element in enumerate(self.matrix_elements): + if self.surrogates: + # the surrogate in the integrand, the full matrix element in + # the second pass, with the random numbers passed in + mat = ms.MatrixElement( + self.surrogates[me_index].index(), + self.particle_count, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in], + [ms.MatrixElement.matrix_element_out], + 1, + False, + self.process.me_frame, + self.process.incoming_count, + ) + self.step_matrix_elements.append(ms.MatrixElement( + matrix_element, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in, ms.MatrixElement.random_color_in, + ms.MatrixElement.random_helicity_in, + ms.MatrixElement.random_diagram_in], + [ms.MatrixElement.matrix_element_out, + ms.MatrixElement.color_index_out, + ms.MatrixElement.helicity_index_out, + ms.MatrixElement.diagram_index_out], + False, + self.process.me_frame, + self.process.incoming_count, + )) + elif matrix_element: mat = ms.MatrixElement( matrix_element, ms.Integrand.matrix_element_inputs, @@ -3288,6 +3413,7 @@ def build_integrands( pdf2=pdf_arg, input_momentum_fraction=True, decay=self.process.is_decay, + return_matrix_element=bool(self.surrogates), ) ) partial_weights = self.process.systematics_enabled() diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 7c1b65efc2..ea618c4862 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -130,6 +130,18 @@ scale = %(custom_functions.scale)s # replaces dynamical_scale_choice pdf = %(custom_functions.pdf)s # replaces the PDF set, alpha_s still comes from [beam] pdf bias = %(custom_functions.bias)s # events are generated with |w| b and written with w / b +[surrogate] +# a fast surrogate of the matrix element in the integrand: a .so library, which can +# contain {subprocess} and {device}, or a Python file defining +# custom_surrogate(momenta, pdg_ids, alpha_s); "" disables it +function = %(surrogate.function)s +# unweight or reweight the events with the full matrix element in a second pass +mode = %(surrogate.mode)s +# unweighting with the surrogate (unweight mode); [generation] sets the second pass +first_pass_max_overweight_truncation = %(surrogate.first_pass_max_overweight_truncation)s +first_pass_freeze_max_weight_after = %(surrogate.first_pass_freeze_max_weight_after)s +param_card = %(surrogate.param_card)s # for a .so surrogate, "" uses Cards/param_card.dat + [multiparticles] $multiparticles diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 8e92ddabc5..c904f8691d 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6904,6 +6904,20 @@ def default_setup(self): self.add_toml_param('custom_functions', 'bias', "", comment="custom event bias") + # ---------------------------- [surrogate] --------------------- + # a fast surrogate of the matrix element for the integrand, followed by + # a second pass with the full matrix element; "" disables it + self.add_toml_param('surrogate', 'function', "", + comment="surrogate matrix element, a .so library or a .py file") + self.add_toml_param('surrogate', 'mode', "unweight", allowed=["unweight", "reweight"], + comment="unweight or reweight the events with the full matrix element") + self.add_toml_param('surrogate', 'first_pass_max_overweight_truncation', 0.01, + comment="max_overweight_truncation of the unweighting with the surrogate") + self.add_toml_param('surrogate', 'first_pass_freeze_max_weight_after', 10000, + comment="freeze_max_weight_after of the unweighting with the surrogate") + self.add_toml_param('surrogate', 'param_card', "", + comment="param card for a .so surrogate; \"\" uses Cards/param_card.dat") + # ----------------------------- [madnis] ----------------------- self.add_toml_param('madnis', 'enable', False, allowed=["auto", True, False], auto=True, comment="train MadNIS networks to improve the sampling; auto decides by process") From 133f045207131b632db22d2ca8bf9aabb47e8b74 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 14:55:16 +0200 Subject: [PATCH 13/88] madspace: unweighting efficiency of each pass in the generation output With weighting steps, the run summary shows the efficiency of the first unweighting, before and after cuts, and of each step relative to the one before, next to the overall one, their product. The channels keep showing the overall efficiency. The accepted counts of the steps are now rescaled like the final count when a maximum of an earlier step rises, and reweighting steps count their events too. --- .../madspace/driver/channel_generator.hpp | 23 ++++- .../madspace/driver/event_generator.hpp | 2 + madspace/src/driver/channel_generator.cpp | 27 +++++- madspace/src/driver/event_generator.cpp | 86 ++++++++++++++++++- madspace/src/driver/weighting_pipeline.cpp | 7 +- 5 files changed, 137 insertions(+), 8 deletions(-) diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 846f20f457..c7434ed177 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -217,9 +217,26 @@ class ChannelEventGenerator { void update_step_max_weight( std::size_t step, const std::vector& values, std::size_t count_target ); - /// Count `count` events accepted by step `step`. - void add_step_accepted(std::size_t step, double count) { - _step_states.at(step - 1).count_accepted += count; + /// Count `count` events accepted by step `step`, which passed the earlier + /// unweightings with the maxima `max_weight` (first step) and + /// `max_weights` (the steps before). + void add_step_accepted( + std::size_t step, + double count, + double max_weight, + const std::vector& max_weights + ); + /// Estimated number of events accepted by the first unweighting and each + /// step so far; without steps, just @ref GeneratorStatus::count_unweighted. + std::vector step_accepted_counts() const { + if (_step_settings.empty()) { + return {_status.count_unweighted}; + } + std::vector counts{_first_step_count}; + for (auto& state : _step_states) { + counts.push_back(state.count_accepted); + } + return counts; } /// Count `count` events accepted by the first unweighting with /// `job_max_weight`, with steps. diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 70b8c6ce57..eaaf1d511a 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -253,6 +253,8 @@ class EventGenerator { ); // run the re- and unweighting steps of all collected events void drain_pipeline(); + // unweighting efficiency of each pass, the first one before and after cuts + std::tuple> pass_efficiencies() const; // start the jobs of the steps whose batches are full std::size_t start_pipeline_jobs(); // Layout of the combined events as read from the channel files diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 16097d9556..1ab28e3498 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -816,6 +816,9 @@ void ChannelEventGenerator::apply_truncation_budget() { if (_max_weight < w) { _status.count_unweighted *= _max_weight / w; _first_step_count *= _max_weight / w; + for (auto& state : _step_states) { + state.count_accepted *= _max_weight / w; + } _max_weight = w; _unweighted_count = 0; _unweighted_accept_count = 0; @@ -1227,8 +1230,12 @@ void ChannelEventGenerator::update_step_max_weight( ++count; if (sum / value > max_truncation) { if (state.max_weight < value) { - _status.count_unweighted *= state.max_weight / value; - state.count_accepted *= state.max_weight / value; + // the events of this and the later steps + double factor = state.max_weight / value; + _status.count_unweighted *= factor; + for (std::size_t i = step - 1; i < _step_states.size(); ++i) { + _step_states.at(i).count_accepted *= factor; + } state.max_weight = value; _unweighted_count = 0; _unweighted_accept_count = 0; @@ -1239,6 +1246,22 @@ void ChannelEventGenerator::update_step_max_weight( large_weights.erase(large_weights.begin() + count, large_weights.end()); } +void ChannelEventGenerator::add_step_accepted( + std::size_t step, + double count, + double max_weight, + const std::vector& max_weights +) { + // the fraction left with the current maxima, like in write_step_events() + double factor = max_weight > 0 ? max_weight / _max_weight : 1.; + for (std::size_t i = 0; i + 1 < step; ++i) { + if (_step_settings.at(i).unweight && max_weights.at(i) > 0) { + factor *= max_weights.at(i) / _step_states.at(i).max_weight; + } + } + _step_states.at(step - 1).count_accepted += count * factor; +} + void ChannelEventGenerator::add_first_step_events( std::size_t count, double job_max_weight ) { diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index 4da3157c8b..a4051d8e5e 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -1425,6 +1425,27 @@ void EventGenerator::print_gen_init() { "Unweighted events:", "Run time:"} ); + if (_pipeline) { + _pretty_box_upper = PrettyBox( + _generate_mode == GenerateMode::fix_max_weights + ? "Integration and max. weight estimation" + : "Integration and unweighting", + 8, + {19, 0}, + offset + ); + _pretty_box_upper.set_column( + 0, + {"Result:", + "Rel. error:", + "Rel. stddev:", + "Number of events:", + "Unweighting eff.:", + " per pass:", + "Unweighted events:", + "Run time:"} + ); + } } _pretty_box_upper.print_first(); if (_channels.size() > 1) { @@ -1432,6 +1453,33 @@ void EventGenerator::print_gen_init() { } } +std::tuple> +EventGenerator::pass_efficiencies() const { + // events accepted per pass, a channel with fewer steps passes the others + std::vector accepted; + for (auto& channel : _channels) { + auto counts = channel->step_accepted_counts(); + if (counts.size() > accepted.size()) { + accepted.resize(counts.size(), 0.); + } + for (std::size_t i = 0; i < accepted.size(); ++i) { + accepted.at(i) += counts.at(std::min(i, counts.size() - 1)); + } + } + // the last pass ends with the counts of the final events, so that the + // product is the overall efficiency + accepted.back() = _status.count_unweighted; + std::vector later_passes; + for (std::size_t i = 1; i < accepted.size(); ++i) { + later_passes.push_back(accepted.at(i) / accepted.at(i - 1)); + } + return { + accepted.at(0) / _status.count_opt, + accepted.at(0) / _status.count_after_cuts_opt, + later_passes + }; +} + void EventGenerator::print_gen_update(bool done) { if (done) { add_timing_data("generate"); @@ -1453,7 +1501,7 @@ void EventGenerator::print_gen_update_pretty(bool done) { _last_print_time = now; bool optimize_only = _generate_mode == GenerateMode::optimize_only; - std::string int_str, rel_str, rsd_str, uweff_str, count_str; + std::string int_str, rel_str, rsd_str, uweff_str, pass_str, count_str; count_str = std::format( "{} before cuts, {} after", format_si_prefix(_status.count), @@ -1472,6 +1520,15 @@ void EventGenerator::print_gen_update_pretty(bool done) { _status.count_unweighted / _status.count_opt, _status.count_unweighted / _status.count_after_cuts_opt ); + if (_pipeline) { + auto [first, first_after_cuts, later] = pass_efficiencies(); + pass_str = + std::format("1: {:.5f} ({:.5f} after cuts)", first, first_after_cuts); + for (std::size_t i = 0; double eff : later) { + pass_str += std::format(", {}: {:.5f}", i + 2, eff); + ++i; + } + } } // the progress of the run: unweighted events, or optimized channels std::string progress_str; @@ -1505,6 +1562,18 @@ void EventGenerator::print_gen_update_pretty(bool done) { _pretty_box_upper.set_column( 1, {int_str, rel_str, rsd_str, count_str, progress_str, time_str} ); + } else if (_pipeline) { + _pretty_box_upper.set_column( + 1, + {int_str, + rel_str, + rsd_str, + count_str, + uweff_str, + pass_str, + progress_str, + time_str} + ); } else { _pretty_box_upper.set_column( 1, {int_str, rel_str, rsd_str, count_str, uweff_str, progress_str, time_str} @@ -1612,11 +1681,23 @@ void EventGenerator::print_gen_update_log(bool done) { ) ); } else { + std::string pass_efficiency_str; + if (_pipeline) { + auto [first, first_after_cuts, later] = pass_efficiencies(); + pass_efficiency_str = std::format( + ", unw. eff. per pass: {:.5f} ({:.5f} after cuts)", + first, + first_after_cuts + ); + for (double eff : later) { + pass_efficiency_str += std::format(", {:.5f}", eff); + } + } Logger::info( std::format( "generating, events: {} / {}, integral: {}, rel. error: {}, " "RSD: {:.3f}, samps: {}, samps. after cuts: {}, " - "unw. eff.: {:.5f}, unw. eff. after cuts: {:.5f}, time: {:%H:%M:%S}", + "unw. eff.: {:.5f}, unw. eff. after cuts: {:.5f}{}, time: {:%H:%M:%S}", format_si_prefix(_status.count_unweighted), format_si_prefix(_status.count_target), format_with_error(_status.mean, _status.error), @@ -1626,6 +1707,7 @@ void EventGenerator::print_gen_update_log(bool done) { format_si_prefix(_status.count_after_cuts), _status.count_unweighted / _status.count_opt, _status.count_unweighted / _status.count_after_cuts_opt, + pass_efficiency_str, std::chrono::round(now - _start_time) ) ); diff --git a/madspace/src/driver/weighting_pipeline.cpp b/madspace/src/driver/weighting_pipeline.cpp index 68b34a2e05..bad6b56b60 100644 --- a/madspace/src/driver/weighting_pipeline.cpp +++ b/madspace/src/driver/weighting_pipeline.cpp @@ -485,7 +485,9 @@ void WeightingPipeline::decide(Chunk& chunk, std::size_t step) { chunk.overweights.at(i) = std::max(1., value / max_weight); chunk.ratio_products.at(i) = std::copysign(1., chunk.ratio_products.at(i)); } - channel->add_step_accepted(step, static_cast(keep.size())); + channel->add_step_accepted( + step, static_cast(keep.size()), chunk.max_weight, chunk.max_weights + ); chunk.max_weights.push_back(max_weight); } else { for (std::size_t i = 0; i < count; ++i) { @@ -493,6 +495,9 @@ void WeightingPipeline::decide(Chunk& chunk, std::size_t step) { keep.push_back(static_cast(i)); } } + channel->add_step_accepted( + step, static_cast(keep.size()), chunk.max_weight, chunk.max_weights + ); chunk.max_weights.push_back(0.); } chunk.ratios.push_back(std::move(ratios)); From 5e24047035a14caf26be6c08b7a9ac59eed22179 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 21:41:53 +0200 Subject: [PATCH 14/88] madspace: weighting steps with matrix elements that return a ratio UMAMI 1.2 adds the output UMAMI_OUT_MATRIX_ELEMENT_RATIO, the ratio of the matrix element to a reference defined by the caller. A weighting step whose matrix elements return it (MatrixElement::matrix_element_ratio_out) uses the ratio directly, and the ratio times the previous matrix element is the reference of the next step. Integrands reject such matrix elements. --- .../madspace/driver/generator_data.hpp | 2 + .../madspace/phasespace/matrix_element.hpp | 5 +- .../madspace/phasespace/weighting_step.hpp | 12 ++- madspace/include/madspace/umami.h | 7 +- madspace/madspace/python_function.py | 1 + madspace/src/compgraphs/instruction.cpp | 1 + madspace/src/cpu/runtime.cpp | 1 + madspace/src/driver/channel_generator.cpp | 2 +- madspace/src/driver/weighting_pipeline.cpp | 11 ++- madspace/src/gpu/runtime.cu | 8 +- madspace/src/phasespace/cross_section.cpp | 8 ++ madspace/src/phasespace/matrix_element.cpp | 18 +++- madspace/src/phasespace/weighting_step.cpp | 24 ++++- madspace/src/python/madspace.cpp | 7 ++ madspace/tests/test_integrand_me_outputs.py | 1 + madspace/tests/test_weighting_steps.py | 93 ++++++++++++++++++- 16 files changed, 181 insertions(+), 20 deletions(-) diff --git a/madspace/include/madspace/driver/generator_data.hpp b/madspace/include/madspace/driver/generator_data.hpp index 3e488d4e6f..447a5aec5f 100644 --- a/madspace/include/madspace/driver/generator_data.hpp +++ b/madspace/include/madspace/driver/generator_data.hpp @@ -121,6 +121,8 @@ struct StepSettings { double max_overweight_truncation = 0.01; /// See @ref GeneratorConfig::freeze_max_weight_after. std::size_t freeze_max_weight_after = 10000; + /// Whether the step returns the ratio instead of the matrix element. + bool ratio = false; }; /// Result of @ref evaluate_step_chain. diff --git a/madspace/include/madspace/phasespace/matrix_element.hpp b/madspace/include/madspace/phasespace/matrix_element.hpp index a74f67a4a7..8d8a5a72f5 100644 --- a/madspace/include/madspace/phasespace/matrix_element.hpp +++ b/madspace/include/madspace/phasespace/matrix_element.hpp @@ -67,7 +67,10 @@ class MatrixElement : public FunctionGenerator { cut_out, ///< Custom cut value, 0 rejects the event. scale_out, ///< Custom energy scale. pdf_out, ///< Custom parton density x f(x, Q). - bias_out ///< Custom event bias. + bias_out, ///< Custom event bias. + /// Ratio of the matrix element to a reference matrix element, for a + /// @ref WeightingStep. + matrix_element_ratio_out }; /** diff --git a/madspace/include/madspace/phasespace/weighting_step.hpp b/madspace/include/madspace/phasespace/weighting_step.hpp index 37be0bc14f..a21f1b36d0 100644 --- a/madspace/include/madspace/phasespace/weighting_step.hpp +++ b/madspace/include/madspace/phasespace/weighting_step.hpp @@ -21,6 +21,11 @@ namespace madspace { * times the overweight factor carried from it, with its own maximum weight. A * reweighting step only multiplies the weights. * + * Instead of the matrix element, the matrix elements of a step can return the + * ratio directly, with the output `matrix_element_ratio_out`. Its matrix element + * is then the ratio times the one of the previous step, which is the reference + * for the ratio of the next step. + * * Typical uses are a fast surrogate of the matrix element in the integrand, * followed by an unweighting step with the full matrix element, or a * leading-color integrand reweighted to full color. @@ -46,6 +51,10 @@ struct WeightingStep { /// earlier steps. Without, an index is taken from the first step whose /// matrix element returns it. std::vector override_outputs; + + /// Whether the matrix elements return the ratio to the previous step + /// instead of the matrix element. + bool returns_ratio() const; }; /** @@ -66,7 +75,8 @@ struct WeightingStep { * element and for the unweighting. * * **Returns** - * - `matrix_element` – `float`, shape `(batch,)`. + * - `matrix_element` – `float`, shape `(batch,)` – the matrix element, or the + * ratio to the previous step if the step returns ratios. * - `color_index`, `helicity_index`, `diagram_index` – `int`, shape `(batch,)` * – the indices the step writes, otherwise the input values. * - `unweight_random` – `float`, shape `(batch,)` – the random number of this diff --git a/madspace/include/madspace/umami.h b/madspace/include/madspace/umami.h index a30959b2a2..4eb2dba380 100644 --- a/madspace/include/madspace/umami.h +++ b/madspace/include/madspace/umami.h @@ -30,7 +30,7 @@ extern "C" { * Minor version number of the UMAMI interface. Between minor versions, new keys for * errors, devices, metadata, inputs and outputs can be added. */ -#define UMAMI_MINOR_VERSION 1 +#define UMAMI_MINOR_VERSION 2 typedef enum { /** operation was executed successfully */ @@ -132,10 +132,13 @@ typedef enum { UMAMI_OUT_PDF, /** value of a custom event bias, type: `double`, shape: `()` */ UMAMI_OUT_BIAS, + /** ratio of the matrix element to a reference matrix element defined by the + * caller, type: `double`, shape: `()` */ + UMAMI_OUT_MATRIX_ELEMENT_RATIO, } UmamiOutputKey; /** Number of values in `UmamiOutputKey` */ -#define UMAMI_OUTPUT_KEY_COUNT (UMAMI_OUT_BIAS + 1) +#define UMAMI_OUTPUT_KEY_COUNT (UMAMI_OUT_MATRIX_ELEMENT_RATIO + 1) /** Implementation-defined pointer to a matrix element instance */ typedef void* UmamiHandle; diff --git a/madspace/madspace/python_function.py b/madspace/madspace/python_function.py index 4d0b8253fc..3202613ea3 100644 --- a/madspace/madspace/python_function.py +++ b/madspace/madspace/python_function.py @@ -36,6 +36,7 @@ "scale": DataType.float, "pdf": DataType.float, "bias": DataType.float, + "matrix_element_ratio": DataType.float, } diff --git a/madspace/src/compgraphs/instruction.cpp b/madspace/src/compgraphs/instruction.cpp index 85bc85fe3b..39e165a5c1 100644 --- a/madspace/src/compgraphs/instruction.cpp +++ b/madspace/src/compgraphs/instruction.cpp @@ -997,6 +997,7 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const { case UMAMI_OUT_SCALE: case UMAMI_OUT_PDF: case UMAMI_OUT_BIAS: + case UMAMI_OUT_MATRIX_ELEMENT_RATIO: dtype = DataType::dt_float; shape = {}; break; diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 9b4a717926..9d24f77720 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -87,6 +87,7 @@ void op_matrix_element( // implemented at LUMI hackathon where the coffee was indeed terrible switch (output_keys[i]) { case UMAMI_OUT_MATRIX_ELEMENT: + case UMAMI_OUT_MATRIX_ELEMENT_RATIO: device.submit([=]() mutable { double* ptr = static_cast(output_ptrs[i]); std::fill(ptr, ptr + batch_size, 1.); diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 1ab28e3498..b96e37595e 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -1123,7 +1123,7 @@ void ChannelEventGenerator::set_weighting_steps( steps.end(), _step_settings.begin(), [](auto& a, auto& b) { - return a.unweight == b.unweight && + return a.unweight == b.unweight && a.ratio == b.ratio && a.max_overweight_truncation == b.max_overweight_truncation && a.freeze_max_weight_after == b.freeze_max_weight_after; } diff --git a/madspace/src/driver/weighting_pipeline.cpp b/madspace/src/driver/weighting_pipeline.cpp index bad6b56b60..46b61dc0d8 100644 --- a/madspace/src/driver/weighting_pipeline.cpp +++ b/madspace/src/driver/weighting_pipeline.cpp @@ -306,6 +306,7 @@ WeightingPipeline::WeightingPipeline( .unweight = step.mode == WeightingStep::unweight, .max_overweight_truncation = step.max_overweight_truncation, .freeze_max_weight_after = step.freeze_max_weight_after, + .ratio = step.returns_ratio(), }); std::array write_indices{}; for (std::size_t i = 0; i < index_outputs.size(); ++i) { @@ -462,8 +463,14 @@ void WeightingPipeline::decide(Chunk& chunk, std::size_t step) { std::vector ratios(count); for (std::size_t i = 0; i < count; ++i) { double previous = chunk.matrix_elements.at(i); - ratios.at(i) = - previous != 0. ? chunk.step_matrix_elements.at(i) / previous : 0.; + double& value = chunk.step_matrix_elements.at(i); + if (settings.ratio) { + // the implied matrix element is the reference of the next step + ratios.at(i) = value; + value *= previous; + } else { + ratios.at(i) = previous != 0. ? value / previous : 0.; + } chunk.ratio_products.at(i) *= ratios.at(i); } std::vector keep; diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index 8e9a06a218..96260efa45 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -40,8 +40,8 @@ void op_matrix_element( std::size_t output_count = locals[instruction.input_indices[2]].index_value(); TensorVec contiguous_inputs(input_count); std::vector input_keys(input_count); - // one more output: the stream (UMAMI_OUT_GPU_STREAM), so that the matrix element runs - // in stream order with the rest of the function, not on the default stream + // one more output: the stream (UMAMI_OUT_GPU_STREAM), so that the matrix element + // runs in stream order with the rest of the function, not on the default stream std::vector output_keys(output_count + 1); std::vector input_ptrs(input_count), output_ptrs(output_count + 1); for (std::size_t i = 0; i < input_count; ++i) { @@ -76,6 +76,7 @@ void op_matrix_element( // implemented at LUMI hackathon where the coffee was indeed terrible switch (output_keys[i]) { case UMAMI_OUT_MATRIX_ELEMENT: + case UMAMI_OUT_MATRIX_ELEMENT_RATIO: thrust::fill_n( thrust_par.on(device.stream()), thrust::device_pointer_cast(static_cast(output_ptrs[i])), @@ -1408,7 +1409,8 @@ void op_histogram( Tensor square_weights_tmp( DataType::dt_float, {padded_size}, device, AllocHint::temporary ); - // the keys reduce_by_key writes: the n_bins bins, the underflow and the overflow bin + // the keys reduce_by_key writes: the n_bins bins, the underflow and the overflow + // bin Tensor reduce_tmp(DataType::dt_float, {n_bins + 2}, device, AllocHint::temporary); launch_kernel( diff --git a/madspace/src/phasespace/cross_section.cpp b/madspace/src/phasespace/cross_section.cpp index 40b4423b49..1d97a319fc 100644 --- a/madspace/src/phasespace/cross_section.cpp +++ b/madspace/src/phasespace/cross_section.cpp @@ -71,6 +71,14 @@ DifferentialCrossSection::DifferentialCrossSection( _input_momentum_fraction(input_momentum_fraction), _decay(decay), _return_matrix_element(return_matrix_element) { + auto& outputs = matrix_element.outputs(); + if (std::find( + outputs.begin(), outputs.end(), MatrixElement::matrix_element_ratio_out + ) != outputs.end()) { + throw std::invalid_argument( + "a matrix element returning a ratio can only be used in a weighting step" + ); + } if (decay) { if (_has_pdf.at(0) || _has_pdf.at(1)) { throw std::invalid_argument("a decay has no beams, so no PDFs"); diff --git a/madspace/src/phasespace/matrix_element.cpp b/madspace/src/phasespace/matrix_element.cpp index dabc80972b..c0ef852d4c 100644 --- a/madspace/src/phasespace/matrix_element.cpp +++ b/madspace/src/phasespace/matrix_element.cpp @@ -143,6 +143,9 @@ MatrixElement::MatrixElement( case bias_out: ret_types.push_back("bias", batch_float); break; + case matrix_element_ratio_out: + ret_types.push_back("matrix_element_ratio", batch_float); + break; default: throw std::invalid_argument("unknown output type"); } @@ -183,7 +186,17 @@ MatrixElement::MatrixElement( mask.clear(); } return mask; - }()) {} + }()) { + auto has_output = [&](MatrixElementOutput output) { + return std::find(outputs.begin(), outputs.end(), output) != outputs.end(); + }; + if (has_output(matrix_element_out) && has_output(matrix_element_ratio_out)) { + throw std::invalid_argument( + "a matrix element returns either matrix_element or matrix_element_ratio, " + "not both" + ); + } +} NamedVector MatrixElement::build_function_impl( FunctionBuilder& fb, const NamedVector& args @@ -320,6 +333,9 @@ NamedVector MatrixElement::build_function_impl( case bias_out: output_key = UMAMI_OUT_BIAS; break; + case matrix_element_ratio_out: + output_key = UMAMI_OUT_MATRIX_ELEMENT_RATIO; + break; } matrix_args.push_back(static_cast(output_key)); matrix_args.push_back(size_arg); diff --git a/madspace/src/phasespace/weighting_step.cpp b/madspace/src/phasespace/weighting_step.cpp index f8f5820ecd..c0e75f7075 100644 --- a/madspace/src/phasespace/weighting_step.cpp +++ b/madspace/src/phasespace/weighting_step.cpp @@ -22,6 +22,13 @@ int output_position( } // namespace +bool WeightingStep::returns_ratio() const { + return !matrix_elements.empty() && + output_position( + matrix_elements.at(0), MatrixElement::matrix_element_ratio_out + ) != -1; +} + WeightingStepFunction::WeightingStepFunction( const WeightingStep& step, std::size_t particle_count, @@ -75,10 +82,12 @@ WeightingStepFunction::WeightingStepFunction( "of the events" ); } - if (output_position(matrix_element, MatrixElement::matrix_element_out) == -1) { + if (output_position(matrix_element, MatrixElement::matrix_element_out) == -1 && + output_position(matrix_element, MatrixElement::matrix_element_ratio_out) == + -1) { throw std::invalid_argument( "the matrix elements of a weighting step must return the matrix " - "element" + "element or its ratio" ); } for (auto input : matrix_element.external_inputs()) { @@ -155,9 +164,14 @@ NamedVector WeightingStepFunction::build_function_impl( // matrix element plus the index outputs, missing ones are empty auto select_outputs = [&](const MatrixElement& matrix_element, const ValueVec& outputs) { - ValueVec ret{outputs.at( - output_position(matrix_element, MatrixElement::matrix_element_out) - )}; + int value_position = + output_position(matrix_element, MatrixElement::matrix_element_out); + if (value_position == -1) { + value_position = output_position( + matrix_element, MatrixElement::matrix_element_ratio_out + ); + } + ValueVec ret{outputs.at(value_position)}; for (auto output : index_outputs) { int position = output_position(matrix_element, output); ret.push_back(position == -1 ? Value() : outputs.at(position)); diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 38aa898a4b..42ed9d1c83 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -556,6 +556,7 @@ PYBIND11_MODULE(_madspace_py, m) { {"scale", UMAMI_OUT_SCALE}, {"pdf", UMAMI_OUT_PDF}, {"bias", UMAMI_OUT_BIAS}, + {"matrix_element_ratio", UMAMI_OUT_MATRIX_ELEMENT_RATIO}, }) { umami_output_keys[name] = key; } @@ -1550,6 +1551,7 @@ PYBIND11_MODULE(_madspace_py, m) { {"scale_out", MatrixElement::scale_out}, {"pdf_out", MatrixElement::pdf_out}, {"bias_out", MatrixElement::bias_out}, + {"matrix_element_ratio_out", MatrixElement::matrix_element_ratio_out}, } ); matrix_element @@ -3436,6 +3438,11 @@ PYBIND11_MODULE(_madspace_py, m) { "override_outputs", &WeightingStep::override_outputs, pydoc::doc("WeightingStep::override_outputs") + ) + .def( + "returns_ratio", + &WeightingStep::returns_ratio, + pydoc::doc("WeightingStep::returns_ratio") ); py::classh(m, "WeightingGroup", pydoc::doc("WeightingGroup")) diff --git a/madspace/tests/test_integrand_me_outputs.py b/madspace/tests/test_integrand_me_outputs.py index b85e6cad15..81b745bdf4 100644 --- a/madspace/tests/test_integrand_me_outputs.py +++ b/madspace/tests/test_integrand_me_outputs.py @@ -32,6 +32,7 @@ def python_matrix_element(context, function, outputs=("matrix_element",)): "color_index": ms.MatrixElement.color_index_out, "helicity_index": ms.MatrixElement.helicity_index_out, "diagram_index": ms.MatrixElement.diagram_index_out, + "matrix_element_ratio": ms.MatrixElement.matrix_element_ratio_out, } return ms.MatrixElement( api.index(), diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py index c5d3b1083a..fef4cbcf8b 100644 --- a/madspace/tests/test_weighting_steps.py +++ b/madspace/tests/test_weighting_steps.py @@ -78,11 +78,17 @@ def generate( integrand_outputs=("matrix_element",), ): """Survey and generate, returns the status, the combined events and the - channels.""" + channels. step_me and step_outputs can be lists, one entry per step.""" os.makedirs(run_dir, exist_ok=True) context, alphas_file = make_context(run_dir, threads) me0 = python_matrix_element(context, integrand_me, integrand_outputs) - me1 = python_matrix_element(context, step_me, step_outputs) + if not isinstance(step_me, list): + step_me = [step_me] * len(step_modes) + step_outputs = [step_outputs] * len(step_modes) + step_mes = [ + python_matrix_element(context, me, outputs) + for me, outputs in zip(step_me, step_outputs) + ] config = make_config() config.target_count = target_count config.cpu_batch_size = 500 @@ -100,8 +106,8 @@ def generate( for i, pt_min in enumerate(pt_mins) ] steps = [ - ms.WeightingStep([me1], mode=mode, override_outputs=list(override_outputs)) - for mode in step_modes + ms.WeightingStep([me], mode=mode, override_outputs=list(override_outputs)) + for me, mode in zip(step_mes, step_modes) ] groups = [ms.WeightingGroup(list(grouped), steps)] if steps else [] generator = ms.EventGenerator( @@ -224,6 +230,85 @@ def surrogate_with_helicity(momenta): assert (events["helicity_index"] == helicity).all() +def true_over_surrogate(momenta): + return true_me(momenta) / surrogate_me(momenta) + + +@pytest.mark.parametrize("step_modes", [["unweight"], ["reweight"]]) +def test_ratio_step(tmp_path, step_modes): + """A step that returns |M_true|^2 / |M_surrogate|^2 gives the same events + as a step that returns |M_true|^2: the pipeline computes the same ratio.""" + results = [ + generate( + str(tmp_path / name), + surrogate_me, + step_modes, + step_me=step_me, + step_outputs=outputs, + ) + for name, step_me, outputs in [ + ("me", true_me, ("matrix_element",)), + ("ratio", true_over_surrogate, ("matrix_element_ratio",)), + ] + ] + (status_me, events_me, _), (status_ratio, events_ratio, _) = results + assert np.array_equal(events_me, events_ratio) + assert status_me.mean == status_ratio.mean + + +def test_ratio_then_matrix_element(tmp_path, reference): + """After a ratio step, the matrix element it implies is the reference of the + next step, so the ratios still multiply to |M_true|^2 / |M_surrogate|^2.""" + + def intermediate_me(momenta): + return true_me(momenta) * (1.0 + 0.2 * np.cos(momenta[:, 2, 2] / 40.0)) + + status, events, _ = generate( + str(tmp_path), + surrogate_me, + ["reweight", "unweight"], + step_me=[lambda p: intermediate_me(p) / surrogate_me(p), true_me], + step_outputs=[("matrix_element_ratio",), ("matrix_element",)], + target_count=20000, + ) + check_compatible(status, events, reference) + weights = events["weight"] + assert np.mean(np.isclose(weights, np.min(weights))) > 0.98 + + +def test_ratio_matrix_element_rules(tmp_path): + context, _ = make_context(str(tmp_path), 1) + ratio_me = python_matrix_element( + context, true_over_surrogate, ("matrix_element_ratio",) + ) + assert ms.WeightingStep([ratio_me]).returns_ratio() + assert not ms.WeightingStep( + [python_matrix_element(context, true_me)] + ).returns_ratio() + with pytest.raises(ValueError, match="weighting step"): + ms.DifferentialCrossSection( + matrix_element=ratio_me, + cm_energy=CM_ENERGY, + running_coupling=None, + energy_scale=ms.CachedScale(), + ) + with pytest.raises(ValueError, match="not both"): + ms.MatrixElement( + ratio_me.matrix_element_index(), + len(PIDS), + [ms.MatrixElement.momenta_in], + [ + ms.MatrixElement.matrix_element_out, + ms.MatrixElement.matrix_element_ratio_out, + ], + ) + groups = [ms.WeightingGroup([0], [ms.WeightingStep([ratio_me])])] + text = ms.weighting_groups_to_json(groups) + loaded = ms.weighting_groups_from_json(text) + assert loaded[0].steps[0].returns_ratio() + assert ms.weighting_groups_to_json(loaded) == text + + def test_invalid_groups(tmp_path): context, alphas_file = make_context(str(tmp_path), 1) me = python_matrix_element(context, true_me) From 370abe871f9a155e75d5b65a435c98241459a2a7 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 22:23:41 +0200 Subject: [PATCH 15/88] madspace: matrix elements with a fixed channel index A non-negative channel passed to MatrixElement is given to the UMAMI channel input as a constant, which is then not an external input. Integrands and weighting steps can so evaluate one channel of a matrix element that treats its channels as independent processes. --- .../madspace/phasespace/matrix_element.hpp | 19 ++++++++++--- madspace/src/phasespace/matrix_element.cpp | 13 +++++++-- madspace/src/python/madspace.cpp | 9 +++++-- madspace/tests/test_python_function.py | 27 +++++++++++++++++++ 4 files changed, 60 insertions(+), 8 deletions(-) diff --git a/madspace/include/madspace/phasespace/matrix_element.hpp b/madspace/include/madspace/phasespace/matrix_element.hpp index 8d8a5a72f5..e30b442b14 100644 --- a/madspace/include/madspace/phasespace/matrix_element.hpp +++ b/madspace/include/madspace/phasespace/matrix_element.hpp @@ -98,6 +98,8 @@ class MatrixElement : public FunctionGenerator { * not commute, so going there straight from * the lab frame would leave a Wigner rotation * behind and rotate the polarisation axes. + * @param channel If non-negative, `channel_in` is passed this + * fixed value instead of being an input. */ MatrixElement( std::size_t matrix_element_index, @@ -107,7 +109,8 @@ class MatrixElement : public FunctionGenerator { std::size_t diagram_count = 1, bool sample_random_inputs = false, const std::vector& me_frame = {}, - std::size_t incoming_count = 2 + std::size_t incoming_count = 2, + me_int_t channel = -1 ); /** * Take @p matrix_element_index, @p particle_count and @p diagram_count from @@ -116,6 +119,9 @@ class MatrixElement : public FunctionGenerator { * @param inputs Quantities to pass in. * @param outputs Quantities to return. * @param sample_random_inputs See the other constructor. + * @param me_frame See the other constructor. + * @param incoming_count See the other constructor. + * @param channel See the other constructor. */ MatrixElement( const MatrixElementApi& matrix_element_api, @@ -123,7 +129,8 @@ class MatrixElement : public FunctionGenerator { const std::vector& outputs = {matrix_element_out}, bool sample_random_inputs = false, const std::vector& me_frame = {}, - std::size_t incoming_count = 2 + std::size_t incoming_count = 2, + me_int_t channel = -1 ) : MatrixElement( matrix_element_api.index(), @@ -133,7 +140,8 @@ class MatrixElement : public FunctionGenerator { matrix_element_api.diagram_count(), sample_random_inputs, me_frame, - incoming_count + incoming_count, + channel ) {}; /// Index of the matrix element in the plugin. std::size_t matrix_element_index() const { return _matrix_element_index; } @@ -147,8 +155,10 @@ class MatrixElement : public FunctionGenerator { const std::vector& outputs() const { return _outputs; } /// Whether the `random_*` inputs are drawn internally. bool sample_random_inputs() const { return _sample_random_inputs; } + /// The fixed channel passed as `channel_in`, or -1 if it is an input. + me_int_t channel() const { return _channel; } /// The inputs that must be supplied by the caller (excludes the ones drawn - /// internally when @p sample_random_inputs is set). + /// internally when @p sample_random_inputs is set and a fixed channel). std::vector external_inputs() const; /// Per-particle 0/1 selector built from me_frame, empty if no boost is /// applied. @@ -169,6 +179,7 @@ class MatrixElement : public FunctionGenerator { std::vector _inputs; std::vector _outputs; bool _sample_random_inputs; + me_int_t _channel; std::vector _frame_mask; std::vector _reference_mask; }; diff --git a/madspace/src/phasespace/matrix_element.cpp b/madspace/src/phasespace/matrix_element.cpp index c0ef852d4c..5f869136e8 100644 --- a/madspace/src/phasespace/matrix_element.cpp +++ b/madspace/src/phasespace/matrix_element.cpp @@ -40,7 +40,8 @@ MatrixElement::MatrixElement( std::size_t diagram_count, bool sample_random_inputs, const std::vector& me_frame, - std::size_t incoming_count + std::size_t incoming_count, + me_int_t channel ) : FunctionGenerator( "MatrixElement", @@ -81,7 +82,9 @@ MatrixElement::MatrixElement( arg_types.push_back("diagram", batch_int); break; case channel_in: - arg_types.push_back("channel", batch_int); + if (channel < 0) { + arg_types.push_back("channel", batch_int); + } break; case particle_count_in: arg_types.push_back("particle_count", batch_int); @@ -159,6 +162,7 @@ MatrixElement::MatrixElement( _particle_count(particle_count), _diagram_count(diagram_count), _sample_random_inputs(sample_random_inputs), + _channel(channel), _frame_mask( me_frame.empty() ? std::vector() : particle_mask(particle_count, me_frame, "me_frame") @@ -274,6 +278,8 @@ NamedVector MatrixElement::build_function_impl( input == random_diagram_in)) { matrix_args.push_back(random.at(random_index)); ++random_index; + } else if (input == channel_in && _channel >= 0) { + matrix_args.push_back(fb.full({_channel, fb.batch_size(args.values())})); } else if (input == momenta_in && !_frame_mask.empty()) { // Evaluate the matrix element in the frame the run card asked for // (me_frame). Only the matrix element sees the boosted momenta: @@ -346,6 +352,9 @@ NamedVector MatrixElement::build_function_impl( std::vector MatrixElement::external_inputs() const { std::vector ret; for (auto input : _inputs) { + if (input == channel_in && _channel >= 0) { + continue; + } if (!_sample_random_inputs || (input != random_color_in && input != random_helicity_in && input != random_diagram_in)) { diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 42ed9d1c83..8f39f4f12b 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1564,7 +1564,8 @@ PYBIND11_MODULE(_madspace_py, m) { std::size_t, bool, const std::vector&, - std::size_t>(), + std::size_t, + me_int_t>(), py::arg("matrix_element_index"), py::arg("particle_count"), // C++ defaults to {momenta_in} / {matrix_element_out}; kept required @@ -1576,6 +1577,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("sample_random_inputs") = false, py::arg("me_frame") = std::vector{}, py::arg("incoming_count") = 2, + py::arg("channel") = -1, pydoc::doc("MatrixElement::MatrixElement") ) .def( @@ -1585,13 +1587,15 @@ PYBIND11_MODULE(_madspace_py, m) { const std::vector&, bool, const std::vector&, - std::size_t>(), + std::size_t, + me_int_t>(), py::arg("matrix_element_api"), py::arg("inputs"), py::arg("outputs"), py::arg("sample_random_inputs") = false, py::arg("me_frame") = std::vector{}, py::arg("incoming_count") = 2, + py::arg("channel") = -1, pydoc::doc("MatrixElement::MatrixElement#2") ) .def( @@ -1611,6 +1615,7 @@ PYBIND11_MODULE(_madspace_py, m) { ) .def("inputs", &MatrixElement::inputs, pydoc::doc("MatrixElement::inputs")) .def("outputs", &MatrixElement::outputs, pydoc::doc("MatrixElement::outputs")) + .def("channel", &MatrixElement::channel, pydoc::doc("MatrixElement::channel")) .def( "external_inputs", &MatrixElement::external_inputs, diff --git a/madspace/tests/test_python_function.py b/madspace/tests/test_python_function.py index 4f290637d5..924094938a 100644 --- a/madspace/tests/test_python_function.py +++ b/madspace/tests/test_python_function.py @@ -91,6 +91,33 @@ def matrix_element(p): assert (np.asarray(helicity) == 3).all() +def test_fixed_channel(context, momenta): + def matrix_element(p, flavor, channel): + return p[:, 2, 0] * (channel + 10 * flavor) + + api = context.add_matrix_element( + ms.PythonFunction( + matrix_element, + ["momenta", "flavor_index", "channel_index"], + ["matrix_element"], + particle_count=4, + ) + ) + inputs = [ + ms.MatrixElement.momenta_in, + ms.MatrixElement.flavor_in, + ms.MatrixElement.channel_in, + ] + generator = ms.MatrixElement( + api.index(), 4, inputs, [ms.MatrixElement.matrix_element_out], channel=3 + ) + assert generator.channel() == 3 + assert generator.external_inputs() == inputs[:2] + flavor = np.full(len(momenta), 2, dtype=np.int32) + me = ms.FunctionRuntime(generator.function(), context)(momenta, flavor) + assert np.allclose(me, momenta[:, 2, 0] * 23) + + @pytest.mark.parametrize( "result, message", [ From 8b581c20fc7a1a1aba21c6881dab4e894315047b Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 22:24:58 +0200 Subject: [PATCH 16/88] madspace: integrands with discrete choices of the phase-space mapping The discrete two-solution choices of a color-ordered mapping were never passed to it by the integrand. They are now sampled uniformly. --- .../include/madspace/phasespace/integrand.hpp | 4 +- madspace/src/phasespace/integrand.cpp | 18 ++++++- madspace/tests/test_integrand_me_outputs.py | 48 +++++++++++++++++++ 3 files changed, 67 insertions(+), 3 deletions(-) diff --git a/madspace/include/madspace/phasespace/integrand.hpp b/madspace/include/madspace/phasespace/integrand.hpp index 3da4e6e9c8..1783bdca30 100644 --- a/madspace/include/madspace/phasespace/integrand.hpp +++ b/madspace/include/madspace/phasespace/integrand.hpp @@ -94,7 +94,9 @@ class Integrand : public FunctionGenerator { }; /** - * @param mapping The phase-space mapping. + * @param mapping The phase-space mapping. Its + * discrete choices, if any, are + * sampled uniformly. * @param diff_xs One differential cross section * per active flavour. * @param adaptive_map Optional continuous adaptive diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index 5e9893e8ef..868d5cfe5a 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -504,8 +504,22 @@ NamedVector Integrand::build_channel_part( chan_index_in_group = fb.full({static_cast(0), batch_size_val}); } - // Apply phase space mapping - auto mapping_result = _mapping.build_forward(fb, {latent}, mapping_conditions); + // Apply phase space mapping, its discrete choices are sampled uniformly + ValueVec mapping_inputs{latent}; + if (std::size_t discrete_dim = _mapping.discrete_dim(); discrete_dim > 0) { + ValueVec choices; + Value discrete_random = + fb.random(batch_size_val, static_cast(discrete_dim)); + for (auto& random : fb.unstack(discrete_random)) { + auto [choice, choice_det] = + fb.sample_discrete(random, static_cast(2)); + choices.push_back(choice); + weights_before_cuts.push_back(choice_det); + } + mapping_inputs.push_back(fb.stack(choices)); + } + auto mapping_result = + _mapping.build_forward(fb, mapping_inputs, mapping_conditions); weights_before_cuts.push_back(mapping_result["det"]); Value momenta = mapping_result["momenta"]; Value x0 = mapping_result["x1"]; diff --git a/madspace/tests/test_integrand_me_outputs.py b/madspace/tests/test_integrand_me_outputs.py index 81b745bdf4..82e400a5af 100644 --- a/madspace/tests/test_integrand_me_outputs.py +++ b/madspace/tests/test_integrand_me_outputs.py @@ -176,3 +176,51 @@ def test_diagram_channel_weights_need_amp2(tmp_path): first_chan_weight_remap=[[0, 1]], first_remapped_chan_count=2, ) + + +@pytest.mark.parametrize("color_order", [[0, 2, 3, 4, 1], [0, 1, 2, 3, 4]]) +def test_color_ordered_discrete_choices(tmp_path, color_order): + """The integrand samples the discrete two-solution choices of a color-ordered + mapping, so it reproduces the phase-space volume of RAMBO.""" + context, alphas_file = make_context(str(tmp_path), 2) + matrix_element = python_matrix_element( + context, lambda momenta: np.ones(len(momenta)) + ) + cuts = ms.Cuts( + [ + ms.CutItem(ms.Observable(PIDS, ms.Observable.obs_pt, [[i]]), min=20.0) + for i in (2, 3, 4) + ] + ) + results = [] + for mode, order in [ + (ms.PhaseSpaceMapping.color_ordered, color_order), + (ms.PhaseSpaceMapping.rambo, None), + ]: + mapping = ms.PhaseSpaceMapping( + [0.0] * len(PIDS), + CM_ENERGY, + leptonic=True, + mode=mode, + cuts=cuts, + color_order=order, + ) + if order is not None: + assert mapping.discrete_dim() > 0 + diff_xs = ms.DifferentialCrossSection( + matrix_element=matrix_element, + cm_energy=CM_ENERGY, + running_coupling=None, + energy_scale=ms.CachedScale(), + ) + integrand = ms.Integrand( + mapping, + [diff_xs], + running_coupling=ms.RunningCoupling(ms.AlphaSGrid(alphas_file)), + energy_scale=ms.EnergyScale(len(PIDS), 91.188), + channel_indices=[0], + ) + results.append(generate(str(tmp_path), context, integrand, str(mode))[0]) + status, rambo_status = results + error = np.hypot(status.error, rambo_status.error) + assert abs(status.mean - rambo_status.mean) < 5 * error From f59cecca997076e5798bb280c1abe4bbc7eb0191 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 22:26:53 +0200 Subject: [PATCH 17/88] madspace: LHE completion of subprocesses without diagrams A subprocess whose matrix element has no diagrams, like a color-ordered one, has no propagators to insert. Its events carry the diagram index -1. --- .../include/madspace/driver/lhe_output.hpp | 6 ++- madspace/src/driver/lhe_output.cpp | 11 +++-- madspace/tests/test_lhe.py | 44 +++++++++++++++++++ 3 files changed, 56 insertions(+), 5 deletions(-) diff --git a/madspace/include/madspace/driver/lhe_output.hpp b/madspace/include/madspace/driver/lhe_output.hpp index fde1385ea2..d580554edb 100644 --- a/madspace/include/madspace/driver/lhe_output.hpp +++ b/madspace/include/madspace/driver/lhe_output.hpp @@ -162,7 +162,8 @@ class LHECompleter { struct SubprocArgs { /// Process id shared by every event of this subprocess. int process_id; - /// Decay topology of each diagram channel. + /// Decay topology of each diagram channel. A subprocess without diagrams + /// has one topology and no permutations and diagram indices. std::vector topologies; /// Final-state particle permutation of each channel. nested_vector3 permutations; @@ -197,7 +198,8 @@ class LHECompleter { * @param event The event to complete, modified in place. * @param subprocess_index Index into the `subproc_args` passed to the * constructor. - * @param diagram_index Sampled diagram. + * @param diagram_index Sampled diagram, -1 for a subprocess without + * diagrams. * @param color_index Sampled color configuration. * @param flavor_index Sampled flavor assignment. * @param helicity_index Sampled helicity configuration. diff --git a/madspace/src/driver/lhe_output.cpp b/madspace/src/driver/lhe_output.cpp index 1e586ac138..245715e552 100644 --- a/madspace/src/driver/lhe_output.cpp +++ b/madspace/src/driver/lhe_output.cpp @@ -529,7 +529,10 @@ void LHECompleter::complete_event_data( if (event.particles.size() != subproc_data.particle_count) { throw std::runtime_error("Invalid particle number for subprocess"); } - if (diagram_index < 0 || diagram_index >= subproc_data.diagram_count) { + // a subprocess without diagrams has no propagators to insert + bool has_diagrams = subproc_data.diagram_count > 0; + if (has_diagrams ? diagram_index < 0 || diagram_index >= subproc_data.diagram_count + : diagram_index != -1) { throw std::runtime_error("Diagram index out of range"); } if (color_index < 0 || color_index >= subproc_data.color_count) { @@ -583,6 +586,9 @@ void LHECompleter::complete_event_data( ++particle_index; } + if (!has_diagrams) { + return; + } auto find_propagators = _propagator_index_and_count.find( cantor_pairing(subprocess_index, diagram_index, color_index) ); @@ -619,8 +625,7 @@ void LHECompleter::complete_event_data( // negative m_min would instead turn it into a large positive floor and // reject nearly everything, so a wider window would write *fewer* // resonances than a narrow one. - double m_min = - std::max(0., propagator.mass - _bw_cutoff * propagator.width); + double m_min = std::max(0., propagator.mass - _bw_cutoff * propagator.width); double m_max = propagator.mass + _bw_cutoff * propagator.width; if (m2 > m_min * m_min && m2 < m_max * m_max) { auto [color, anti_color] = prop_color; diff --git a/madspace/tests/test_lhe.py b/madspace/tests/test_lhe.py index c1c14fb07f..9df98bfbe8 100644 --- a/madspace/tests/test_lhe.py +++ b/madspace/tests/test_lhe.py @@ -461,3 +461,47 @@ def test_particle_momentum_kwargs_match_attributes(): for name in ("p_x", "p_y", "p_z"): with pytest.raises(TypeError): ms.LHEParticle(**{name: 1.0}) + + +def chain_topology(masses): + """t-channel chain topology without propagator information.""" + n_out = len(masses) - 2 + vertices = [["i0", "p0", "o0"]] + vertices += [[f"p{i - 1}", f"p{i}", f"o{i}"] for i in range(1, n_out - 1)] + vertices.append(["i1", f"p{n_out - 2}", f"o{n_out - 1}"]) + diagram = ms.Diagram( + masses[:2], masses[2:], [ms.Propagator() for _ in range(n_out - 1)], vertices + ) + return ms.Topology(diagram) + + +def test_subprocess_without_diagrams(): + """Without diagrams, events carry the diagram index -1 and get no + intermediate particles.""" + pdg_ids = [21, 21, 21, 21, 21] + color_flows = [ + [[502, 501], [503, 502], [503, 504], [505, 501], [505, 504]], + [[501, 502], [502, 503], [504, 503], [501, 505], [504, 505]], + ] + subproc_args = ms.SubprocArgs( + process_id=3, + topologies=[chain_topology([0.0] * 5)], + permutations=[[]], + diagram_indices=[[]], + diagram_color_indices=[[]], + color_flows=color_flows, + pdg_color_types={21: 8}, + helicities=[[1, 1, -1, -1, 1], [-1, -1, 1, 1, -1]], + pdg_ids=[[pdg_ids]], + ) + lhe_completer = ms.LHECompleter([subproc_args], bw_cutoff=BW_CUTOFF) + event = build_event(np.ones((5, 4)) * [100.0, 0.0, 0.0, 0.0]) + lhe_completer.complete_event_data(event, 0, -1, 1, 0, 1, ms.MixMaxRandom(1)) + assert event.process_id == 3 + assert [p.pdg_id for p in event.particles] == pdg_ids + assert [(p.color, p.anti_color) for p in event.particles] == [ + tuple(flow) for flow in color_flows[1] + ] + assert [p.spin for p in event.particles] == [-1, -1, 1, 1, -1] + with pytest.raises(RuntimeError, match="Diagram index out of range"): + lhe_completer.complete_event_data(event, 0, 0, 0, 0, 0, ms.MixMaxRandom(1)) From cf7f1b7e9ddef6d96e4353fa6e3b6b7430f4f55a Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:08:22 +0200 Subject: [PATCH 18/88] mg7: split the launcher into modules launch.py stays the entry point and re-exports what the launcher modules define. The code is moved without changes: common (madspace import, logging, helpers), observables, subprocess_madmatrix, madgraph_process, options, model_widths, postprocess, driver and cli. model_widths does not import madspace, so the run interface loads the model without it. --- madgraph/iolibs/template_files/mg7/cli.py | 470 ++ madgraph/iolibs/template_files/mg7/common.py | 200 + madgraph/iolibs/template_files/mg7/driver.py | 186 + madgraph/iolibs/template_files/mg7/launch.py | 4754 +---------------- .../template_files/mg7/madgraph_process.py | 1966 +++++++ .../iolibs/template_files/mg7/model_widths.py | 155 + .../iolibs/template_files/mg7/observables.py | 178 + madgraph/iolibs/template_files/mg7/options.py | 86 + .../iolibs/template_files/mg7/postprocess.py | 380 ++ .../template_files/mg7/run_interface.py | 2 +- .../mg7/subprocess_madmatrix.py | 1131 ++++ tests/unit_tests/interface/test_mg7_launch.py | 7 +- 12 files changed, 4828 insertions(+), 4687 deletions(-) create mode 100644 madgraph/iolibs/template_files/mg7/cli.py create mode 100644 madgraph/iolibs/template_files/mg7/common.py create mode 100644 madgraph/iolibs/template_files/mg7/driver.py create mode 100644 madgraph/iolibs/template_files/mg7/madgraph_process.py create mode 100644 madgraph/iolibs/template_files/mg7/model_widths.py create mode 100644 madgraph/iolibs/template_files/mg7/observables.py create mode 100644 madgraph/iolibs/template_files/mg7/options.py create mode 100644 madgraph/iolibs/template_files/mg7/postprocess.py create mode 100644 madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py diff --git a/madgraph/iolibs/template_files/mg7/cli.py b/madgraph/iolibs/template_files/mg7/cli.py new file mode 100644 index 0000000000..d096cca0d3 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/cli.py @@ -0,0 +1,470 @@ +"""The MG7> command line interface and its card editing.""" + +import argparse +import os +import re +import resource + +from madgraph.iolibs.template_files.mg7.common import logger +from madgraph.interface.extended_cmd import Cmd +from madgraph.various.banner import RunCardMG7 +from madgraph.iolibs.template_files.mg7 import npy_to_lhe +from madgraph.iolibs.template_files.mg7.driver import force_lhe_output_if_needed, run_generation +from madgraph.iolibs.template_files.mg7.model_widths import compute_auto_widths, load_process_model +from madgraph.iolibs.template_files.mg7.options import load_mg5_options +from madgraph.iolibs.template_files.mg7.postprocess import _off, _setup_logging + +class MG7Cmd(Cmd): + """The mg7 run interface: the command object that drives a launch. + + It plays two roles, and they are the same object on purpose. + + 1. It is the *mother* of the merged switch/card question (the thing + ``AskRunEditCard`` calls back into for ``keep_cards``/``do_open``/ + ``compute_widths``). + 2. It is the child command interface MG5 registers via + ``define_child_cmd_interface`` when you type ``launch`` on an mg7 output. + + Role 2 is what makes role 1 behave. ``Cmd.ask`` resolves a scripted answer + through ``self.check_answer_in_input_file``, where ``self`` is the mother -- + so once MG5 hands us its ``inputfile`` and registers itself as our + ``mother``, the question reads the very lines that follow ``launch`` in the + user's command file, and any line it does not recognise goes back up to MG5 + through ``store_line`` instead of being silently consumed. This is exactly + how a madevent output already behaves; before this class existed the mg7 + branch shelled out to ``bin/generate_events`` and had to guess which + following lines belonged to the run, which quietly swallowed every ``set``. + """ + + prompt = 'MG7> ' + + def __init__(self, me_dir='.', options=None, *args, **opts): + super().__init__(*args, **opts) + self.me_dir = os.path.abspath(me_dir) if me_dir != '.' else '.' + # The rest of the launcher is written against the process directory as + # the working directory, so do_generate_events chdirs there; but at + # construction time MG5 has not moved, hence the explicit me_dir. + self.options = load_mg5_options( + self.me_dir if self.me_dir != '.' else None) + if options: + # MG5's own resolved tool paths win: they are the ones the user + # configured in the session that is doing the launching. + self.options.update({k: v for k, v in options.items() + if k in self.options and v is not None}) + self.plugin_path = [] + self.proc_characteristics = {'grouped_matrix': False, 'limitations': []} + self._model = None # lazily imported by get_model() + + # ------------------------------------------------------------------ + # the contract the card question expects from its mother + # ------------------------------------------------------------------ + def keep_cards(self, need_card=[], ignore=[]): + from madgraph.interface.common_run_interface import CommonRunCmd + return CommonRunCmd.keep_cards(self, need_card, ignore) + + def do_open(self, line): + from madgraph.interface.common_run_interface import CommonRunCmd + CommonRunCmd.do_open(self, line) + + def check_open(self, args): + from madgraph.interface.common_run_interface import CommonRunCmd + CommonRunCmd.check_open(self, args) + + def get_model(self): + """The model this process was generated with (an ``import_ufo`` model). + + This is the hook the card question goes through for everything that + needs the model rather than the card: 'update dependent' (which resets + the masses/widths the model computes from the free parameters) and the + generic auto-width handler both call ``mother_interface.get_model()``. + Without it the launch question could only warn that it had failed to + update the dependent parameters, and MadSpin/the shower then read + whatever inconsistent values the param_card still held. + + None when the process recorded no model (an output made before + SubProcesses/model.txt existed) -- the caller warns in that case. + """ + if self._model is None: + me_dir = self.me_dir if self.me_dir != '.' else None + # False (not None) so a model that cannot be imported is not + # retried at every question. + self._model = load_process_model(me_dir) or False + return self._model or None + + def do_compute_widths(self, line): + # The interactive card editor delegates 'auto' width computation to + # the mother interface. Reuse the runtime helper (madgraph subprocess + # + the model stored at output time). ``line`` looks like + # " --path= [--nlo]"; we only need the card path. + m = re.search(r'--path=(\S+)', line or "") + path = m.group(1) if m else os.path.join("Cards", "param_card.dat") + compute_auto_widths(path) + # return an empty mapping: the caller iterates out.items() for the + # small-width treatment, which mg7 does not apply. + return {} + + # ------------------------------------------------------------------ + # running + # ------------------------------------------------------------------ + def do_generate_events(self, line): + """generate_events [-f] [-n NAME] [--laststep=parton] + + Run the mg7 generation and the selected post-processing tools.""" + + opts = self._parse_run_options(line) + cwd = os.getcwd() + target = self.me_dir if self.me_dir != '.' else cwd + try: + os.chdir(target) + _setup_logging() + if opts['name']: + self._set_run_name(opts['name']) + switch = {} + if not opts['force']: + switch = ask_edit_cards(mother=self) + switch = self._apply_laststep(switch, opts) + force_lhe_output_if_needed(switch) + _raise_open_file_limit() + run_generation(switch) + finally: + os.chdir(cwd) + + # ``launch`` is the name MG5 users type; keep it working here too. + do_launch = do_generate_events + + def do_npy_to_lhe(self, line): + """npy_to_lhe [RUN_NAME] [-o OUTPUT] [--seed N] [--no-gzip] + + Write the LHE file (events.lhe.gz) of a run that produced npy events, + e.g. to shower it. Without a run name, the most recent such run.""" + import shlex + me_dir = self.me_dir if self.me_dir != '.' else os.getcwd() + try: + npy_to_lhe.main(shlex.split(line), me_dir=me_dir) + except SystemExit as error: + if error.code not in (None, 0): + logger.error(str(error.code)) + + def do_quit(self, line): + """Leave the mg7 run interface.""" + return super().do_quit(line) + + do_exit = do_quit + + # ------------------------------------------------------------------ + # helpers + # ------------------------------------------------------------------ + def _parse_run_options(self, line): + """Understand the subset of MG5's `launch` options that mean something + for an mg7 run, and say so out loud for the ones that do not.""" + out = {'force': False, 'name': '', 'laststep': ''} + args = self.split_arg(line) + index = 0 + while index < len(args): + arg = args[index] + index += 1 + if arg in ('-f', '--force'): + out['force'] = True + elif arg in ('-n', '--name') and index < len(args): + out['name'] = args[index] + index += 1 + elif arg.startswith(('-n=', '--name=')): + out['name'] = arg.split('=', 1)[1] + elif arg in ('-s', '--laststep') and index < len(args): + out['laststep'] = args[index] + index += 1 + elif arg.startswith(('-s=', '--laststep=')): + out['laststep'] = arg.split('=', 1)[1] + elif arg.startswith('-'): + logger.warning( + "'%s' is not supported by the mg7 output and is ignored; " + "the equivalent settings live in Cards/run_card.toml", arg) + elif arg and not out['name']: + out['name'] = arg + return out + + def _set_run_name(self, name): + """Honour `launch -n NAME`: the mg7 run directory is + Events/_NN, so the name is a run_card setting.""" + path = os.path.join("Cards", "run_card.toml") + run_card = RunCardMG7(path, consistency=False) + if run_card["run"]["run_name"] != name: + run_card["run"]["run_name"] = name + run_card.write(path) + logger.info("run name set to '%s'", name) + + def _apply_laststep(self, switch, opts): + """`--laststep=parton` means: generate the events, run nothing on them. + + The mg7 tool selection is the switch dict the question returns, so + stopping at parton level is switching those entries off. + """ + if opts['laststep'] not in ('parton', 'auto', ''): + logger.warning( + "--laststep=%s is not supported by the mg7 output; " + "select the programs to run in the launch question instead", + opts['laststep']) + return switch + if opts['laststep'] != 'parton' or not switch: + return switch + switch = dict(switch) + for key in ('shower', 'detector', 'analysis', 'madspin', 'reweight'): + if not _off(switch.get(key, 'OFF')): + logger.info("--laststep=parton: not running %s", key) + switch[key] = 'OFF' + return switch + + +def build_selector_cmd(mother=None): + """Build the (monkey-patched) mother command + merged switch/card selector + used by the mg7 output. Returns the selector *class* and a mother instance + understood by AskRun/AskforEditCard. + + ``mother`` lets an in-process caller supply its own :class:`MG7Cmd` -- the + one MG5 registered as its child -- so that the question resolves scripted + answers against MG5's own command file (see :class:`MG7Cmd`). When it is + omitted a standalone instance is built, which is what ``bin/generate_events`` + gets. + """ + + from madgraph.interface.madevent_interface import AskRunEditCard + + from madgraph.various import banner as _banner_mod + from madgraph.various import misc as _misc + + class MG7Selector(AskRunEditCard): + """Merged switch/card question for the mg7 output. + + The mg7 run_card is a TOML file, so define_paths/init_run/do_set are + overridden *on this class* (not globally on AskforEditCard) to treat it + as such. Scoping them here is important: MadSpin/reweight spawn their own + legacy madevent runs in the same process, and a global patch would break + their (run_card.dat based) card editing.""" + + # param_card + the TOML run_card are always offered + always_cards = ['param_card.dat', 'run_card.toml'] + optional_cards = [] + + def define_paths(self, **opt): + super().define_paths(**opt) + self.paths["run"] = os.path.join(self.me_dir, "Cards", "run_card.toml") + self.paths["run_card.toml"] = os.path.join(self.me_dir, "Cards", "run_card.toml") + # the TOML run_card uses its own default file (concrete defaults + # written at output time); this powers "set default". + self.paths["run_default"] = os.path.join(self.me_dir, "Cards", "run_card_default.toml") + + def init_run(self, cards): + # Make sure the run_card is loaded as a RunCardMG7 (the generic + # editor does not recognise run_card.toml), else "set " fails. + out = super().init_run(cards) + if not isinstance(getattr(self, "run_card", None), RunCardMG7): + toml_path = self.paths.get("run") or os.path.join( + self.me_dir, "Cards", "run_card.toml") + if os.path.exists(toml_path): + try: + # allow_scan so a run_card holding scan:[...] values loads + with _misc.TMP_variable(_banner_mod.RunCard, "allow_scan", True): + self.run_card = RunCardMG7(toml_path, consistency="warning") + self.run_set = list(self.run_card.keys()) + except Exception as err: + logger.warning("could not load %s: %s", toml_path, err) + if isinstance(getattr(self, "run_card", None), RunCardMG7): + self.run_card.allow_scan = True + return getattr(self, "run_set", out) + + def do_set(self, line, *args, **kwargs): + # madevent-style shortcuts (lhc/lep/fixed_scale/no_parton_cut), cut + # editing, energy units and arithmetic/mass expressions on the TOML + # run_card, intercepted before delegating to the generic editor. + targs = self.split_arg(line) + run_card = getattr(self, "run_card", None) + if isinstance(run_card, RunCardMG7) and targs: + start = 1 if targs[0] == "run_card" else 0 + if len(targs) > start: + name = targs[start] + nlow = name.lower() + rest = " ".join(targs[start + 1:]).split("#")[0].strip() + masses = run_card.get_mass_shortcuts(getattr(self, "param_card", None)) + + if nlow in ("no_parton_cut", "nocut", "no_cut"): + run_card.remove_all_cut() + logger.info("removing all cuts from the run_card.toml") + self.modified_card.add("run") + return + # [histograms] is a list of observables, not a parameter, + # so "set histograms ..." can only mean the whole section. + if nlow in ("histograms", "histogram", "plots"): + value = rest.lower() + if value in ("off", "none", "no", "false", "0"): + run_card.remove_all_histograms() + logger.info( + "removing all histograms from the run_card.toml") + self.modified_card.add("run") + return + if value in ("default", "on"): + if run_card.restore_default_histograms( + self.paths.get("run_default")): + logger.info( + "restored the %d histograms written at " + "output time", + len(run_card["histograms"])) + self.modified_card.add("run") + else: + logger.warning( + "no run_card_default.toml to restore the " + "histograms from") + return + logger.warning( + "'set histograms %s': only OFF (remove them all) " + "and default (put the generated ones back) are " + "understood; edit the [histograms] section of " + "Cards/run_card.toml for anything else", + rest or "(no value)") + return + if nlow in ("lhc", "lep", "ilc", "lcc") and rest: + ecm = run_card.set_collider(nlow, rest, masses) + logger.info("set %s collider: e_cm = %s GeV", nlow, ecm) + self.modified_card.add("run") + return + if nlow == "fixed_scale" and rest: + val = run_card.set_fixed_scale(rest, masses) + logger.info("set fixed scales to %s GeV", val) + self.modified_card.add("run") + return + + if rest and run_card.is_cut_name(name): + cut, bound, val = run_card.set_cut(name, run_card.evaluate(rest, masses)) + logger.info("modify cut %s.%s of the run_card.toml to %s", cut, bound, val) + self.modified_card.add("run") + return + + # 'iseed' is the madevent spelling of the mg7 run.seed, + # but the two disagree on how to ask for a random seed: + # madevent uses 0, mg7 uses -1 (0 being a perfectly good + # fixed seed there). Translate here rather than in + # _LO_SCALAR_MAP, which also drives the LO -> MG7 run_card + # conversion and would silently pin those runs to seed 0. + if rest and nlow == 'iseed': + value = run_card.evaluate(rest, masses) + try: + value = int(value) + except (TypeError, ValueError): + value = None + if value is None: + logger.warning("ignoring 'set iseed %s': not an integer", rest) + return + seed = -1 if value == 0 else value + run_card.set('run.seed', seed, user=True) + logger.info("set iseed (mg7 run.seed) of the run_card.toml to %s%s", + seed, " (random)" if seed == -1 else "") + self.modified_card.add("run") + return + + # legacy madevent run_card names -> mg7 "section.key", so a + # madevent-style launch script ("set nevents 500", "set + # use_syst F", "set bwcutoff 10", ...) edits the mg7 + # run_card.toml verbatim. RunCardMG7._LO_SCALAR_MAP is the + # same rename table used by the LO->MG7 run_card conversion. + if rest and nlow in run_card._LO_SCALAR_MAP \ + and nlow not in [k.lower() for k in run_card.keys()]: + target = run_card._LO_SCALAR_MAP[nlow] + run_card.set(target, rest, user=True) + logger.info("set %s (mg7 %s) of the run_card.toml to %s", + nlow, target, run_card[target]) + self.modified_card.add("run") + return + + if rest and nlow in [k.lower() for k in run_card.keys()]: + current = run_card[nlow] + if isinstance(current, (int, float)) and not isinstance(current, bool): + resolved = run_card.evaluate(rest, masses) + if not isinstance(resolved, str): + prefix = "run_card " if start == 1 else "" + line = "%s%s %s" % (prefix, name, resolved) + return super().do_set(line, *args, **kwargs) + + return MG7Selector, (mother if mother is not None else MG7Cmd()) + + +def ask_edit_cards(mother=None) -> dict: + """Single (MadDM-style) question letting the user both pick which programs + to run after generation and edit the associated cards. Returns the switch + dict describing the selected tools. + + ``mother`` is the :class:`MG7Cmd` asking the question. Passing the instance + MG5 registered as its child is what lets a scripted ``launch`` answer this + question from MG5's own command file. + """ + + selector_class, mother = build_selector_cmd(mother) + # path_msg is what makes Cmd.check_answer_in_input_file accept a bare path as + # an answer (its "elif path:" branch), so that a scripted launch can hand the + # question a card/banner path -- as the question itself advertises -- and have + # it replace the corresponding card. Without it the path is rejected ("This + # answer is not valid for current question") and the default is used instead. + try: + switch, question = mother.ask('', '0', [], path_msg='enter path', + ask_class=selector_class, + mode='auto', line_args=[], force=False, + return_instance=True) + except BaseException: + # interrupted (e.g. ctrl-C) before the prune below could run: prune + # here too, using the instance ask() stashed pre-interrupt, or the + # eagerly-materialised tool cards get misread as "left ON" next time. + instance = getattr(mother, '_last_ask_instance', None) + if instance is not None and hasattr(instance, 'active_cards'): + try: + mother.keep_cards(instance.active_cards()) + except Exception as error: + logger.debug('could not prune tool cards after aborted launch question: %s', error) + raise + switch = dict(switch) + prune_unselected_tool_cards(mother, question, switch) + return switch + + +def prune_unselected_tool_cards(mother, question, switch) -> None: + """Hide the cards of the tools that were NOT selected in the question. + + The question materialises every candidate tool card (copying the + *_default.dat) so that it can offer them all for edition; madevent's + ask_run_configuration prunes them again afterwards (keep_cards on the + selected cards), but the mg7 launcher had no such cleanup. The materialised + cards therefore persisted, and a later re-launch of the same output saw them + all present and defaulted every switch to ON (set_default_ keys off + card presence). Prune here too, using the question's own switch->card map, so + a re-launch reflects the previous selection instead of turning everything on. + """ + try: + keep = list(question.always_cards) + for spec in question.switch_cards: + if spec['on'](switch): + keep.append(spec['card']) + mother.keep_cards(keep) + except Exception as error: + logger.debug('could not prune unselected tool cards: %s', error) + + +def _raise_open_file_limit() -> None: + """Remove the soft limit on the number of open files: it can be quite low + on some systems and the event-combination step opens one file per channel.""" + soft_lim, hard_lim = resource.getrlimit(resource.RLIMIT_NOFILE) + try: + resource.setrlimit(resource.RLIMIT_NOFILE, (hard_lim, hard_lim)) + except (ValueError, OSError) as error: + logger.debug("could not raise the open-file limit: %s", error) + + +def main() -> None: + """``bin/generate_events`` entry point. + + It runs the same :class:`MG7Cmd` command MG5's in-process ``launch`` uses, + so the two entry points cannot drift apart; the only difference is that + nothing here supplies a ``mother``/``inputfile``, so the question falls back + to reading this process's own stdin. + """ + parser = argparse.ArgumentParser() + parser.add_argument("-f", action="store_false", dest="ask_edit_cards") + args = parser.parse_args() + + MG7Cmd().run_cmd("generate_events" if args.ask_edit_cards else "generate_events -f") diff --git a/madgraph/iolibs/template_files/mg7/common.py b/madgraph/iolibs/template_files/mg7/common.py new file mode 100644 index 0000000000..253f4d06d9 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/common.py @@ -0,0 +1,200 @@ +"""Shared state of the mg7 launcher: the madspace import, logging and small helpers.""" + +import logging +import os +import re +import subprocess +import sys +import time + +# Install madspace on first use, then import it. The bootstrap lives in its own +# module (which must not import madspace) so that an in-process caller can run +# it *before* importing this one -- see bootstrap.ensure_madspace. +from madgraph.iolibs.template_files.mg7.bootstrap import ( + ensure_madspace as _ensure_madspace, + drop_install_path as _drop_install_path, + MADSPACE_DIR as _MADSPACE_DIR, +) + +_ensure_madspace() + +import madspace as ms + +# madspace is imported; stop its install directory from shadowing the caller's +# yaml/packaging/... for the rest of what is now MG5's own session. +_drop_install_path() + +_source_hash = subprocess.run( + [sys.executable, str(_MADSPACE_DIR / "source_hash.py")], + capture_output=True, text=True, check=True, +).stdout.strip() +if _source_hash != ms.SOURCE_HASH: + print() + print( + "\033[1m\033[31mWARNING\033[39m: madspace source and installed binaries " + "are not compatible (source hash mismatch) — consider recompiling " + "madspace (e.g. `install madspace -y`)\033[0m" + ) + print() + +logger = logging.getLogger("madgraph7") +LOG_LEVEL_MAP = { + ms.Logger.LogLevel.level_debug: logging.DEBUG, + ms.Logger.LogLevel.level_info: logging.INFO, + ms.Logger.LogLevel.level_warning: logging.WARNING, + ms.Logger.LogLevel.level_error: logging.ERROR, +} +def ms_log_handler(level: ms.Logger.LogLevel, message: str): + logger.log(LOG_LEVEL_MAP[level], message) +ms.Logger.set_log_handler(ms_log_handler) + + +def get_start_time(): + return time.time(), time.process_time() + + +def format_time(t: int, centi: bool = False): + hours, t = divmod(t, 3600) + minutes, seconds = divmod(t, 60) + if centi: + return f"{int(hours):02}:{int(minutes):02}:{seconds:02.2f}" + else: + return f"{int(hours):02}:{int(minutes):02}:{seconds:02.0f}" + + +def resolve_verbosity(verbosity: str) -> str: + """Resolve the run_card "auto" verbosity to "pretty"/"log" depending on + whether stdout is attached to a terminal; other values pass through + unchanged.""" + if verbosity == "auto": + return "pretty" if sys.stdout.isatty() else "log" + return verbosity + + +def resolve_seed(seed: int) -> int: + """Resolve the run_card "seed": -1 draws a fresh 64-bit seed via + os.urandom, any other value is used as-is.""" + if seed == -1: + return int.from_bytes(os.urandom(8), "big") + return seed + + +def device_type_of(device_name: str) -> str: + """The device type of one 'device' run_card entry (the optional ':' + suffix is stripped).""" + return device_name.split(":")[0] + + +def backend_of(device_name: str, cpu_mode: str) -> str: + """The build BACKEND that serves one 'device' run_card entry. + + ``cpu_mode`` describes the SIMD width of the CPU code and is therefore + meaningless for the cuda/hip devices: those build the backend named after + the device itself. + """ + device_type = device_type_of(device_name) + return cpu_mode if device_type == "cpu" else device_type + + +_PRECISION_2_FPTYPE = {"all64": "d", "all32": "f", "color32": "m", "denom64": "v"} + + +def fptype_of(precision: str) -> str: + return _PRECISION_2_FPTYPE.get(str(precision).lower(), str(precision).lower()) + + +def resolve_auto_backend(build_path: str) -> str: + """Ask the matrix-element Makefile to resolve the ``auto`` cpu_mode. + + Given the produced shared library will have its name taken from the resolved + backend name, we need to make sure the detection is taking place correctly + and catch any possible error. + """ + command = ["make", "-n", "BACKEND=auto", "detect-backend"] + try: + result = subprocess.run( + command, + cwd=build_path, + capture_output=True, + text=True, + check=True, + ) + except OSError as exc: + raise RuntimeError( + f"Could not run make to resolve cpu_mode='auto' in '{build_path}': {exc}" + ) from exc + except subprocess.CalledProcessError as exc: + output = "\n".join( + part.strip() for part in (exc.stdout, exc.stderr) if part.strip() + ) + detail = f"\nmake output:\n{output}" if output else "" + raise RuntimeError( + f"Could not resolve cpu_mode='auto' in '{build_path}': " + f"Exit status {exc.returncode}.{detail}" + ) from exc + + match = re.search( + r"^BACKEND=(\S+) \(was auto\)$", result.stdout, re.MULTILINE + ) + if match is None or match.group(1) == "auto": + output = "\n".join( + part.strip() for part in (result.stdout, result.stderr) if part.strip() + ) + detail = f"\nmake output:\n{output}" if output else "" + raise RuntimeError( + f"Could not resolve cpu_mode='auto' in '{build_path}': " + f"make failed to report a backend.{detail}" + ) + return match.group(1) + + +def me_parameters(run_card): + """Run-time parameters of the matrix-element libraries, passed to each + instance through umami_set_parameter by ms.Context.load_matrix_element: + the window of the $-excluded propagators is the run card bw_cutoff (the + one the phase space uses around the resonances too), and with + [generation] interference_helicity = "summed" the helicity sum as the + |M|^2 of an interference (see interference_helicities_summed). That one is + only sent when asked for: a library made before it existed refuses an + unknown parameter.""" + parameters = {"bwcutoff": float(run_card["phasespace"]["bw_cutoff"])} + if interference_helicities_summed(run_card): + parameters["interference_helicity_summed"] = 1. + return parameters + + +def interference_helicities_summed(run_card): + """Whether [generation] interference_helicity is "summed". + + For an interference (squared split orders dropping a component), a helicity + can contribute negatively. "exact", the default, draws the helicity of each + event on |T_i| and gives the event sign(T_i) * sum_j |T_j|, as madevent does: + the helicity of the event is exact, at the cost of a larger sum|w| (about 2x + the events for the same precision). "summed" keeps the helicity sum sum_j T_j + as the weight (lower variance), and the helicity then means nothing: it is + written as 9 in the LHE (lhe_helicities).""" + return run_card["generation"].get("interference_helicity", "exact") == "summed" + + +def lhe_helicities(meta, run_card): + """The helicity table the LHE completer writes for a subprocess: its own, + or 9 for every particle of an interference subprocess whose |M|^2 is the + helicity sum (interference_helicities_summed).""" + helicities = meta["helicities"] + if meta.get("interference") and interference_helicities_summed(run_card): + return [[9] * len(row) for row in helicities] + return helicities + + +def lhe_weight_info(status): + """(IDWTUP, XMAXUP) for the LHE block of a run with this generator + status. The events are unweighted to +-sigma_abs, keeping the sign of the + integrand, and XMAXUP is that unit weight rather than the signed cross + section. A sample that can hold negative weights -- an interference -- + declares it with IDWTUP = -4, as madevent does: the cross section is then + the mean of the signed event weights. With -3 a shower takes it as + |XSECUP| times the mean sign instead (Pythia8 does), which is not the + cross section, and +3 is for positive weights only. gridpack.py applies the + same rule inline, since it must not import this module.""" + negative = status.mean < 0 or status.mean_abs > abs(status.mean) * (1 + 1e-12) + return (-4 if negative else 3), status.mean_abs diff --git a/madgraph/iolibs/template_files/mg7/driver.py b/madgraph/iolibs/template_files/mg7/driver.py new file mode 100644 index 0000000000..ed1d974ce9 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/driver.py @@ -0,0 +1,186 @@ +"""Runs, scans and the entry point of the mg7 launcher.""" + +import gc +import logging +import os +import re +import shutil + +from madgraph.iolibs.template_files.mg7.common import logger +from madgraph.various.banner import RunCardMG7 +from madgraph.iolibs.template_files.mg7.model_widths import compute_auto_widths +from madgraph.iolibs.template_files.mg7.postprocess import run_lhe_postprocessing, run_selected_tools, selected_tools +from madgraph.iolibs.template_files.mg7.madgraph_process import MadgraphProcess + +def _release_channel_generators(process) -> None: + """Drop the per-channel event/weight generators so their intermediate + .npy files are deleted right away. madspace only frees them when the + owning Python objects are collected, and process.event_generator / + process.phasespaces sit in a reference cycle that plain refcounting + never breaks -- only an explicit gc.collect() does, promptly, instead of + leaving it to whenever the cyclic GC next runs on its own.""" + process.event_generator = None + process.phasespaces = None + gc.collect() + + +def run_single(switch=None) -> "MadgraphProcess": + """Run a single generation and return the process (for its result).""" + compute_auto_widths() + process = MadgraphProcess() + process.survey() + process.train_madnis() + process.generate_events() + if process.gridpack_run_mode != "regular": + return process + # run_card-driven LHE post-processing (displaced vertex + systematics) + run_lhe_postprocessing(process) + # run the post-processing tools (Pythia8/Delphes/MadSpin/reweight/analysis) + # selected in the merged question above on the generated events. + if switch: + run_selected_tools(switch, process) + return process + + +def detect_run_scan(run_card_path): + """Return a banner.RunCardIterator if the run_card contains scan:[...] + values, else None.""" + from madgraph.various import banner as banner_mod + from madgraph.various import misc as _misc + with _misc.TMP_variable(banner_mod.RunCard, 'allow_scan', True): + rc = banner_mod.RunCard(run_card_path, consistency=False) + if getattr(rc, 'scan_set', None): + return banner_mod.RunCardIterator(run_card_path) + return None + + +def detect_param_scan(param_card_path): + """Return a ParamCardIterator if the param_card contains scan:[...] values, + else None.""" + if not os.path.exists(param_card_path): + return None + from models import check_param_card as param_card_mod + it = param_card_mod.ParamCardIterator(param_card_path) + for block in it.order: + for param in block: + if isinstance(param.value, str) and param.value.strip().lower().startswith('scan'): + return it + return None + + +def run_scan(iterator, card_path, switch=None) -> None: + """Iterate over all scan points, running a full generation for each and + accumulating the results, then write the scan summary. Works for both the + run_card (RunCardIterator) and the param_card (ParamCardIterator); their + interface (__iter__/write/store_entry/get_next_name/write_summary) is the + same. The scan card is restored afterwards.""" + import tomllib + with open(os.path.join("Cards", "run_card.toml"), "rb") as f: + run_name = tomllib.load(f).get("run", {}).get("run_name", "run") + + from models import check_param_card as param_card_mod + is_param_scan = isinstance(iterator, param_card_mod.ParamCardIterator) + + backup = card_path + ".scan_bak" + shutil.copy(card_path, backup) + try: + for i, point in enumerate(iterator): + point.write(card_path) + logger.info("=== scan point %d ===", i + 1) + process = run_single(switch) + # use the run directory the process actually created, so the + # per-point params.dat written by write_summary has a home + name = os.path.basename(process.run_path) + if is_param_scan: + # pass the (possibly auto-width-updated) card so the summary + # records the model-computed width for each 'auto' entry + iterator.store_entry(name, process.get_result(), + param_card_path=card_path) + else: + iterator.store_entry(name, process.get_result()) + _release_channel_generators(process) + os.makedirs("Events", exist_ok=True) + summary = os.path.join("Events", "scan_%s.txt" % run_name) + iterator.write_summary(summary) + logger.info("scan results written to %s", summary) + finally: + shutil.move(backup, card_path) + + +def run_generation(switch=None) -> None: + """Run the generation, expanding a scan over the run_card or the param_card + when one is present (scanning both simultaneously is not allowed).""" + run_card_path = os.path.join("Cards", "run_card.toml") + param_card_path = os.path.join("Cards", "param_card.dat") + run_iter = detect_run_scan(run_card_path) + param_iter = detect_param_scan(param_card_path) + if run_iter and param_iter: + raise RuntimeError( + "Scanning simultaneously over the run_card and the param_card is " + "not allowed. Please keep the scan:[...] entries in only one card.") + if run_iter: + run_scan(run_iter, run_card_path, switch) + elif param_iter: + run_scan(param_iter, param_card_path, switch) + else: + _release_channel_generators(run_single(switch)) + + +def _lhe_needed_by(switch, card) -> "str | None": + """Why the run needs an LHE event file, or None when it does not: every + post-processing tool run_selected_tools drives and the run_card + [postprocessing] steps read the LHE file. ``card`` is the raw TOML content + of the run_card.""" + tools = selected_tools(switch) if switch else [] + if tools: + return tools[0] + cfg = card.get("postprocessing", {}) + try: + if float(cfg.get("time_of_flight", -1.0)) >= 0: + return "time_of_flight" + except (TypeError, ValueError): + pass + # the legacy systematics.py path only runs when madspace did not compute + # the weights itself (see run_lhe_postprocessing) + if cfg.get("systematics") \ + and not card.get("systematics", {}).get("enable", True): + return "systematics" + return None + + +def force_lhe_output_if_needed(switch) -> None: + """The events are written as npy by default; switch the run_card to the + LHE format when a selected post-processing needs an LHE file. The card is + rewritten rather than refused, so that asking for a shower in the launch + question is enough -- the user does not have to know that it constrains + output_format. + + The card is read as raw TOML and only its output_format line is edited: a + run_card scan ("scan:[...]" values) is not a valid RunCardMG7 and must + survive untouched for run_scan.""" + import tomllib + path = os.path.join("Cards", "run_card.toml") + if not os.path.exists(path): + return + with open(path, "rb") as f: + card = tomllib.load(f) + output_format = card.get("run", {}).get( + "output_format", RunCardMG7()["run"]["output_format"]) + if output_format == "lhe": + return + reason = _lhe_needed_by(switch, card) + if reason is None: + return + with open(path) as f: + text = f.read() + text, count = re.subn(r'(?m)^(\s*output_format\s*=\s*)("[^"]*"|\'[^\']*\'|\S+)', + r'\g<1>"lhe"', text, count=1) + if not count: + text, count = re.subn(r'(?m)^\[run\][^\n]*\n', '\\g<0>output_format = "lhe"\n', + text, count=1) + if not count: + text += '\n[run]\noutput_format = "lhe"\n' + with open(path, "w") as f: + f.write(text) + logging.getLogger("madevent").info( + "output_format set to 'lhe' (required by %s).", reason) diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index b871e59039..ceb4420a49 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -1,4685 +1,73 @@ -import argparse -import gc -import os -import sys -import time -from datetime import timedelta -from pathlib import Path -import glob -import shutil -import json -import subprocess -import re -import logging -from dataclasses import dataclass, field -from typing import Literal, NamedTuple -import resource - -# Install madspace on first use, then import it. The bootstrap lives in its own -# module (which must not import madspace) so that an in-process caller can run -# it *before* importing this one -- see bootstrap.ensure_madspace. -from madgraph.iolibs.template_files.mg7.bootstrap import ( - ensure_madspace as _ensure_madspace, - drop_install_path as _drop_install_path, - MG_ROOT as _MG_ROOT, - MADSPACE_DIR as _MADSPACE_DIR, - INSTALL_DIR as _INSTALL_DIR, +"""Entry point of the mg7 launcher, re-exporting its modules.""" + +from madgraph.iolibs.template_files.mg7.common import * # noqa: F401,F403 +from madgraph.iolibs.template_files.mg7.common import ms, logger # noqa: F401 +from madgraph.iolibs.template_files.mg7.observables import ( # noqa: F401 + CutItem, + HistItem, + _MERGED_PID_REPRESENTATIVE, + clean_pids, + _SIGNED_OBSERVABLES, + _flavor_options, + observable_pids, + build_observable, + decay_products, + pid_is_qcd, ) - -_ensure_madspace() - -import madspace as ms - -# madspace is imported; stop its install directory from shadowing the caller's -# yaml/packaging/... for the rest of what is now MG5's own session. -_drop_install_path() -from models.check_param_card import ParamCard -from madgraph.iolibs.template_files.mg7 import hwu_output -from madgraph.iolibs.template_files.mg7 import plots -from madgraph.iolibs.template_files.mg7 import systematics_summary -from madgraph.iolibs.template_files.mg7 import npy_to_lhe -from madgraph.iolibs.template_files.mg7.custom_function_loader import ( - load_custom_function, - load_surrogate, +from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import ( # noqa: F401 + Channel, + PhaseSpace, + MultiChannelData, + build_topologies, + build_multi_channel_data, + MadgraphSubprocess, +) +from madgraph.iolibs.template_files.mg7.madgraph_process import ( # noqa: F401 + MadgraphProcess, +) +from madgraph.iolibs.template_files.mg7.options import ( # noqa: F401 + _ROOTED_OPTIONS, + load_mg5_options, + lhapdf_paths, +) +from madgraph.iolibs.template_files.mg7.model_widths import ( # noqa: F401 + _model_hash_warned, + _model_cache, + read_stored_model, + _proc_characteristic, + load_process_model, + compute_auto_widths, +) +from madgraph.iolibs.template_files.mg7.postprocess import ( # noqa: F401 + _find_event_file, + _find_or_convert_event_file, + _report_failure, + _off, + _has_effective_handler, + _TOOL_LOGGING_READY, + _setup_logging, + selected_tools, + run_selected_tools, + _add_time_of_flight, + _lhapdf_config_path, + _run_systematics, + run_lhe_postprocessing, +) +from madgraph.iolibs.template_files.mg7.driver import ( # noqa: F401 + _release_channel_generators, + run_single, + detect_run_scan, + detect_param_scan, + run_scan, + run_generation, + _lhe_needed_by, + force_lhe_output_if_needed, +) +from madgraph.iolibs.template_files.mg7.cli import ( # noqa: F401 + MG7Cmd, + build_selector_cmd, + ask_edit_cards, + prune_unselected_tool_cards, + _raise_open_file_limit, + main, ) -from madgraph.various.banner import RunCardMG7 -from madgraph.various import misc -from madgraph.interface.extended_cmd import Cmd - -_source_hash = subprocess.run( - [sys.executable, str(_MADSPACE_DIR / "source_hash.py")], - capture_output=True, text=True, check=True, -).stdout.strip() -if _source_hash != ms.SOURCE_HASH: - print() - print( - "\033[1m\033[31mWARNING\033[39m: madspace source and installed binaries " - "are not compatible (source hash mismatch) — consider recompiling " - "madspace (e.g. `install madspace -y`)\033[0m" - ) - print() - -logger = logging.getLogger("madgraph7") -LOG_LEVEL_MAP = { - ms.Logger.LogLevel.level_debug: logging.DEBUG, - ms.Logger.LogLevel.level_info: logging.INFO, - ms.Logger.LogLevel.level_warning: logging.WARNING, - ms.Logger.LogLevel.level_error: logging.ERROR, -} -def ms_log_handler(level: ms.Logger.LogLevel, message: str): - logger.log(LOG_LEVEL_MAP[level], message) -ms.Logger.set_log_handler(ms_log_handler) - - -def get_start_time(): - return time.time(), time.process_time() - - -def format_time(t: int, centi: bool = False): - hours, t = divmod(t, 3600) - minutes, seconds = divmod(t, 60) - if centi: - return f"{int(hours):02}:{int(minutes):02}:{seconds:02.2f}" - else: - return f"{int(hours):02}:{int(minutes):02}:{seconds:02.0f}" - - -def resolve_verbosity(verbosity: str) -> str: - """Resolve the run_card "auto" verbosity to "pretty"/"log" depending on - whether stdout is attached to a terminal; other values pass through - unchanged.""" - if verbosity == "auto": - return "pretty" if sys.stdout.isatty() else "log" - return verbosity - - -def resolve_seed(seed: int) -> int: - """Resolve the run_card "seed": -1 draws a fresh 64-bit seed via - os.urandom, any other value is used as-is.""" - if seed == -1: - return int.from_bytes(os.urandom(8), "big") - return seed - - -def device_type_of(device_name: str) -> str: - """The device type of one 'device' run_card entry (the optional ':' - suffix is stripped).""" - return device_name.split(":")[0] - - -def backend_of(device_name: str, cpu_mode: str) -> str: - """The build BACKEND that serves one 'device' run_card entry. - - ``cpu_mode`` describes the SIMD width of the CPU code and is therefore - meaningless for the cuda/hip devices: those build the backend named after - the device itself. - """ - device_type = device_type_of(device_name) - return cpu_mode if device_type == "cpu" else device_type - - -_PRECISION_2_FPTYPE = {"all64": "d", "all32": "f", "color32": "m", "denom64": "v"} - - -def fptype_of(precision: str) -> str: - return _PRECISION_2_FPTYPE.get(str(precision).lower(), str(precision).lower()) - - -def resolve_auto_backend(build_path: str) -> str: - """Ask the matrix-element Makefile to resolve the ``auto`` cpu_mode. - - Given the produced shared library will have its name taken from the resolved - backend name, we need to make sure the detection is taking place correctly - and catch any possible error. - """ - command = ["make", "-n", "BACKEND=auto", "detect-backend"] - try: - result = subprocess.run( - command, - cwd=build_path, - capture_output=True, - text=True, - check=True, - ) - except OSError as exc: - raise RuntimeError( - f"Could not run make to resolve cpu_mode='auto' in '{build_path}': {exc}" - ) from exc - except subprocess.CalledProcessError as exc: - output = "\n".join( - part.strip() for part in (exc.stdout, exc.stderr) if part.strip() - ) - detail = f"\nmake output:\n{output}" if output else "" - raise RuntimeError( - f"Could not resolve cpu_mode='auto' in '{build_path}': " - f"Exit status {exc.returncode}.{detail}" - ) from exc - - match = re.search( - r"^BACKEND=(\S+) \(was auto\)$", result.stdout, re.MULTILINE - ) - if match is None or match.group(1) == "auto": - output = "\n".join( - part.strip() for part in (result.stdout, result.stderr) if part.strip() - ) - detail = f"\nmake output:\n{output}" if output else "" - raise RuntimeError( - f"Could not resolve cpu_mode='auto' in '{build_path}': " - f"make failed to report a backend.{detail}" - ) - return match.group(1) - - -def me_parameters(run_card): - """Run-time parameters of the matrix-element libraries, passed to each - instance through umami_set_parameter by ms.Context.load_matrix_element: - the window of the $-excluded propagators is the run card bw_cutoff (the - one the phase space uses around the resonances too), and with - [generation] interference_helicity = "summed" the helicity sum as the - |M|^2 of an interference (see interference_helicities_summed). That one is - only sent when asked for: a library made before it existed refuses an - unknown parameter.""" - parameters = {"bwcutoff": float(run_card["phasespace"]["bw_cutoff"])} - if interference_helicities_summed(run_card): - parameters["interference_helicity_summed"] = 1. - return parameters - - -def interference_helicities_summed(run_card): - """Whether [generation] interference_helicity is "summed". - - For an interference (squared split orders dropping a component), a helicity - can contribute negatively. "exact", the default, draws the helicity of each - event on |T_i| and gives the event sign(T_i) * sum_j |T_j|, as madevent does: - the helicity of the event is exact, at the cost of a larger sum|w| (about 2x - the events for the same precision). "summed" keeps the helicity sum sum_j T_j - as the weight (lower variance), and the helicity then means nothing: it is - written as 9 in the LHE (lhe_helicities).""" - return run_card["generation"].get("interference_helicity", "exact") == "summed" - - -def lhe_helicities(meta, run_card): - """The helicity table the LHE completer writes for a subprocess: its own, - or 9 for every particle of an interference subprocess whose |M|^2 is the - helicity sum (interference_helicities_summed).""" - helicities = meta["helicities"] - if meta.get("interference") and interference_helicities_summed(run_card): - return [[9] * len(row) for row in helicities] - return helicities - - -def lhe_weight_info(status): - """(IDWTUP, XMAXUP) for the LHE block of a run with this generator - status. The events are unweighted to +-sigma_abs, keeping the sign of the - integrand, and XMAXUP is that unit weight rather than the signed cross - section. A sample that can hold negative weights -- an interference -- - declares it with IDWTUP = -4, as madevent does: the cross section is then - the mean of the signed event weights. With -3 a shower takes it as - |XSECUP| times the mean sign instead (Pythia8 does), which is not the - cross section, and +3 is for positive weights only. gridpack.py applies the - same rule inline, since it must not import this module.""" - negative = status.mean < 0 or status.mean_abs > abs(status.mean) * (1 + 1e-12) - return (-4 if negative else 3), status.mean_abs - - -@dataclass -class Channel: - phasespace_mapping: ms.PhaseSpaceMapping - adaptive_mapping: ms.Flow | ms.VegasMapping - discrete_sym: ms.DiscreteSampler | ms.DiscreteFlow | None - discrete_flavor: ms.DiscreteSampler | ms.DiscreteFlow | None - channel_weight_indices: list[int] | None - name: str - active_flavors: list[int] - event_generator: ms.ChannelEventGenerator | None = None - - -@dataclass -class PhaseSpace: - mode: Literal["multichannel", "flat", "both"] - channels: list[Channel] - symfact: list[int | None] - first_chan_weight_remap: list[list[int]] = field(default_factory=list) - first_remapped_chan_count: int = 0 - second_chan_weight_remap: list[int] = field(default_factory=list) - second_remapped_chan_count: int = 0 - prop_chan_weights: ms.PropagatorChannelWeights | None = None - subchan_weights: ms.SubchannelWeights | None = None - cwnet: ms.ChannelWeightNetwork | None = None - - -class MultiChannelData(NamedTuple): - amp2_remaps: list[list[int]] - symfact: list[int | None] - topologies: list[list[ms.Topology]] - permutations: list[list[list[int]]] - channel_indices: list[list[int]] - channel_weight_indices: list[list[list[int]]] - diagram_indices: list[list[int]] - diagram_color_indices: list[list[list[int]]] - diagram_propagator_pdgs: list[list[list[int]]] - active_flavors: list[list[list[int]]] - qcd_s_channel_count: list[int] - - -@dataclass -class CutItem: - observable_kwargs: dict - min: float - max: float - mode: str - - -@dataclass -class HistItem: - observable_kwargs: dict - min: float - max: float - bin_count: int - # the distribution of the event weight itself (the reserved "weight" key - # of [histograms]); it is read off the event record, not computed from - # the momenta, so it has no observable - name: str = "" - from_weight: bool = False - - -class MadgraphProcess: - def __init__(self): - self.madnis_trained = False - self.lhe_completer = None - self.load_cards() - self.init_backend() - self.init_event_dir() - self.init_context() - self.init_custom_functions() - self.init_surrogate() - self.init_cuts() - self.init_histograms() - self.init_generator_config() - self.init_beam() - self.init_subprocesses() - - def load_cards(self) -> None: - self.run_card = RunCardMG7(os.path.join("Cards", "run_card.toml")) - # Resolved once so every generator built during this run shares the same - # seed; the concrete value (even if randomly drawn) is recorded in each - # generator's info.json. - self.run_seed = resolve_seed(self.run_card["run"]["seed"]) - self.param_card_path = os.path.join("Cards", "param_card.dat") - self.param_card = ParamCard(self.param_card_path) - with open(os.path.join("SubProcesses", "subprocesses.json")) as f: - self.subprocess_data = json.load(f) - if self.run_card["phasespace"]["merge_subprocesses"]: - with open(os.path.join("SubProcesses", "merged_subprocesses.json")) as f: - self.merged_subprocess_data = json.load(f) - else: - self.merged_subprocess_data = None - - self.init_decay_mode() - self.init_me_frame() - if interference_helicities_summed(self.run_card): - count = sum(1 for meta in self.subprocess_data if meta.get("interference")) - if count: - logger.info( - "interference_helicity = summed: the weight of the %d interference " - "subprocess(es) is the helicity sum, and their helicities are " - "written as 9 in the LHE", count) - - def init_me_frame(self) -> None: - """Resolve the run card's me_frame into the list of external particles - whose momentum sum defines the frame the matrix element is evaluated in. - - Empty (the default) means no boost: the matrix element sees the momenta - in the frame they are generated in. Unlike madevent, which hands its - matrix element momenta that are already in the partonic centre of mass, - madspace generates them in the lab frame, so asking for madevent's own - frame ([1, 2]) does mean an actual boost here. It changes nothing for a - Lorentz invariant matrix element (only the rounding), but it is what - defines the polarisation axes of a polarised one. - - The particle numbers are checked here rather than only in madspace, - which would not see them until the integrands are built -- long after - the matrix-element libraries have been compiled. Every subprocess is - checked against its own particles: an output can mix multiplicities - (p p > z j and p p > z j j, say, for MLM merging), and madspace builds - one matrix element per subprocess, each refusing a particle it does not - have. That is stricter than madevent, which ignores the particles a - process does not have and so changes the frame from one multiplicity - to the next without a word. - """ - me_frame = list(self.run_card["run"]["me_frame"]) - pdgs = [clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) - for meta in self.subprocess_data] - counts = sorted({len(pdg_list) for pdg_list in pdgs}) - for meta, pdg_list in zip(self.subprocess_data, pdgs): - for index in me_frame: - if 1 <= index <= len(pdg_list): - continue - process = "%s > %s" % ( - " ".join(str(pid) for pid in meta["incoming"]), - " ".join(str(pid) for pid in meta["outgoing"]), - ) - mixed = ( - f"; this output mixes subprocesses with " - f"{', '.join(str(count) for count in counts)} external " - "particles, and me_frame has to name particles every one " - "of them has" - if len(counts) > 1 else "" - ) - raise ValueError( - f"me_frame particle {index} out of range: subprocess " - f"{process} has {len(pdg_list)} external particles{mixed}" - ) - if len(me_frame) == 1: - # A massless particle has no rest frame: its momentum squares to - # zero, so the boost divides by it and every momentum comes back as - # a NaN, silently. Name the mistake instead. - massless = [pdg_list[me_frame[0] - 1] for pdg_list in pdgs - if self.get_mass(pdg_list[me_frame[0] - 1]) == 0.] - if massless: - raise ValueError( - f"me_frame = {me_frame} asks for the rest frame of particle " - f"{me_frame[0]}, which is massless (pdg " - f"{sorted(set(massless))}) and has none. Name a massive " - "particle, or several particles whose sum is massive." - ) - self.me_frame = me_frame - self.incoming_count = 1 if self.is_decay else 2 - - def init_decay_mode(self) -> None: - """Decide whether this directory is a decay (1 -> n) or a collision. - - Read off the exported process rather than the run card: the two cannot - then disagree. MadSpin generates its decay matrix elements this way. - """ - incoming_counts = { - len(clean_pids(meta["incoming"])) for meta in self.subprocess_data - } - if incoming_counts - {1, 2}: - raise ValueError( - f"processes with {sorted(incoming_counts)} incoming particles " - "are not supported" - ) - if len(incoming_counts) > 1: - raise ValueError( - "cannot mix decays and collisions in one output directory" - ) - self.is_decay = incoming_counts == {1} - if not self.is_decay: - self.decaying_mass = None - return - - masses = { - self.get_mass(clean_pids(meta["incoming"])[0]) - for meta in self.subprocess_data - } - if len(masses) > 1: - raise ValueError( - f"decaying particles have different masses: {sorted(masses)}" - ) - self.decaying_mass = masses.pop() - if self.decaying_mass <= 0.0: - raise ValueError("the decaying particle must have a non-zero mass") - self.drop_closed_channels() - - def drop_closed_channels(self) -> None: - """Remove subprocesses the decaying particle is too light to produce. - - A multiparticle decay definition enumerates every vertex the model - allows, closed ones included: "t > b w+, w+ > all all" yields - t > b t b~ (and b W+ Z, b W+ h, ...), which need more mass than the top - has. Their partial width is exactly zero, but the phase-space mapping - has no physical point to hand back -- the invariant's lower bound ends - up above its upper bound -- so it produces NaN momenta and poisons the - whole integral. Drop them here instead. - """ - kept, dropped = [], [] - for meta in self.subprocess_data: - total = sum( - self.get_mass(pid) for pid in clean_pids(meta["outgoing"]) - ) - if total < self.decaying_mass: - kept.append(meta) - else: - dropped.append((meta["outgoing"], total)) - if dropped and self.merged_subprocess_data is not None: - # merged_subprocesses.json indexes the *unfiltered* subprocess - # list, so dropping entries here would silently shift every index - # it holds. Refuse rather than mis-map. - raise ValueError( - "merge_subprocesses is not supported for a decay directory " - "with kinematically closed channels" - ) - if dropped: - for outgoing, total in dropped: - logger.info( - "skipping closed decay channel -> %s (needs %.4g GeV, " - "the decaying particle has %.4g GeV)", - outgoing, total, self.decaying_mass, - ) - if not kept: - raise ValueError( - "every decay channel is kinematically closed: the decaying " - f"particle's mass is {self.decaying_mass} GeV" - ) - self.subprocess_data = kept - - def init_backend(self) -> None: - ms.set_simd_vector_size(self.run_card["run"]["simd_vector_size"]) - - def init_event_dir(self) -> None: - output_dir = self.run_card["gridpack"]["output_dir"] - temp_dir = self.run_card["gridpack"]["temp_output_dir"] - if output_dir: - self.run_path = os.path.abspath(os.path.expanduser(output_dir)) - os.makedirs(self.run_path, exist_ok=True) - else: - run_name = self.run_card["run"]["run_name"] - os.makedirs("Events", exist_ok=True) - run_dir_prefix = os.path.join("Events", f"{run_name}_") - existing_run_dirs = glob.glob(f"{run_dir_prefix}*") - run_index = 1 - for run_dir in existing_run_dirs: - run_index_str = run_dir[len(run_dir_prefix):] - if run_index_str.isnumeric(): - run_index = max(run_index, int(run_index_str) + 1) - while True: - try: - self.run_path = f"{run_dir_prefix}{run_index:02d}" - os.mkdir(self.run_path) - break - except FileExistsError: - run_index += 1 - # temporary npy files of the channel generators - if temp_dir: - self.temp_path = os.path.abspath(os.path.expanduser(temp_dir)) - os.makedirs(self.temp_path, exist_ok=True) - else: - self.temp_path = self.run_path - # Absolute on purpose. StatusFile is a C++ object that finishes its - # write from its destructor (rename info.json.tmp -> info.json), and - # that destructor runs whenever Python gets round to collecting the - # process object. When the launcher runs in MG5's process, an - # interrupted run is collected *after* the launch command has restored - # the working directory, and a relative path would then resolve - # somewhere else -- the rename throws a C++ filesystem_error out of a - # destructor, which is an immediate abort. - self.status_file = ms.StatusFile( - os.path.abspath(os.path.join(self.run_path, "info.json"))) - - def init_context(self) -> None: - device_names = self.run_card["run"]["device"] - self.contexts = [] - self.device_types = [] - self.devices = [] - self.pool_sizes = [] - for device_name in device_names: - if ":" in device_name: - device_type, device_index_str = device_name.split(":") - device_index = int(device_index_str) - else: - device_type = device_name - device_index = 0 - self.device_types.append(device_type) - if device_type == "cuda": - device = ms.cuda_device(device_index) - pool_size = self.run_card["run"]["gpu_thread_pool_size"] - elif device_type == "hip": - device = ms.hip_device(device_index) - pool_size = self.run_card["run"]["gpu_thread_pool_size"] - else: - device = ms.cpu_device() - pool_size = self.run_card["run"]["cpu_thread_pool_size"] - self.devices.append(device) - self.pool_sizes.append(pool_size) - self.contexts.append(ms.Context(device=device, thread_count=pool_size)) - - CUSTOM_FUNCTION_KINDS = ("cut", "scale", "pdf", "bias") - - def init_custom_functions(self) -> None: - """Load the functions of [custom_functions], building the libraries - first. They are loaded before the matrix elements, so that they have the - same index in every context and in a gridpack made from this run.""" - self.custom_functions = {} - self.custom_function_paths = [] - for kind in self.CUSTOM_FUNCTION_KINDS: - path = self.run_card["custom_functions"][kind] - if not path: - continue - build_dir = f"custom_{kind}" - if path.endswith(".so") and os.path.isfile(os.path.join(build_dir, "Makefile")): - self.build_custom_function(build_dir) - if not os.path.isfile(path): - raise RuntimeError( - f"[custom_functions] {kind} = {path!r}: file not found. " - "Create a .so library or a .py file with bin/init_custom_func." - ) - apis = load_custom_function(ms, self.contexts, self.device_types, kind, path) - if len({api.index() for api in apis}) != 1: - raise RuntimeError(f"custom {kind} function loaded at different indices") - self.custom_functions[kind] = apis[0] - self.custom_function_paths.append([kind, path]) - logger.info("Using the custom %s function %s", kind, path) - - def init_surrogate(self) -> None: - """Check the [surrogate] settings. The surrogates are loaded per - subprocess, by MadgraphSubprocess.""" - settings = self.run_card["surrogate"] - self.surrogate = None - if not settings["function"]: - return - path = settings["function"] - if path.endswith(".py") and not os.path.isfile(path): - raise RuntimeError(f"[surrogate] function = {path!r}: file not found") - if self.run_card["phasespace"]["sde_strategy"] != "denominators": - raise RuntimeError( - "[surrogate] needs [phasespace] sde_strategy = \"denominators\", " - "as the surrogate does not return channel weights" - ) - if self.run_card["run"]["dummy_matrix_element"]: - raise RuntimeError("[surrogate] cannot be used with dummy_matrix_element") - if self.run_card["gridpack"]["save_gridpack"]: - raise RuntimeError("[surrogate] does not support gridpacks yet") - self.surrogate = dict(settings) - if not self.surrogate["param_card"]: - self.surrogate["param_card"] = self.param_card_path - logger.info( - "Using the surrogate %s, %s with the full matrix element", - path, - "unweighting" if settings["mode"] == "unweight" else "reweighting", - ) - - def weighting_groups(self, subprocess_channels: list[list[int]]) -> list: - """One ms.WeightingGroup per subprocess with a [surrogate]: the second - pass with the full matrix element over its channels, given by their - indices in the event generator.""" - if self.surrogate is None: - return [] - gen_args = self.run_card["generation"] - return [ - ms.WeightingGroup( - channel_indices, - [ms.WeightingStep( - subproc.step_matrix_elements, - mode=self.surrogate["mode"], - max_overweight_truncation=gen_args["max_overweight_truncation"], - freeze_max_weight_after=gen_args["freeze_max_weight_after"], - )], - ) - for subproc, channel_indices in zip(self.subprocesses, subprocess_channels) - if channel_indices - ] - - @staticmethod - def build_custom_function(build_dir: str) -> None: - """Run make in a custom_ directory, logging to compile.log.""" - log_path = os.path.join(build_dir, "compile.log") - with open(log_path, "w") as log: - returncode = subprocess.call( - ["make"], cwd=build_dir, stdout=log, stderr=subprocess.STDOUT - ) - if returncode != 0: - raise RuntimeError( - f"Compilation of {build_dir} failed with exit status {returncode}; " - f"see {log_path} for details" - ) - - def custom_function(self, kind: str, particle_count: int, pid_options: list): - """The ms.CustomFunction of [custom_functions] , or None.""" - api = self.custom_functions.get(kind) - if api is None: - return None - function_type = getattr(ms.CustomFunction, kind) - return ms.CustomFunction(api, function_type, particle_count, pid_options) - - def parse_observable(self, name: str, order_observable: str) -> dict: - parts = name.split("-") - sum_momenta = False - sum_observable = False - ordered = False - multiparticles = self.run_card["multiparticles"] - - if len(parts) == 0: - raise ValueError("Invalid observable name") - elif len(parts) == 1: - # event-level observables - obs_name = parts[0] - select_pids = [] - else: - if parts[-1] == "sum": - sum_observable = True - obs_name = parts[-2] - selection = parts[:-2] - elif parts[-2] == "sum": - sum_momenta = True - obs_name = parts[-1] - selection = parts[:-2] - else: - obs_name = parts[-1] - selection = parts[:-1] - select_pids = [] - order_indices = [] - for mp_name in selection: - mp_parts = mp_name.split("_") - if mp_parts[-1].isnumeric(): - order_indices.append(int(mp_parts[-1])) - select_pids.append(multiparticles["_".join(mp_parts[:-1])]) - ordered = True - else: - order_indices.append(0) - select_pids.append(multiparticles[mp_name]) - - return dict( - observable=obs_name, - select_pids=select_pids, - sum_momenta=sum_momenta, - sum_observable=sum_observable, - order_observable=order_observable if ordered else None, - order_indices=order_indices if ordered else [], - ignore_incoming=True, - name=name, - ) - - def init_cuts(self) -> None: - inf = float("inf") - order_observable = self.run_card["cuts"].get("order_by", "pt") - self.cut_data = [ - CutItem( - observable_kwargs=self.parse_observable(key, order_observable), - min=values.get("min", -inf), - max=values.get("max", inf), - mode=values.get("mode", "all"), - ) - for key, values in self.run_card["cuts"].items() - if key != "order_by" - ] - - # [histograms] key that means "the distribution of the event weight", - # normalised to the unit weight sigma_abs, rather than an observable of the - # momenta: an unweighted sample is a spike at 1 (and at -1 for its negative - # weights) and a partially unweighted one shows its spread. - weight_histogram_key = "weight" - - def init_histograms(self) -> None: - order_observable = self.run_card["histograms"].get("order_by", "pt") - self.hist_data = [] - for key, values in self.run_card["histograms"].items(): - if key == "order_by": - continue - from_weight = key == self.weight_histogram_key - self.hist_data.append(HistItem( - observable_kwargs=( - None if from_weight - else self.parse_observable(key, order_observable)), - min=values["min"], - max=values["max"], - bin_count=values["bin_count"], - name=key, - from_weight=from_weight, - )) - - def ensure_pdf_set(self, pdf_set: str) -> misc.LhapdfPaths: - """Make sure the requested LHAPDF set is available, downloading it into - the first writable PDF directory if needed, and return the (possibly - updated) LHAPDF resolution. Nothing is downloaded when the set is - already present.""" - paths = lhapdf_paths() - if paths.find_set(pdf_set): - return paths - if not paths.config: - logger.debug("no usable lhapdf-config: cannot download %s", pdf_set) - return paths - if not paths.download_path: - logger.debug("no writable PDF directory: cannot download %s", pdf_set) - return paths - try: - from madgraph.interface.common_run_interface import CommonRunCmd - os.makedirs(paths.download_path, exist_ok=True) - logger.info("PDF set %s not found; downloading into %s", - pdf_set, paths.download_path) - CommonRunCmd.install_lhapdf_pdfset_static( - paths.config, paths.download_path, pdf_set) - except Exception as err: - logger.warning("Could not download PDF set %s: %s", pdf_set, err) - return paths - global _LHAPDF - _LHAPDF = paths.with_data_path(paths.download_path) - return _LHAPDF - - def _no_pdf_message(self, pdf_set: str) -> str: - """Explain a missing PDF set in terms of the settings that fix it.""" - searched = os.pathsep.join(self.lhapdf.data_paths) or \ - "(no PDF set directory could be located)" - if self.lhapdf.config: - found = "using lhapdf-config: %s" % self.lhapdf.config - else: - found = ("no usable lhapdf-config was found, so the set could not " - "be downloaded automatically") - return ( - "LHAPDF set %r, requested by [beam] pdf in Cards/run_card.toml, " - "was not found.\n" - " searched: %s\n" - " %s\n" - "Fix this in one of the following ways:\n" - " * install LHAPDF for this MadGraph: MG7> install lhapdf6\n" - " * point MadGraph at an existing one:\n" - " MG7> set lhapdf /path/to/lhapdf-config\n" - " MG7> save options\n" - " * or set it for this directory only, by adding\n" - " lhapdf = /path/to/lhapdf-config\n" - " to %s\n" - " * or set $LHAPDF_DATA_PATH to a directory containing %r." - % (pdf_set, searched, found, - os.path.join("Cards", "me5_configuration.txt"), pdf_set)) - - def init_beam(self) -> None: - beam_args = self.run_card["beam"] - - # built lazily by build_systematics (needs the subprocess data) - self.systematics = None - self.systematics_data = None - self.systematics_context = None - self.event_histograms = None - self.event_histograms_context = None - - if self.is_decay: - # No beams: the total energy is the decaying particle's mass, and - # "leptonic" is what the mappings call "no parton luminosity". - self.e_cm = self.decaying_mass - self.leptonic = True - else: - self.e_cm = beam_args["e_cm"] - self.leptonic = beam_args["leptonic"] - - dynamical_scales = { - "transverse_energy": ms.EnergyScale.transverse_energy, - "transverse_mass": ms.EnergyScale.transverse_mass, - "half_transverse_mass": ms.EnergyScale.half_transverse_mass, - "partonic_energy": ms.EnergyScale.partonic_energy, - } - if beam_args["dynamical_scale_choice"] in dynamical_scales: - dynamical_scale_type = dynamical_scales[beam_args["dynamical_scale_choice"]] - else: - raise ValueError("Unknown dynamical scale choice") - self.scale_kwargs = dict( - dynamical_scale_type=dynamical_scale_type, - ren_scale_fixed=beam_args["fixed_ren_scale"], - fact_scale_fixed=beam_args["fixed_fact_scale"], - ren_scale=beam_args["ren_scale"], - fact_scale1=beam_args["fact_scale1"], - fact_scale2=beam_args["fact_scale2"], - scale_factor=float(beam_args.get("scale_factor", 1.)), - ) - if self.is_decay and self.custom_functions.pop("scale", None) is not None: - logger.warning("custom scale ignored: a decay uses the decaying mass") - if self.is_decay and self.custom_functions.pop("pdf", None) is not None: - logger.warning("custom pdf ignored: a decay has no parton densities") - if self.is_decay: - # One scale is available for a decay -- the decaying mass -- so use - # it, fixed, whatever the card asks for. - self.scale_kwargs.update( - ren_scale_fixed=True, - fact_scale_fixed=True, - ren_scale=self.decaying_mass, - fact_scale1=self.decaying_mass, - fact_scale2=self.decaying_mass, - scale_factor=1., - ) - self.pdf_grid = None - self.lhapdf = None - self.pdf_dir = None - self.fixed_alphas_info = self.write_fixed_alphas_info() - self.alphas_grid = ms.AlphaSGrid(self.fixed_alphas_info) - for context in self.contexts: - self.alphas_grid.initialize_globals(context) - self.running_coupling = ms.RunningCoupling(self.alphas_grid) - return - - pdf_set = beam_args["pdf"] - self.pdf_set = pdf_set - self.fixed_alphas_info = None - self.lhapdf = self.ensure_pdf_set(pdf_set) - self.pdf_dir = self.lhapdf.find_set(pdf_set) - if self.pdf_dir is None: - raise RuntimeError(self._no_pdf_message(pdf_set)) - self.alphas_grid = ms.AlphaSGrid(os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}.info")) - if "pdf" in self.custom_functions: - # alpha_s still comes from the PDF set of the run card - self.pdf_grid = ms.CustomPdf(self.custom_functions["pdf"]) - else: - self.pdf_grid = ms.PdfGrid( - os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}_0000.dat") - ) - for context in self.contexts: - if not isinstance(self.pdf_grid, ms.CustomPdf): - self.pdf_grid.initialize_globals(context) - self.alphas_grid.initialize_globals(context) - self.running_coupling = ms.RunningCoupling(self.alphas_grid) - - # ------------------------------------------------------------------ - # scale / PDF systematics - # ------------------------------------------------------------------ - def systematics_enabled(self) -> bool: - """The native scale/PDF weights are requested ([systematics] enable, or - the legacy [generation] systematics alias).""" - try: - enabled = bool(self.run_card["systematics"]["enable"]) - except Exception: - enabled = False - if not enabled: - try: - enabled = bool(self.run_card["generation"]["systematics"]) - except Exception: - enabled = False - if enabled and "pdf" in getattr(self, "custom_functions", {}): - # the variations re-evaluate the PDF set, which a custom PDF replaces - if not getattr(self, "_warned_custom_pdf_systematics", False): - logger.warning("systematics are disabled with a custom pdf") - self._warned_custom_pdf_systematics = True - return False - return enabled - - @staticmethod - def pdf_set_info(info_path: str) -> dict: - """Read SetIndex, NumMembers, ErrorType and SetDesc from an LHAPDF - .info file (missing keys get neutral defaults).""" - info = {"SetIndex": 0, "NumMembers": 1, "ErrorType": "", "SetDesc": ""} - try: - with open(info_path) as f: - for line in f: - key, sep, value = line.partition(":") - key = key.strip() - if not sep or key not in info: - continue - value = value.strip().strip('"').strip("'") - if key in ("SetIndex", "NumMembers"): - info[key] = int(value) - else: - info[key] = value - except OSError as err: - logger.warning("could not read %s: %s", info_path, err) - return info - - def resolve_pdf_set_name(self, entry: str) -> tuple[str, int | None]: - """Turn a [systematics] pdf entry ("", "", "@", - "@") into (set name, member or None) using the LHAPDF - pdfsets.index when an id is given.""" - member = None - if "@" in entry: - entry, member_str = entry.split("@", 1) - member = int(member_str) - entry = entry.strip() - if not entry.isdigit(): - return entry, member - lhaid = int(entry) - index_path = os.path.join(self.pdf_dir or "", "pdfsets.index") - sets = [] - try: - with open(index_path) as f: - for line in f: - parts = line.split() - if len(parts) >= 2 and parts[0].isdigit(): - sets.append((int(parts[0]), parts[1])) - except OSError: - raise ValueError( - "cannot resolve LHAPDF id %s: no %s" % (lhaid, index_path)) - for set_id, name in sets: - if set_id == lhaid: - return name, member - # an id inside a set's member range - sets.sort() - for (set_id, name), (next_id, _) in zip(sets, sets[1:] + [(None, None)]): - if set_id < lhaid and (next_id is None or lhaid < next_id): - return name, lhaid - set_id if member is None else member - raise ValueError("cannot resolve LHAPDF id %s" % lhaid) - - def resolve_pdf_variations(self) -> list: - """Expand the [systematics] pdf entries into ms.PdfMemberSpec objects - (downloading the sets when needed). The nominal member itself is never - listed: the calculator adds it as the central weight of its group.""" - entries = self.run_card["systematics"]["pdf"] - if isinstance(entries, str): - entries = [entries] - specs = [] - seen = set() - for raw in entries: - entry = str(raw).strip() - if not entry or entry.lower() in ("none", "no", "false"): - continue - if entry.lower() == "central": - continue - if entry.lower() == "errorset": - set_name, member = self.pdf_set, None - else: - set_name, member = self.resolve_pdf_set_name(entry) - variation_lhapdf = self.ensure_pdf_set(set_name) - variation_dir = variation_lhapdf.find_set(set_name) - if variation_dir is None: - logger.warning("PDF set %s not available; its variation is skipped", set_name) - continue - info_path = os.path.join(variation_dir, set_name, f"{set_name}.info") - if not os.path.exists(info_path): - logger.warning("PDF set %s not available; its variation is skipped", set_name) - continue - info = self.pdf_set_info(info_path) - is_nominal = set_name == self.pdf_set - if member is not None: - members = [member] - else: - members = list(range(info["NumMembers"])) - if is_nominal: - members = members[1:] - if not members: - logger.warning( - "PDF set %s has a single member: no PDF uncertainty can be " - "computed from it. Use a set with error members in " - "[systematics] pdf (e.g. a PDF4LHC or NNPDF error set).", - set_name) - for m in members: - if is_nominal and m == 0: - continue - if (set_name, m) in seen: - continue - seen.add((set_name, m)) - grid_file = os.path.join(variation_dir, set_name, f"{set_name}_{m:04d}.dat") - if not os.path.exists(grid_file): - logger.warning("PDF member file %s not found; skipped", grid_file) - continue - specs.append(ms.PdfMemberSpec( - set_name=set_name, set_lhaid=info["SetIndex"], member=m, - grid_file=grid_file, info_file=info_path, - error_type=info["ErrorType"], description=info["SetDesc"], - )) - return specs - - def build_systematics_args(self) -> list: - """One ms.SubprocessSystArgs per (unmerged) subprocess: the alpha_s power - of |M|^2 and the beam parton ids of every flavor index, indexed like the - LHECompleter (subprocess_index / flavor_index of the combined events).""" - args = [] - for meta in self.subprocess_data: - beam_pdgs = [ - [int(flavor["options"][0][0]), int(flavor["options"][0][1])] - for flavor in meta["flavors"] - ] - args.append(ms.SubprocessSystArgs( - qcd_power=int(meta.get("qcd_power", -1)), beam_pdgs=beam_pdgs)) - return args - - def build_systematics(self): - """Create the ms.SystematicsCalculator from the [systematics] section, or - None when disabled. Also keeps the JSON description needed by gridpacks.""" - if not self.systematics_enabled(): - self.systematics = None - return None - syst = self.run_card["systematics"] - config = ms.SystematicsConfig() - config.mur = [float(v) for v in syst["mur"]] - config.muf = [float(v) for v in syst["muf"]] - config.together = bool(syst["together"]) - config.dyn_scales = self.resolve_dynamical_scales() - # alternative dynamical scales get the same factor as the nominal one - config.scale_factor = float(self.run_card["beam"].get("scale_factor", 1.)) - config.write_inputs = bool(syst["write_inputs"]) - config.has_pdf = not self.leptonic - # Without parton luminosity -- a decay, or leptonic beams -- there is no - # nominal PDF *set* to describe and no PDF variation to compute, so - # has_pdf gates the nominal-set metadata and the member list below. - # The alpha_s grid is a separate matter: the mu_R variations always need - # one, so its .info file has to be recorded whatever has_pdf says. - # init_beam() skips the PDF lookup for a decay only -- leptonic beams - # still resolve a set and take alpha_s from it -- so the path follows - # self.pdf_dir, and falls back to the constant grid written for a decay. - if self.pdf_dir is not None: - info_path = os.path.join(self.pdf_dir, self.pdf_set, f"{self.pdf_set}.info") - else: - info_path = self.fixed_alphas_info - if config.has_pdf: - info = self.pdf_set_info(info_path) - config.nominal_set_name = self.pdf_set - config.nominal_lhaid = info["SetIndex"] - config.nominal_error_type = info["ErrorType"] - config.nominal_description = info["SetDesc"] - config.pdf_members = self.resolve_pdf_variations() - else: - config.pdf_members = [] - args = self.build_systematics_args() - # the PDFs, alpha_s and (for mixed-order subprocesses) the matrix - # elements are evaluated with the batched madspace functions on this - # CPU context - self.systematics_context = ms.Context(device=ms.cpu_device(), thread_count=1) - matrix_elements, flavor_remap, me_backend = self.build_systematics_matrix_elements( - args, self.systematics_context) - self.systematics = ms.SystematicsCalculator( - config, args, self.pdf_grid, self.alphas_grid, - self.systematics_context, matrix_elements, flavor_remap) - for warning in self.systematics.warnings: - logger.warning("systematics: %s", warning) - self.systematics_data = { - "config": json.loads(config.to_json()), - "subproc_args": [json.loads(a.to_json()) for a in args], - "nominal_grid_file": os.path.join( - self.pdf_dir, self.pdf_set, f"{self.pdf_set}_0000.dat") - if config.has_pdf else None, - "nominal_info_file": info_path, - # matrix elements re-evaluated for the mixed-order subprocesses - "me_paths": [meta["me_path"] for meta in self.subprocess_data], - "me_backend": me_backend, - "flavor_remap": flavor_remap, - "me_frame": self.me_frame, - "incoming_count": self.incoming_count, - } - return self.systematics - - _DYN_SCALE_CODES = { - "transverse_energy": 1, "transverse_mass": 2, - "half_transverse_mass": 3, "partonic_energy": 4, - } - - def resolve_dynamical_scales(self) -> list: - """[systematics] dynamical_scale entries (names or systematics.py codes - 1-4) as codes, without the choice the events were generated with.""" - try: - entries = self.run_card["systematics"]["dynamical_scale"] - except Exception: - return [] - if isinstance(entries, str): - entries = [entries] - beam = self.run_card["beam"] - generated = None - if not (beam["fixed_ren_scale"] and beam["fixed_fact_scale"]): - generated = self._DYN_SCALE_CODES.get(beam["dynamical_scale_choice"]) - codes = [] - for raw in entries: - entry = str(raw).strip().lower() - if not entry: - continue - if entry.isdigit(): - code = int(entry) - elif entry in self._DYN_SCALE_CODES: - code = self._DYN_SCALE_CODES[entry] - else: - raise ValueError("unknown dynamical scale choice %r in [systematics] " - "dynamical_scale" % raw) - if code not in (1, 2, 3, 4): - raise ValueError("dynamical scale code %s must be 1-4" % code) - if code == generated or code in codes: - continue - codes.append(code) - return codes - - def build_systematics_matrix_elements(self, args, context): - """For the subprocesses whose |M|^2 mixes several powers of alpha_s - (qcd_power = -1), load the matrix element into the systematics CPU - context so that the renormalisation scale variations can re-evaluate it - at the varied alpha_s. Returns ([MatrixElement or None per subprocess], - flavor remap, backend name); empty when nothing is needed or no CPU - library exists.""" - need = [i for i, a in enumerate(args) if a.qcd_power < 0] - empty = ([], [], None) - if not need or self.run_card["run"]["dummy_matrix_element"]: - return empty - cpu_backends = [b for b in self.backends if not b.startswith(("cuda", "hip"))] - if not cpu_backends: - logger.warning("systematics: no CPU matrix-element library available, the " - "renormalisation scale variations of the mixed-order " - "subprocesses cannot be computed") - return empty - backend = cpu_backends[0] - matrix_elements = [] - flavor_remap = [] - for i, meta in enumerate(self.subprocess_data): - flavor_remap.append([int(f["index"]) for f in meta["flavors"]]) - if i not in need: - matrix_elements.append(None) - continue - api = context.load_matrix_element( - meta["me_path"].format(device=backend), self.param_card_path, - me_parameters(self.run_card)) - matrix_elements.append(ms.MatrixElement( - api, - [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, - ms.MatrixElement.flavor_in], - [ms.MatrixElement.matrix_element_out], - False, - self.me_frame, - self.incoming_count, - )) - logger.info("systematics: re-evaluating the matrix element of %d mixed-order " - "subprocess(es) for the renormalisation scale variations", len(need)) - return matrix_elements, flavor_remap, backend - - def build_event_histograms(self): - """ms.EventHistograms for the [histograms] observables, filled with the - written events and all their variation weights (info.json - "event_histograms"); None without histograms.""" - if not self.hist_data \ - or not self.run_card["run"]["postprocessing_histograms"]: - self.event_histograms = None - return None - context = ms.Context(device=ms.cpu_device(), thread_count=1) - specs = [ - ms.EventHistogramSpec( - name=item.name, min=item.min, max=item.max, - bin_count=item.bin_count, from_weight=item.from_weight) - for item in self.hist_data - ] - from_momenta = [item for item in self.hist_data if not item.from_weight] - observables = [] - for meta in self.subprocess_data: - if not from_momenta: - # nothing to evaluate on the momenta (only the weight is - # histogrammed): no observable runtime to build - observables.append(None) - continue - all_pids = clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) - values = ms.ObservableValues([ - build_observable( - observable_pids(meta, item.observable_kwargs), - len(meta["incoming"]), item.observable_kwargs - ) - for item in from_momenta - ]) - observables.append(ms.SubprocessObservables(values, len(all_pids))) - # the weight histograms are drawn in units of the unit weight sigma_abs - # (the cross section of a sample without negative weights) - self.event_histograms_context = context - self.event_histograms = ms.EventHistograms( - context, specs, observables, - reference_weight=self._mean_event_weight()) - return self.event_histograms - - def _mean_event_weight(self) -> float: - """The unit weight of the unweighted events, sigma_abs, used as the - unit of the weight histograms: the cross section of a positive sample, - and for one with negative weights (an interference) what puts them at - +-1 rather than dividing by a signed, possibly small, total. 0 (the raw - weight) when it is not available.""" - try: - return float(self.event_generator.status().mean_abs) - except Exception as error: - logger.debug("no cross section for the weight histograms: %s", error) - return 0. - - def log_systematics_summary(self) -> None: - """Print the scale/PDF uncertainties on the total cross section (the - per-variation cross sections are in events.weights.json / info.json).""" - if self.systematics is None or self.systematics.weight_count == 0: - return - summary = json.loads(self.systematics.summary()) - # the percentages come from systematics_summary, which get_result() - # reads too: the box and the scan row must not disagree, so neither - # computes them itself - xsec = systematics_summary.nominal_cross_section(summary) - if xsec is None: - return - def format_variation(up, down): - # right-justify the signed numbers (not the sign alone) so the % - # values line up across rows, with extra spacing between the two - return f"{'+%.3g' % up:>6}% {'-%.3g' % down:>6}%" - - scale_count = sum(1 for v in self.systematics.variations if v.is_scale) - member_count = len(self.systematics.members) - rows = [("Variations per event:", - f"{self.systematics.weight_count} ({scale_count} scale, {member_count} PDF members)")] - for pdf in summary.get("pdf", []): - rows.append(("PDF set:", f"{pdf['pdf_set']}, {pdf['error_type']}")) - rows.append(("Original cross-section:", f"{xsec} pb")) - scale = systematics_summary.scale_percentages(summary) - if scale is not None: - rows.append(("Scale variation:", format_variation(*scale))) - for _entry, up, down in systematics_summary.pdf_percentages(summary): - rows.append(("PDF variation:", format_variation(up, down))) - if self.event_generator_config.verbosity == ms.Verbosity.pretty: - box = ms.PrettyBox("Systematics", len(rows), [24, 0]) - box.set_column(0, [label for label, _ in rows]) - box.set_column(1, [value for _, value in rows]) - box.print_first() - else: - for label, value in rows: - logger.info("systematics, %s %s", label, value) - - def write_fixed_alphas_info(self) -> str: - """Write a minimal LHAPDF ``.info`` holding a constant alpha_s. - - A decay has no beams, so there is no PDF set to take alpha_s from -- - and demanding one (possibly downloading it) just to evaluate a coupling - would be absurd. The renormalisation scale of a decay is fixed at the - decaying mass anyway, so a constant alpha_s is the right answer, not an - approximation: take it from the param card, exactly as the Fortran - decay matrix elements do. - """ - # SMINPUTS entry 3 is alpha_s(m_Z), the same value the Fortran - # parameter setup feeds to G. - alpha_s = float(self.param_card.get_value("sminputs", 3)) - path = os.path.join(self.run_path, "fixed_alphas.info") - # AlphaSGrid interpolates in log(q^2) across the grid, so give it a - # comfortable number of nodes rather than the bare minimum of three. - q_values = [10 ** (i / 4.0) for i in range(-4, 21)] - with open(path, "w") as f: - f.write("SetDesc: constant alpha_s for a decay (no beams)\n") - f.write("AlphaS_Qs: [%s]\n" % ", ".join(f"{q:g}" for q in q_values)) - f.write( - "AlphaS_Vals: [%s]\n" - % ", ".join(f"{alpha_s:g}" for _ in q_values) - ) - logger.info( - "decay mode: fixed alpha_s = %g at mu = %g GeV", alpha_s, self.e_cm - ) - return path - - def needs_systematics_matrix_elements(self) -> bool: - """The systematics will re-evaluate a matrix element per scale point. - - True when the native weights are on and some subprocess mixes alpha_s - powers (qcd_power < 0), i.e. when build_systematics_matrix_elements() - loads a matrix element. Conservative: it does not check that a CPU - backend exists, so it can be True where that call ends up loading - nothing. - """ - if not self.systematics_enabled(): - return False - if self.run_card["run"]["dummy_matrix_element"]: - return False - return any(int(meta.get("qcd_power", -1)) < 0 - for meta in self.subprocess_data) - - def init_generator_config(self) -> None: - run_args = self.run_card["run"] - gen_args = self.run_card["generation"] - vegas_args = self.run_card["vegas"] - cfg = ms.GeneratorConfig() - cfg.target_count = gen_args["events"] - cfg.vegas_damping = vegas_args["damping"] - cfg.max_overweight_truncation = gen_args["max_overweight_truncation"] - cfg.freeze_max_weight_after = gen_args["freeze_max_weight_after"] - if self.surrogate is not None and self.surrogate["mode"] == "unweight": - # [generation] applies to the second pass, which gives the events - cfg.max_overweight_truncation = ( - self.surrogate["first_pass_max_overweight_truncation"] - ) - cfg.freeze_max_weight_after = ( - self.surrogate["first_pass_freeze_max_weight_after"] - ) - cfg.start_batch_size = vegas_args["start_batch_size"] - cfg.max_batch_size = vegas_args["max_batch_size"] - cfg.survey_min_iters = gen_args["survey_min_iters"] - cfg.survey_max_iters = gen_args["survey_max_iters"] - cfg.survey_target_precision = gen_args["survey_target_precision"] - cfg.optimization_patience = vegas_args["optimization_patience"] - cfg.optimization_threshold = vegas_args["optimization_threshold"] - cfg.cpu_batch_size = gen_args["cpu_batch_size"] - cfg.gpu_batch_size = gen_args["gpu_batch_size"] - cfg.verbosity = resolve_verbosity(run_args["verbosity"]) - cfg.combine_thread_count = run_args["combine_thread_pool_size"] - if self.needs_systematics_matrix_elements() and cfg.combine_thread_count != 1: - # SystematicsCalculator::compute() runs inside the combine workers, - # and a matrix element's per-thread process instances live in a - # ThreadResource indexed by ThreadPool::thread_index() -- the - # *calling* thread's index. A parallel combine pool would therefore - # need one process instance per worker; building N of them to index - # them safely is not worth it, the more so as - # SystematicsCalculator::matrix_elements() already serialises every - # evaluation on its own mutex, so that parallelism was not being - # used anyway. Run the combine stage single-threaded instead. - logger.info( - "systematics: the combine stage runs single-threaded because a " - "subprocess needs its matrix element re-evaluated for the " - "renormalisation scale variations") - cfg.combine_thread_count = 1 - cfg.cut_efficiency_threshold = gen_args["cut_efficiency_threshold"] - cfg.max_cut_repetitions = gen_args["max_cut_repetitions"] - self.event_generator_config = cfg - self.event_generator = None - - def init_subprocesses(self) -> None: - self.backends = self.compile_matrix_elements() - self.subprocesses = [] - if self.merged_subprocess_data is None: - for subproc_id, meta in enumerate(self.subprocess_data): - self.subprocesses.append(MadgraphSubprocess(self, meta, subproc_id)) - else: - for subproc_id, meta in enumerate(self.merged_subprocess_data): - self.subprocesses.append( - MadgraphSubprocess(self, meta, subproc_id, self.subprocess_data) - ) - - def compile_matrix_elements(self) -> list[str]: - """Build the matrix-element library of every subprocess, and return the - build backend of each requested device (in the same order). - - A 'cpu' device builds the backend named by the 'cpu_mode' entry; the - cuda/hip devices ignore cpu_mode and build their own backend. A - cpu_mode of 'auto' is resolved to the actual SIMD backend chosen on - this machine. - - SubProcesses/makefile is a dispatcher over the P* directories, so a - single 'make -j N' there builds all the subprocesses at once with one - shared pool of N jobs: no subprocess is built with N jobs while the - others wait, and none is limited to N/#subprocesses jobs either. - """ - device_names = self.run_card["run"]["device"] - if not isinstance(device_names, list): - device_names = [device_names] - cpu_mode = self.run_card["run"]["cpu_mode"] - backends = [backend_of(name, cpu_mode) for name in device_names] - if not self.subprocess_data: - return backends - - first_proc_path = self.subprocess_data[0]["path"] - subproc_path = os.path.dirname(first_proc_path) - - fptype = fptype_of(self.run_card["run"]["precision"]) - - # Resolve cpu_mode='auto' once (the build rules pick the best SIMD - # backend available here), so all subprocesses agree on the library names. - auto_backend = None - if "auto" in backends: - auto_backend = resolve_auto_backend(first_proc_path) - logger.info("cpu_mode 'auto' resolved as '%s'", auto_backend) - resolved = [ - auto_backend if backend == "auto" else backend - for backend in backends - ] - - nb_core = self.run_card["run"]["cpu_thread_pool_size"] - if not nb_core or nb_core < 0: - nb_core = os.cpu_count() or 1 - - log_path = os.path.join(self.run_path, "compile_subprocesses.log") - for backend in resolved: - missing = [ - meta for meta in self.subprocess_data - if not os.path.isfile(meta["me_path"].format(device=backend)) - ] - if not missing: - continue - logger.info( - f"Start compilation of SubProcesses for device '{backend}' " - f"({len(missing)} subprocess(es), {nb_core} parallel job(s)), " - f"see log detail in {log_path}" - ) - start_time = time.time() - self.make_subprocesses( - subproc_path, - [f"BACKEND={backend}", f"FPTYPE={fptype}", "USEBUILDDIR=1"], - nb_core, log_path, - ) - logger.info( - f"Compilation of SubProcesses done in {time.time() - start_time:.1f} s" - ) - return resolved - - @staticmethod - def make_subprocesses( - subproc_path: str, args: list[str], nb_core: int, log_path: str - ) -> None: - """Run 'make -j nb_core' in SubProcesses/, appending the full build - output to log_path (one run per device, so the file is appended to).""" - command = ["make", f"-j{nb_core}"] + args - with open(log_path, "a") as log: - log.write(f"\n$ cd {subproc_path} && {' '.join(command)}\n") - log.flush() - returncode = subprocess.call( - command, cwd=subproc_path, stdout=log, stderr=subprocess.STDOUT - ) - if returncode != 0: - raise RuntimeError( - f"Compilation of the SubProcesses failed with exit status " - f"{returncode}; see {log_path} for details" - ) - - def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenerator: - channel_generators = [] - subprocess_channels = [[] for _ in self.subprocesses] - for i, (subproc, phasespace) in enumerate(zip(self.subprocesses, phasespaces)): - for integrand, channel in zip( - subproc.build_integrands(phasespace), phasespace.channels - ): - if channel.event_generator is None: - channel.event_generator = ms.ChannelEventGenerator( - contexts=self.contexts, - integrand=integrand, - event_file=os.path.join(self.temp_path, f"events.{i}.{channel.name}.npy"), - weight_file=os.path.join(self.temp_path, f"weights.{i}.{channel.name}.npy"), - config=self.event_generator_config, - subprocess_index=i, - name=f"{i}.{channel.name}", - histograms=subproc.histograms - ) - subprocess_channels[i].append(len(channel_generators)) - channel_generators.append(channel.event_generator) - - event_generator = ms.EventGenerator( - contexts=self.contexts, - channels=channel_generators, - status_file=self.status_file, - config=self.event_generator_config, - seed=self.run_seed, - groups=self.weighting_groups(subprocess_channels), - ) - unused_globals = ( - set(self.contexts[0].global_names()) - event_generator.used_globals() - ) - for context in self.contexts: - for global_name in unused_globals: - context.delete_global(global_name) - return event_generator - - def survey_phasespaces( - self, phasespaces: list[PhaseSpace | None], survey_pass: int = 0 - ) -> ms.EventGenerator | None: - ps_filtered = [ps for ps in phasespaces if ps is not None] - if len(ps_filtered) == 0: - return None - event_generator = self.build_event_generator(ps_filtered) - event_generator.survey(survey_pass) - return event_generator - - def survey(self) -> None: - # survey_pass distinguishes the survey() calls below: "both" mode can - # re-survey a channel carried over unchanged from the multichannel pass - # into the final (simplified) pass, and both passes schedule jobs on the - # same underlying ChannelEventGenerator. The explicit pass index keeps each - # pass's job seeds independent of the other passes' job counts, rather than - # depending on call history. - phasespace_mode = self.run_card["phasespace"]["mode"] - if self.is_decay and phasespace_mode != "multichannel": - # The flat mapping is built from a synthetic two-incoming diagram, - # which a decay has no counterpart for -- and 'both'/'auto' end up - # simplifying towards it. Decays have few channels, so the - # multichannel phase space is the right one regardless. - logger.info( - "decay mode: using the multichannel phase space instead of " - "'%s'", phasespace_mode - ) - phasespace_mode = "multichannel" - if phasespace_mode in ["multichannel", "both", "auto"]: - self.phasespaces = [ - subproc.build_multichannel_phasespace() - for subproc in self.subprocesses - ] - self.event_generator = self.survey_phasespaces(self.phasespaces, 0) - elif phasespace_mode == "flat": - self.phasespaces = [ - subproc.build_flat_phasespace() - for subproc in self.subprocesses - ] - self.event_generator = self.survey_phasespaces(self.phasespaces, 0) - else: - raise ValueError("Unknown phasespace mode") - - channel_status = self.event_generator.channel_status() - chan_offset = 0 - madnis_enabled = False - for subproc, ps in zip(self.subprocesses, self.phasespaces): - mean = 0. - variance = 0. - count_opt = 0 - for status in channel_status[chan_offset:chan_offset + len(ps.channels)]: - mean += status.mean_abs - variance += status.error_abs**2 - count_opt += status.count_opt - # A subprocess whose matrix element is identically zero (an FCNC - # channel with every Wilson coefficient at zero, say) integrates to - # exactly zero, so there is no relative spread to speak of. Report 0 - # rather than dividing by it: that asks for the smallest networks and - # leaves MadNIS off, which is what an empty subprocess wants. - rsd = (variance * count_opt)**0.5 / mean if mean else 0. - subproc.set_madnis_auto_settings(rsd) - chan_offset += len(ps.channels) - if subproc.madnis_settings["enable"]: - madnis_enabled = True - if not ( - phasespace_mode == "both" or (phasespace_mode == "auto" and madnis_enabled) - ): - return - - phasespaces_multi = self.phasespaces - cross_sections = [] - index = 0 - for phasespace in phasespaces_multi: - channel_count = len(phasespace.channels) - # rank the channels by the integral of |w|: a channel whose signed - # integral cancels (an interference) can still carry most of the - # variance, and folding it into the flat channel would blow it up - cross_sections.append([ - status.mean_abs - for status in channel_status[index:index + channel_count] - ]) - index += channel_count - - self.phasespaces = [ - subproc.simplify_phasespace(ps_multi, cross_secs) - for subproc, ps_multi, cross_secs in zip( - self.subprocesses, phasespaces_multi, cross_sections - ) - ] - if any( - ps_multi is not ps_both - for ps_multi, ps_both in zip(phasespaces_multi, self.phasespaces) - ): - # distinct survey_pass: a channel carried over unchanged from the - # multichannel pass (pass 0) into this resurvey must not share its - # seed stream with that earlier pass. - self.event_generator = self.survey_phasespaces(self.phasespaces, 1) - - def train_madnis(self) -> None: - madnis_args = self.run_card["madnis"] - if not any(subproc.madnis_settings["enable"] for subproc in self.subprocesses): - return - - gen_args = self.run_card["generation"] - run_args = self.run_card["run"] - - verbosity = resolve_verbosity(run_args["verbosity"]) - madnis_phasespaces = [] - training_args = [] - self.event_generator = None - for subproc, phasespace in zip(self.subprocesses, self.phasespaces): - for channel in phasespace.channels: - channel.event_generator = None - - config = ms.MadnisConfig() - config.learning_rate = subproc.madnis_settings["lr"] - config.batches = subproc.madnis_settings["train_batches"] - config.log_interval = madnis_args["log_interval"] - config.integration_history_length = madnis_args["integration_history_length"] - config.channel_dropping_interval = madnis_args["channel_dropping_interval"] - config.channel_dropping_threshold = madnis_args["channel_dropping_threshold"] - config.cpu_generator_batch_size = gen_args["cpu_batch_size"] - config.gpu_generator_batch_size = gen_args["gpu_batch_size"] - config.gpu_generator_batch_granularity = madnis_args["gpu_generator_batch_granularity"] - config.generator_target_size_factor = madnis_args["generator_target_size_factor"] - config.batch_size_offset = madnis_args["batch_size_offset"] - config.batch_size_per_channel = subproc.madnis_settings["batch_size_per_channel"] - config.uniform_channel_ratio = madnis_args["uniform_channel_ratio"] - config.lr_schedule = madnis_args["lr_scheduler"] - config.adam_beta1 = madnis_args["adam_beta1"] - config.adam_beta2 = madnis_args["adam_beta2"] - config.adam_eps = madnis_args["adam_eps"] - config.adam_weight_decay = madnis_args["adam_weight_decay"] - config.grad_clip_threshold = madnis_args["grad_clip_threshold"] - config.buffer_capacity = madnis_args["buffer_capacity"] - config.minimum_buffer_size = madnis_args["minimum_buffer_size"] - config.buffered_steps_fraction = madnis_args["buffered_steps_fraction"] - config.buffer_skip_batches = madnis_args["buffer_skip_batches"] - config.buffer_unweighting_quantile = madnis_args["buffer_unweighting_quantile"] - config.fixed_cwnet_fraction = subproc.madnis_settings["fixed_cwnet_fraction"] - config.softclip_threshold = madnis_args["softclip_threshold"] - config.compressed_channel_weight_count = madnis_args["compressed_channel_weight_count"] - phasespace = subproc.build_madnis(phasespace) - madnis_phasespaces.append(phasespace) - training_args.append( - ms.TrainingArgs( - config=config, - integrands=subproc.build_integrands( - phasespace, - madnis_training=True, - drop_cuts_and_rescale=madnis_args["drop_zero_integrands"] - ), - cwnet=phasespace.cwnet, - ) - ) - - gen_context = self.contexts[0] - opt_context = ms.Context( - device=self.devices[0], thread_count=self.pool_sizes[0] - ) - opt_context.copy_globals_from(gen_context) - - madnis_training = ms.MultiMadnisTraining( - generator_context=gen_context, - optimizer_context=opt_context, - training_args=training_args, - verbosity=verbosity, - status_file=self.status_file, - # Reuses the run's resolved seed (also used by build_event_generator()). - # Only the single-channel CPU sample-generation path is currently seeded - # -- buffered training and GPU multi-channel batches are still - # non-deterministic. - seed=self.run_seed, - ) - madnis_training.train() - for phasespace, active_channels in zip( - madnis_phasespaces, madnis_training.active_channels() - ): - phasespace.channels = [ - phasespace.channels[index] for index in active_channels - ] - del madnis_training - del opt_context - self.phasespaces = madnis_phasespaces - for context in self.contexts[1:]: - context.copy_globals_from(self.contexts[0]) - self.event_generator = self.build_event_generator(madnis_phasespaces) - self.madnis_trained = True - - def update_madnis_status_single( - self, batch: int, batch_target: int, loss: float, lr: float, channel_count: int - ) -> None: - now = time.time() - if batch + 1 < batch_target: - if now - self.last_update_time < 0.1: - return - self.last_update_time = now - progress_bar = ms.format_progress((batch + 1) / batch_target, 52) - time_diff = now - self.madnis_wall_time - time_str = f"{format_time(time_diff)}" - else: - progress_bar = "" - wall_diff = now - self.madnis_wall_time - cpu_diff = time.process_time() - self.madnis_cpu_time - time_str = ( - f"{format_time(wall_diff, centi=True)} wall, " - f"{format_time(cpu_diff, centi=True)} cpu" - ) - batch_str = f"{batch + 1} / {batch_target}" - self.madnis_box.set_column(1, [ - f"{batch_str:<15} {progress_bar}", - f"{loss:>.4f}", - f"{channel_count}", - time_str - ]) - self.madnis_box.print_update() - - def update_madnis_status_multi( - self, - subproc_id: int, - batch: int, - batch_target: int, - loss: float, - lr: float, - channel_count: int - ) -> None: - now = time.time() - subproc_count = len(self.subprocesses) - if batch + 1 < batch_target: - if now - self.last_update_time < 0.1: - return - self.last_update_time = now - progress_bar = ms.format_progress((batch + 1) / batch_target, 34) - progress_bar_all = ms.format_progress( - (subproc_id * batch_target + batch + 1) / (subproc_count * batch_target), - 52 - ) - time_str = f"{format_time(now - self.madnis_wall_time)}" - subproc_str = f"{subproc_id} / {subproc_count}" - elif subproc_id < subproc_count - 1: - progress_bar = "" - progress_bar_all = ms.format_progress( - ((subproc_id + 1) * batch_target + 1) / (subproc_count * batch_target), - 52 - ) - time_str = f"{format_time(now - self.madnis_wall_time)}" - subproc_str = f"{subproc_id} / {subproc_count}" - else: - progress_bar = "" - progress_bar_all = "" - wall_diff = now - self.madnis_wall_time - cpu_diff = time.process_time() - self.madnis_cpu_time - time_str = ( - f"{format_time(wall_diff, centi=True)} wall, " - f"{format_time(cpu_diff, centi=True)} cpu" - ) - subproc_str = f"{subproc_count} / {subproc_count}" - batch_str = f"{batch + 1} / {batch_target}" - self.madnis_upper_box.set_column(1, [ - f"{subproc_str:<15} {progress_bar_all}", - time_str, - ]) - self.madnis_lower_box.set_row(subproc_id + 1, [ - f"{subproc_id}", - f"{loss:>.4f}", - f"{channel_count}", - f"{batch_str:<15} {progress_bar}", - ]) - self.madnis_upper_box.print_update() - self.madnis_lower_box.print_update() - - @property - def gridpack_run_mode(self) -> str: - """How this run ends: a gridpack run may skip the events (see the - [gridpack] run_mode run card option); any other run is a regular one.""" - if not self.run_card["gridpack"]["save_gridpack"]: - return "regular" - return self.run_card["gridpack"]["run_mode"] - - def finish_without_events(self) -> None: - """The end of a gridpack run that writes no events: only what the - gridpack needs is prepared, then it is stored.""" - if self.gridpack_run_mode == "minimal": - # a trained MadNIS needs nothing more, a VEGAS grid has to converge - if not self.madnis_trained: - self.event_generator.optimize() - else: - self.event_generator.fix_max_weights() - self.build_systematics() - self.save_gridpack() - - def generate_events(self) -> None: - if self.gridpack_run_mode != "regular": - self.finish_without_events() - return - start_time = get_start_time() - self.event_generator.generate() - output_format = self.run_card["run"]["output_format"] - systematics = self.build_systematics() - histograms = self.build_event_histograms() - if output_format == "compact_npy": - self.lhe_completer = None - self.event_generator.combine_to_compact_npy( - os.path.join(self.run_path, "events.npy"), systematics, histograms - ) - self.write_lhe_header(systematics) - self.save_lhe_completer() - elif output_format == "lhe_npy": - self.lhe_completer = self.build_lhe_completer() - self.event_generator.combine_to_lhe_npy( - os.path.join(self.run_path, "events.npy"), self.lhe_completer, - systematics, histograms - ) - self.write_lhe_header(systematics) - elif output_format == "lhe": - self.lhe_completer = self.build_lhe_completer() - lhe_path = os.path.join(self.run_path, "events.lhe") - self.event_generator.combine_to_lhe( - lhe_path, self.lhe_completer, - self.build_lhe_meta(), systematics, histograms - ) - # Ship the LHE compressed by default. These files are large and - # very compressible, madevent has always stored its events - # gzipped, and every consumer here already accepts either form - # (see _find_event_file). misc.gzip replaces events.lhe with - # events.lhe.gz, and switches to an external multithreaded tool - # above 256 MB. - misc.gzip(lhe_path) - else: - raise ValueError("Unknown output format") - self.write_hwu(histograms) - self.make_plots(histograms) - if systematics is not None: - self.write_systematics_sidecar() - self.log_systematics_summary() - self.save_gridpack() - - def save_lhe_completer(self) -> None: - """Keep the LHE completer next to compact_npy events, so npy_to_lhe - can complete them into an LHE file later (a shower or MadSpin requested - after the run). Never fatal to the run.""" - try: - self.build_lhe_completer().save( - os.path.join(self.run_path, npy_to_lhe.COMPLETER_FILE)) - except Exception as err: - logger.warning("could not save the LHE completer; these events " - "cannot be converted to LHE later: %s", err) - - hwu_file_name = "MADatLO.HwU" - plot_dir_name = "plots" - - def make_plots(self, histograms) -> None: - """Draw the event-sample histograms into Events//plots, with the - scale and PDF bands info.json carries, and say where they landed. - - This does not go through the HwU file: the bands are already computed - per bin by madspace, so drawing them directly is both fewer steps and - more than histograms.py can do without the LHAPDF python module. - """ - if histograms is None or not self.run_card["run"]["make_plots"]: - return - out_dir = os.path.join(self.run_path, self.plot_dir_name) - try: - data = json.loads(histograms.to_json(self.systematics)) - written = plots.render(data, out_dir) - except plots.BackendMissing as error: - logger.info("matplotlib is not available (%s): no plots. The " - "distributions are in %s", - error, os.path.join(self.run_path, "info.json")) - return - except Exception as error: - logger.warning("could not draw the histograms: %s", error) - return - if not written: - return - logger.info("%d plot(s) stored in %s", len(written), out_dir) - - def write_hwu(self, histograms) -> None: - """Write the event-sample histograms next to the events in the HwU - format (MADatLO.HwU), the one an aMC@NLO run writes as MADatNLO.HwU - and madgraph/various/histograms.py reads: an mg7 distribution can then - be overlaid on an NLO one without converting anything in between. The - numbers are the ones info.json carries, bands included, so this is off - by default -- it is a second copy in another format, for the runs that - are going to be plotted.""" - if histograms is None or not self.run_card["run"]["write_hwu"]: - return - try: - data = json.loads(histograms.to_json(self.systematics)) - summary = (json.loads(self.systematics.summary()) - if self.systematics is not None else None) - text = hwu_output.to_hwu(data, summary) - except Exception as error: - logger.warning("could not write the HwU histograms: %s", error) - return - if not text: - return - with open(os.path.join(self.run_path, self.hwu_file_name), "w") as f: - f.write(text) - - def write_systematics_sidecar(self) -> None: - """Describe the variation weights next to the event file - (events.weights.json): ids, scale factors, PDF set/member and the - text, so npy consumers get the same metadata as the LHE - header, plus the per-variation cross sections.""" - if self.systematics is None: - return - data = json.loads(self.systematics.summary()) - data["initrwgt"] = self.systematics.initrwgt() - data["columns"] = ["rwgt_%d" % i for i in self.systematics.weight_ids] - with open(os.path.join(self.run_path, "events.weights.json"), "w") as f: - json.dump(data, f, indent=1) - - @staticmethod - def _histogram_mean(hist_min, hist_max, bin_values): - """Cross-section-weighted mean of a histogrammed observable over its - range. ``bin_values`` carries the underflow and overflow bins first and - last, as both madspace histogram kinds do; they have no position and - are left out.""" - values = list(bin_values)[1:-1] - total = sum(values) - if not values or total == 0: - return None - width = (hist_max - hist_min) / len(values) - return sum(v * (hist_min + (i + 0.5) * width) - for i, v in enumerate(values)) / total - - def _histogram_means(self) -> dict: - """name -> mean of every [histograms] entry, from the post-processing - histograms (the final event sample) when they were filled, otherwise - from the weighted integration histograms.""" - if self.event_histograms is not None: - data = json.loads(self.event_histograms.to_json(self.systematics)) - hists = [(h["name"], h["min"], h["max"], h["bin_values"]) - for h in data] - else: - hists = [(h.name, h.min, h.max, h.bin_values) - for h in self.event_generator.histograms()] - means = {} - for name, hist_min, hist_max, values in hists: - mean = self._histogram_mean(hist_min, hist_max, values) - if mean is not None: - means[name] = mean - return means - - def get_result(self) -> dict: - """Return the run result: cross-section (pb) with MC error, the number - of (unweighted) events, and the mean of every observable declared in the - [histograms] section. Used to build the scan summary.""" - status = self.event_generator.status() - result = {'cross(pb)': status.mean, 'error(pb)': status.error, - 'nb_event': status.count_unweighted} - try: - if self.systematics is not None and self.systematics.weight_count: - summary = json.loads(self.systematics.summary()) - scale = systematics_summary.scale_percentages(summary) - if scale is not None: - result['scale_up(%)'], result['scale_down(%)'] = scale - pdf = systematics_summary.pdf_percentages(summary) - if pdf: - # one column pair; the nominal set's entry comes first - _entry, up, down = pdf[0] - result['pdf_up(%)'], result['pdf_down(%)'] = up, down - except Exception as err: - logger.warning("could not extract the systematics summary: %s", err) - try: - for name, mean in self._histogram_means().items(): - result['<%s>' % name] = mean - except Exception as err: - logger.warning("could not extract observable means: %s", err) - return result - - def _beam_info(self): - """Return (beam_pdg_ids, beam_energies) for the LHE block. - - `incoming` holds the *partonic* initial state (e.g. gluons), not the - beam particle, so hadronic beams are protons (2212); leptonic beams are - the incoming leptons themselves.""" - half_e = float(self.e_cm) / 2. - if self.is_decay: - # No beams. LHE has no way to say that, so report the decaying - # particle at rest as a single "beam"; the second slot is empty. - data = self.subprocess_data[0] - pdg = clean_pids(data["incoming"])[0] - return [pdg, 0], [float(self.e_cm), 0.0] - if not self.leptonic: - # hadronic collider: proton beams (p-pbar is not distinguished) - return [2212, 2212], [half_e, half_e] - data = self.subprocess_data[0] - incoming = data["incoming"] - flavor0 = data["flavors"][0]["options"][0] - init_pdgs = flavor0[:len(incoming)] - beam_pdgs = [pdg if abs(code) in (81, 82) else code - for code, pdg in zip(incoming, init_pdgs)] - return beam_pdgs, [half_e, half_e] - - def _lhapdf_id(self): - """Central LHAPDF id of the beam PDF set (read from its .info SetIndex), - or -1 for a leptonic beam (no PDF).""" - if self.leptonic or not self.pdf_dir: - return -1 - pdf_set = self.run_card["beam"]["pdf"] - info = os.path.join(self.pdf_dir, pdf_set, "%s.info" % pdf_set) - try: - for line in open(info): - if line.strip().startswith("SetIndex:"): - return int(line.split(":", 1)[1].strip()) - except Exception as err: - logger.warning("could not read LHAPDF id from %s: %s", info, err) - return -1 - - def build_lhe_meta(self, systematics=None): - """Build the LHE header/ metadata: the param_card () and the - run_card.toml () headers plus the beam/PDF/cross-section info - needed by downstream tools (systematics, MadSpin, ...). - - `combine_to_lhe` injects the header itself, so callers that - go through it should leave `systematics` unset; pass it only when - building meta for a writer that bypasses that injection (e.g. - write_lhe_header).""" - beam_pdgs, energies = self._beam_info() - lhaid = self._lhapdf_id() - pdf_group = -1 if self.leptonic else 0 - status = self.event_generator.status() - xsec, err = status.mean, status.error - weight_mode, max_weight = lhe_weight_info(status) - with open(self.param_card_path) as f: - param_text = f.read() - with open(os.path.join("Cards", "run_card.toml")) as f: - run_text = f.read() - headers = [] - # generation commands (model + process): read by MadSpin/reweight/... - proc_card = os.path.join("Cards", "proc_card_mg5.dat") - if os.path.exists(proc_card): - with open(proc_card) as f: - headers.append(ms.LHEHeader(name="MG5ProcCard", content=f.read())) - headers.append(ms.LHEHeader(name="slha", content=param_text)) - headers.append(ms.LHEHeader(name="MG7RunCard", content=run_text)) - # The resolved seed (even when the run_card requested a random one via - # seed = -1), so the run can be reproduced from the LHE file alone. - headers.append(ms.LHEHeader(name="MG7Seed", content=str(self.run_seed))) - if systematics is not None and systematics.weight_ids: - headers.append(ms.LHEHeader( - name="initrwgt", content=systematics.initrwgt(), escape_content=False - )) - return ms.LHEMeta( - beam1_pdg_id=beam_pdgs[0], beam2_pdg_id=beam_pdgs[1], - beam1_energy=energies[0], beam2_energy=energies[1], - beam1_pdf_authors=pdf_group, beam2_pdf_authors=pdf_group, - beam1_pdf_id=lhaid, beam2_pdf_id=lhaid, - weight_mode=weight_mode, - # positional: the pybind arg name for max_weight is non-kwarg-safe - processes=[ms.LHEProcess(xsec, err, max_weight, 1)], - headers=headers, - ) - - def write_lhe_header(self, systematics) -> None: - """Write header.lhe next to events.npy: the
/ blocks - (run card, param card, beam/PDF, cross section) that the npy formats - otherwise drop, with no events.""" - header_path = os.path.join(self.run_path, "header.lhe") - writer = ms.LHEFileWriter(header_path, self.build_lhe_meta(systematics)) - del writer # closes the file (writes the closing tag) - - def build_lhe_completer(self): - all_mcdata = ( - [subproc.build_multi_channel_data() for subproc in self.subprocesses] - if self.merged_subprocess_data is None else - [build_multi_channel_data(meta, self) for meta in self.subprocess_data] - ) - subproc_args = [ - ms.SubprocArgs( - topologies = [topo[0] for topo in mcdata.topologies], - permutations = mcdata.permutations, - diagram_indices = mcdata.diagram_indices, - diagram_color_indices = mcdata.diagram_color_indices, - diagram_propagator_pdgs = mcdata.diagram_propagator_pdgs, - color_flows = meta["color_flows"], - pdg_color_types = { - int(key): value - for key, value in meta["pdg_color_types"].items() - }, - helicities = lhe_helicities(meta, self.run_card), - pdg_ids = [flavor["options"] for flavor in meta["flavors"]], - ) - for mcdata, meta in zip(all_mcdata, self.subprocess_data) - ] - return ms.LHECompleter( - subproc_args=subproc_args, - bw_cutoff=self.run_card["phasespace"]["bw_cutoff"] - ) - - def save_gridpack(self) -> None: - if not self.run_card["gridpack"]["save_gridpack"]: - return - - gridpack_path = os.path.join(self.run_path, "gridpack") - data_path = os.path.join(gridpack_path, "data") - events_path = os.path.join(gridpack_path, "Events") - os.mkdir(gridpack_path) - os.mkdir(data_path) - os.mkdir(events_path) - self.contexts[0].save_globals(os.path.join(data_path, "globals.tar")) - - # one json object per channel, joined by hand so that the channel - # json is not parsed and dumped again - channel_entries = [ - f"{json.dumps(channel.status().name)}:{channel.to_json(self.gridpack_run_mode == "fix_max_weight")}" - for channel in self.event_generator.channels() - ] - with open(os.path.join(data_path, "channels.json"), "w") as f: - f.write("{" + ",".join(channel_entries) + "}") - - lib_path = os.path.join(gridpack_path, "lib") - if self.run_card["gridpack"]["include_source"]: - os.mkdir(lib_path) - shutil.copytree("src", os.path.join(gridpack_path, "src")) - shutil.copytree("SubProcesses", os.path.join(gridpack_path, "SubProcesses")) - shutil.copytree("backend", os.path.join(gridpack_path, "backend")) - else: - shutil.copytree("lib", lib_path) - for _, path in self.custom_function_paths: - target = path if os.path.isabs(path) else os.path.join(gridpack_path, path) - if not os.path.exists(target): - os.makedirs(os.path.dirname(target), exist_ok=True) - shutil.copy(path, target) - - if self.run_card["gridpack"]["include_madspace_source"]: - shutil.copytree( - _MADSPACE_DIR, - os.path.join(gridpack_path, "madspace"), - ignore=shutil.ignore_patterns("build", "install"), - ) - - if self.run_card["gridpack"]["include_madspace"]: - shutil.copytree( - _INSTALL_DIR / "madspace", - os.path.join(gridpack_path, "madspace", "install", "madspace"), - ignore=shutil.ignore_patterns("include"), - ) - - matrix_elements = [] - for subproc in self.subprocess_data: - me_path = subproc["me_path"] - matrix_elements.append(me_path) - - cards_path = os.path.join(gridpack_path, "Cards") - os.mkdir(cards_path) - shutil.copy(os.path.join("Cards", "param_card.dat"), cards_path) - # Full run card with a header noting it is read-only in gridpack context. - import io as _io - _buf = _io.StringIO() - self.run_card.write(_buf) - _header = ( - "# This is the run card used to generate this gridpack.\n" - "# Modifying this file will have no effect on gridpack execution.\n" - "# To change event-generation settings, edit grid_run_card.toml.\n\n" - ) - with open(os.path.join(cards_path, "run_card.toml"), 'w') as _f: - _f.write(_header + _buf.getvalue()) - # Minimal card containing only the settings used by generate_events. - # The gridpack ships the libraries that were actually built, and their - # names carry the resolved backend, so the card it runs from has to name - # that backend as well: a cpu_mode of 'auto' would send the gridpack - # looking for a ..._auto.so that was never built. The full run_card.toml - # written just above keeps 'auto', since that is what was asked for. - device_names = self.run_card["run"]["device"] - if not isinstance(device_names, list): - device_names = [device_names] - resolved_cpu_mode = None - for name, backend in zip(device_names, getattr(self, "backends", [])): - if name.split(":")[0] == "cpu": - resolved_cpu_mode = backend - break - previous_cpu_mode = self.run_card["run"]["cpu_mode"] - if resolved_cpu_mode is not None: - self.run_card["run"]["cpu_mode"] = resolved_cpu_mode - # A gridpack is run to feed a shower or detector chain, which reads - # LHE, so it writes LHE by default whatever format this run used - # (bin/generate_events --output_format still selects an npy output). - previous_output_format = self.run_card["run"]["output_format"] - self.run_card["run"]["output_format"] = "lhe" - # the directories of this run mean nothing on the machine running the gridpack - previous_dirs = (self.run_card["gridpack"]["temp_output_dir"], - self.run_card["gridpack"]["output_dir"]) - self.run_card["gridpack"]["temp_output_dir"] = "" - self.run_card["gridpack"]["output_dir"] = "" - try: - self.run_card.write_gridpack_card( - os.path.join(cards_path, "grid_run_card.toml")) - finally: - self.run_card["run"]["cpu_mode"] = previous_cpu_mode - self.run_card["run"]["output_format"] = previous_output_format - (self.run_card["gridpack"]["temp_output_dir"], - self.run_card["gridpack"]["output_dir"]) = previous_dirs - - bin_path = os.path.join(gridpack_path, "bin") - os.mkdir(bin_path) - gen_events_file = os.path.join(bin_path, "generate_events") - template_dir = os.path.dirname(__file__) - shutil.copy(os.path.join(template_dir, "gridpack.py"), gen_events_file) - shutil.copy(os.path.join(template_dir, "custom_function_loader.py"), bin_path) - os.chmod(gen_events_file, 0o755) - shutil.copy(os.path.join(template_dir, "gridpack_setup.py"), bin_path) - # streaming needs the fixed maximum weights of this mode - if self.gridpack_run_mode == "fix_max_weight": - shutil.copy(os.path.join(template_dir, "event_stream.py"), bin_path) - - data = { - "matrix_elements": matrix_elements, - # loaded before the matrix elements, as in the run - "custom_functions": self.custom_function_paths, - "me_parameters": me_parameters(self.run_card), - "source_hash": ms.SOURCE_HASH, - } - with open(os.path.join(data_path, "data.json"), "w") as f: - json.dump(data, f) - - if self.lhe_completer is None: - self.lhe_completer = self.build_lhe_completer() - self.lhe_completer.save(os.path.join(data_path, "lhe.json")) - self.save_gridpack_lhe_meta(os.path.join(data_path, "lhe_meta.json")) - if self.systematics_data is not None: - systematics_data = dict(self.systematics_data) - # A decay takes alpha_s from a constant grid written into the run - # directory rather than from an LHAPDF set, so the gridpack cannot - # look that file up: ship it and point at the shipped copy. - if self.fixed_alphas_info and \ - systematics_data["nominal_info_file"] == self.fixed_alphas_info: - name = os.path.basename(self.fixed_alphas_info) - shutil.copy(self.fixed_alphas_info, os.path.join(data_path, name)) - systematics_data["nominal_info_file"] = os.path.join("data", name) - with open(os.path.join(data_path, "systematics.json"), "w") as f: - json.dump(systematics_data, f) - if self.run_card["gridpack"]["compress"]: - shutil.make_archive(gridpack_path, "gztar", self.run_path, "gridpack") - shutil.rmtree(gridpack_path) - - def save_gridpack_lhe_meta(self, path) -> None: - """The run-independent part of the LHE header (cards, beams, PDF), for - the gridpack to write complete LHE headers: it adds the cross section - and seed of its own run, which is why those two are left out here.""" - meta = self.build_lhe_meta() - data = {key: getattr(meta, key) for key in ( - "beam1_pdg_id", "beam2_pdg_id", "beam1_energy", "beam2_energy", - "beam1_pdf_authors", "beam2_pdf_authors", "beam1_pdf_id", - "beam2_pdf_id", "weight_mode")} - data["headers"] = [ - {"name": h.name, "content": h.content, - "escape_content": h.escape_content} - for h in meta.headers if h.name != "MG7Seed"] - with open(path, "w") as f: - json.dump(data, f) - - def get_mass(self, pid: int) -> float: - return self.param_card.get_value("mass", pid) - - def get_width(self, pid: int) -> float: - return self.param_card.get_value("width", pid) - - -# Flavor-merged legs carry a group id instead of a pdg. Every member of a group -# shares the same mass -- that is what makes them mergeable -- so any member is -# a valid representative for the kinematics. -_MERGED_PID_REPRESENTATIVE = { - 81: 1, # light quarks d u s c - 82: 11, # charged leptons e mu - 83: 12, # neutrinos ve vm vt -} - - -def clean_pids(pids: list[int]) -> list[int]: - pids_out = [] - for pid in pids: - pid = abs(pid) - if pid in _MERGED_PID_REPRESENTATIVE: - pid = _MERGED_PID_REPRESENTATIVE[pid] - elif 81 <= pid <= 99: - # Reserved for flavor merging. Anything outside this window is a - # real pdg (BSM models use codes in the millions), so let it pass. - raise ValueError( - f"unknown flavor-merged particle id {pid}; add its " - "representative to _MERGED_PID_REPRESENTATIVE" - ) - pids_out.append(pid) - return pids_out - - -# observables that tell a particle from its antiparticle, and so need the signed -# pdg ids of the actual flavours rather than clean_pids' unsigned representatives -_SIGNED_OBSERVABLES = {"sfos_pair_mass"} - - -def _flavor_options(meta: dict, unmerged_meta: dict | None = None) -> list[list[int]]: - """Every flavour assignment (signed pdg ids, incoming then outgoing) a - subprocess evaluates with the same momenta.""" - if "flavors" in meta: - return [ - option for flavor in meta["flavors"] for option in flavor["options"] - ] - return [ - option - for subproc in meta.get("subprocesses", []) - for option in _flavor_options(unmerged_meta[subproc]) - ] - - -def observable_pids( - meta: dict, observable_kwargs: dict, unmerged_meta: dict | None = None -) -> list[int]: - """The pdg ids an observable of this subprocess is built with. - - Most observables only select particles by group, which the unsigned - representatives of clean_pids do. One that pairs a particle with its - antiparticle needs the real signed flavours: those of the first flavour - assignment, provided every assignment of the subprocess forms the same - pairs -- the cut is applied once to momenta that all of them share. - """ - all_pids = clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) - if observable_kwargs.get("observable") not in _SIGNED_OBSERVABLES: - return all_pids - options = _flavor_options(meta, unmerged_meta) - if not options: - raise ValueError( - f"{observable_kwargs.get('name')}: no flavour information for " - f"subprocess {meta['incoming']} > {meta['outgoing']}" - ) - n_in = len(meta["incoming"]) - groups = [set(group) for group in observable_kwargs["select_pids"]] - if len(groups) == 1: - groups = groups * 2 - - def sfos_pairs(pids): - out = range(n_in, len(pids)) - return frozenset( - frozenset((i, j)) - for i in out if pids[i] in groups[0] - for j in out if pids[j] in groups[1] - if i != j and pids[i] != 0 and pids[i] == -pids[j] - ) - - reference = sfos_pairs(options[0]) - if any(sfos_pairs(option) != reference for option in options[1:]): - raise ValueError( - f"{observable_kwargs.get('name')}: the flavours grouped into subprocess " - f"{meta['incoming']} > {meta['outgoing']} pair the particles " - "differently, so a same-flavour opposite-sign cut cannot be applied " - "to their common momenta; regenerate the process after " - "'set apply_flavor_grouping False'" - ) - return list(options[0]) - - -def build_observable(pids: list[int], n_in: int, observable_kwargs: dict): - """ms.Observable of a subprocess with `n_in` incoming particles. - - madspace takes the first two momenta to be the beams (ignore_incoming), so - for a 1 -> N decay it would drop the first decay product from every - selection. There the selection runs over all the momenta instead, with the - incoming particle set to pdg id 0, which is in no group. - """ - if n_in == 2: - return ms.Observable(pids, **observable_kwargs) - pids = [0] * n_in + list(pids[n_in:]) - return ms.Observable(pids, **dict(observable_kwargs, ignore_incoming=False)) - - -def decay_products(meta: dict, unmerged_meta: dict | None = None) -> set[int]: - """Outgoing positions (0-based, beams excluded) of the particles that come - from an on-shell (decay-chain) propagator, as MadEvent's check_decay marks - them. Every channel of a subprocess shares its decay chains, so the first - channel tells, read through its first diagram's permutation.""" - if not meta["channels"]: - return set() - channel = meta["channels"][0] - topo_channel = channel - if "vertices" not in channel: - topo_channel = unmerged_meta[channel["subprocess"]]["channels"][channel["channel"]] - on_shell = topo_channel.get("on_shell_propagators", []) - if not on_shell: - return set() - vertices = topo_channel["vertices"] - # the closing vertex creates no propagator - created = {vertex[-1]: vertex[:-1] for vertex in vertices[:-1]} - - def legs(name): - if name in created: - return set().union(*(legs(child) for child in created[name])) - return {name} - - n_in = len(meta["incoming"]) - n_out = len(meta["outgoing"]) - products = set() - for prop in on_shell: - external = legs(f"p{prop}") - incoming = {leg for leg in external if leg.startswith("i")} - outgoing = {int(leg[1:]) for leg in external if leg.startswith("o")} - if not incoming: - products |= outgoing - elif len(incoming) == n_in: - # built from the beams: the decay is everything else - products |= set(range(n_out)) - outgoing - # event[i] = topology[permutation[i]] - permutation = channel["diagrams"][0]["permutation"] - return { - i - n_in for i in range(n_in, len(permutation)) - if permutation[i] - n_in in products - } - - -def pid_is_qcd(pid: int): - return abs(pid) in [21, 1, 2, 3, 4, 5, 6, 81] - - -def build_topologies( - incoming_masses: list[float], - outgoing_masses: list[float], - channel: dict, - process: MadgraphProcess -) -> list[ms.Topology]: - propagators = [] - for i, (pid, signed_pid) in enumerate(zip( - clean_pids(channel["propagators"]), channel["propagators"] - )): - mass = process.get_mass(pid) - width = process.get_width(pid) - bw_cutoff = process.run_card["phasespace"]["bw_cutoff"] - if i in channel["on_shell_propagators"]: - e_min = mass - bw_cutoff * width - e_max = mass + bw_cutoff * width - else: - e_min = 0 - e_max = 0 - # $-excluded propagator: the matrix element vanishes within bw_cutoff - # widths of the pole, so do not pile the channel's points up there - # (a flat density in that window instead of the Breit-Wigner peak). - # Channels written before the exporter recorded them have none. - flat_window = 0. - if i in channel.get("dollar_propagators", ()) and width > 0: - flat_window = float(bw_cutoff) - propagators.append(ms.Propagator( - mass=mass, - width=width, - integration_order=0, - e_min=e_min, - e_max=e_max, - pdg_id=signed_pid, - flat_window=flat_window, - )) - vertices = channel["vertices"] - diag = ms.Diagram( - incoming_masses, outgoing_masses, propagators, vertices - ) - return ms.Topology.topologies(diag) - - -def build_multi_channel_data( - meta: dict, process: MadgraphProcess, unmerged_meta: dict | None = None -) -> MultiChannelData: - incoming_masses = [ - process.get_mass(pid) for pid in clean_pids(meta["incoming"]) - ] - outgoing_masses = [ - process.get_mass(pid) for pid in clean_pids(meta["outgoing"]) - ] - - if unmerged_meta is None: - diagram_count = meta["diagram_count"] - amp2_remaps = [[-1] * diagram_count] - else: - amp2_remaps = [ - [-1] * unmerged_meta[subproc]["diagram_count"] - for subproc in meta["subprocesses"] - ] - symfact = [] - topologies = [] - permutations = [] - channel_indices = [] - channel_weight_indices = [] - diagram_indices = [] - diagram_color_indices = [] - diagram_propagator_pdgs = [] - active_flavors = [] - channel_index = 0 - qcd_s_channel_count = [] - - for channel in meta["channels"]: - if unmerged_meta is None: - topo_channel = channel - else: - topo_subproc = channel["subprocess"] - topo_channel_index = channel["channel"] - topo_channel = unmerged_meta[topo_subproc]["channels"][topo_channel_index] - chan_topologies = build_topologies( - incoming_masses, outgoing_masses, topo_channel, process - ) - topo_count = len(chan_topologies) - if topo_count == 0: - continue - - topo = chan_topologies[0] - s_chan_count = [0] * len(topo.decays) - non_qcd = [False] * len(topo.decays) - pdg_ids = [decay.pdg_id for decay in topo.decays] - for i, pid in zip(topo.outgoing_indices, meta["outgoing"]): - pdg_ids[i] = pid - for decay in reversed(topo.decays): - if len(decay.child_indices) == 0: - continue - if decay.index == 0 and topo.t_propagator_count > 0: - s_chan_count[0] = sum(s_chan_count[i] for i in decay.child_indices) - break - - is_qcd = pid_is_qcd(pdg_ids[decay.index]) and all( - pid_is_qcd(pdg_ids[i]) for i in decay.child_indices - ) - is_non_qcd = decay.on_shell or not is_qcd or any( - non_qcd[i] for i in decay.child_indices - ) - non_qcd[decay.index] = is_non_qcd - s_chan_count[decay.index] = ( - 0 if is_non_qcd else sum(s_chan_count[i] for i in decay.child_indices) + 1 - ) - qcd_s_channel_count.append(s_chan_count[0]) - - diagrams = channel["diagrams"] - chan_permutations = [d["permutation"] for d in diagrams] - if unmerged_meta is None: - amp2_remaps[0][diagrams[0]["diagram"]] = channel_index - else: - for amp2_remap, diag in zip(amp2_remaps, diagrams[0]["diagram"]): - if diag != -1: - amp2_remap[diag] = channel_index - - channel_index_first = channel_index - symfact_index_first = len(symfact) - channel_index += 1 - symfact.extend([None] * topo_count) - for d in diagrams[1:]: - if unmerged_meta is None: - amp2_remaps[0][d["diagram"]] = channel_index - else: - for amp2_remap, diag in zip(amp2_remaps, d["diagram"]): - if diag != -1: - amp2_remap[diag] = channel_index - channel_index += 1 - symfact.extend(range(symfact_index_first, symfact_index_first + topo_count)) - - topologies.append(chan_topologies) - permutations.append(chan_permutations) - channel_indices.append(list(range(channel_index_first, channel_index))) - channel_weight_indices.append([ - [ - symfact_index_first + topo_index + i * topo_count - for i in range(len(chan_permutations)) - ] - for topo_index in range(topo_count) - ]) - diagram_indices.append([d["diagram"] for d in diagrams]) - if unmerged_meta is None: - diagram_color_indices.append([d["active_colors"] for d in diagrams]) - diagram_propagator_pdgs.append( - [d["propagator_pdgs"] for d in diagrams] - ) - active_flavors.append([d["active_flavors"] for d in diagrams]) - - return MultiChannelData( - amp2_remaps, - symfact, - topologies, - permutations, - channel_indices, - channel_weight_indices, - diagram_indices, - diagram_color_indices, - diagram_propagator_pdgs, - active_flavors, - qcd_s_channel_count, - ) - - -class MadgraphSubprocess: - def __init__( - self, - process: MadgraphProcess, - meta: dict, - subproc_id: int, - unmerged_meta: dict | None = None - ): - self.process = process - self.meta = meta - self.subproc_id = subproc_id - self.multi_channel_data = None - - self.unmerged_meta = None - if unmerged_meta is None: - api_path_formats = [self.meta["me_path"]] - member_metas = [self.meta] - else: - api_path_formats = [] - member_metas = [] - for subproc in self.meta["subprocesses"]: - submeta = unmerged_meta[subproc] - api_path_formats.append(submeta["me_path"]) - member_metas.append(submeta) - if len(api_path_formats) == 1: - self.meta = submeta - else: - self.unmerged_meta = unmerged_meta - - # The libraries were all built up front by MadgraphProcess.compile_matrix_elements - all_api_paths = [ - [api_path_format.format(device=backend) for backend in self.process.backends] - for api_path_format in api_path_formats - ] - - self.incoming_masses = [ - self.process.get_mass(pid) for pid in clean_pids(self.meta["incoming"]) - ] - self.outgoing_masses = [ - self.process.get_mass(pid) for pid in clean_pids(self.meta["outgoing"]) - ] - self.particle_count = len(self.incoming_masses) + len(self.outgoing_masses) - # MadEvent's cut_decays = F: the decay products of on-shell propagators - # are left out of every cut selection (a pair cut then needs both - # members to be cut), while the histograms still see them - if self.process.run_card["phasespace"]["cut_decays"]: - uncut = set() - else: - uncut = decay_products(self.meta, self.unmerged_meta) - n_in = len(self.incoming_masses) - - def cut_pids(observable_kwargs): - pids = observable_pids(self.meta, observable_kwargs, self.unmerged_meta) - # pdg id 0 is in no group, so the particle is never selected - return [0 if i - n_in in uncut else pid for i, pid in enumerate(pids)] - - # custom cuts and scales see the PDG ids of the first flavor, like the - # built-in cuts, which select by flavor-independent groups - representative_pids = [self.flavor_pids()[0]] - custom_cut = self.process.custom_function( - "cut", self.particle_count, representative_pids - ) - if len(self.process.cut_data) > 0: - self.cuts = ms.Cuts([ - ms.CutItem( - observable=build_observable( - cut_pids(cut_item.observable_kwargs), - n_in, cut_item.observable_kwargs - ), - min=cut_item.min, - max=cut_item.max, - mode=cut_item.mode, - ) - for cut_item in self.process.cut_data - ], custom_cut) - elif custom_cut is not None: - self.cuts = ms.Cuts(self.particle_count, custom_cut) - else: - self.cuts = None - # the integration histograms are functions of the momenta; the weight - # distribution is a property of the final event sample, so it is only - # filled by MadgraphProcess.build_event_histograms at combine time - # They cost an observable evaluation per integration point, and the - # plots/HwU are drawn from the post-processing ones instead, so they - # are only filled on request ([run] weighted_histograms). - momentum_hists = [item for item in self.process.hist_data - if not item.from_weight] \ - if self.process.run_card["run"]["weighted_histograms"] else [] - self.histograms = ( - ms.ObservableHistograms([ - ms.HistItem( - observable=build_observable( - observable_pids( - self.meta, hist_item.observable_kwargs, self.unmerged_meta - ), - n_in, hist_item.observable_kwargs - ), - min=hist_item.min, - max=hist_item.max, - bin_count=hist_item.bin_count, - ) - for hist_item in momentum_hists - ]) - if momentum_hists - else None - ) - - self.scale = ms.EnergyScale( - particle_count=self.particle_count, - custom_scale=self.process.custom_function( - "scale", self.particle_count, representative_pids - ), - **self.process.scale_kwargs - ) - - if self.process.run_card["run"]["dummy_matrix_element"]: - self.matrix_elements = [None] * len(all_api_paths) - else: - self.matrix_elements = [] - for api_paths in all_api_paths: - for context, api_path in zip(self.process.contexts, api_paths): - mat = context.load_matrix_element( - api_path, self.process.param_card_path, - me_parameters(self.process.run_card) - ) - self.matrix_elements.append(mat) - self.surrogates = [] - if self.process.surrogate is not None: - for path_format, member_meta in zip(api_path_formats, member_metas): - self.surrogates.append(self.load_surrogate(path_format, member_meta)) - - def load_surrogate(self, me_path_format: str, member_meta: dict): - """The [surrogate] of one matrix element of the subprocess.""" - name = re.sub( - r"^libmadmatrix_(.*)_\{device\}\.so$", r"\1", os.path.basename(me_path_format) - ) - # PDG ids for each flavor index of the matrix element - flavors = member_meta["flavors"] - pdg_ids = [flavors[0]["options"][0]] * (max(f["index"] for f in flavors) + 1) - for flavor in flavors: - pdg_ids[flavor["index"]] = flavor["options"][0] - settings = self.process.surrogate - apis = load_surrogate( - ms, - self.process.contexts, - self.process.backends, - self.process.device_types, - settings["function"], - name, - pdg_ids, - settings["param_card"], - ) - if len({api.index() for api in apis}) != 1: - raise RuntimeError("surrogate loaded at different indices") - return apis[0] - - def flavor_pids(self) -> list[list[int]]: - """PDG ids of the external particles for every flavor of the - integrand, in the order of its flavor index.""" - pids = [] - for flav in self.meta["flavors"]: - if self.unmerged_meta is not None: - subproc_index = self.meta["subprocesses"][flav["subprocess"]] - flav = self.unmerged_meta[subproc_index]["flavors"][flav["flavor"]] - pids.append(flav["options"][0]) - return pids - - def build_multi_channel_data(self) -> MultiChannelData: - if self.multi_channel_data is not None: - return self.multi_channel_data - self.multi_channel_data = build_multi_channel_data( - self.meta, self.process, self.unmerged_meta - ) - return self.multi_channel_data - - def build_multichannel_phasespace(self) -> PhaseSpace: - mcdata = self.build_multi_channel_data() - channel_count = sum(len(topos) for topos in mcdata.topologies) - drop_threshold = self.process.run_card["phasespace"]["drop_qcd_s_channel"] - if drop_threshold >= 0 and channel_count > drop_threshold: - mcdata = self.drop_qcd_s_channels(mcdata) - mcdata = self.drop_empty_channels(mcdata) - - channels = [] - t_channel_mode = self.t_channel_mode( - self.process.run_card["phasespace"]["t_channel"] - ) - for channel_id, (chan_topologies, chan_permutations, chan_indices, active_flavors) in enumerate(zip( - mcdata.topologies, mcdata.permutations, mcdata.channel_weight_indices, - mcdata.active_flavors - )): - topo_count = len(chan_topologies) - for topo_index, (topo, indices) in enumerate(zip(chan_topologies, chan_indices)): - mapping = ms.PhaseSpaceMapping( - chan_topologies[0], - self.process.e_cm, - t_channel_mode=t_channel_mode, - cuts=self.cuts, - invariant_power=self.process.run_card["phasespace"]["invariant_power"], - permutations=chan_permutations, - leptonic=self.process.leptonic, - ) - prefix = f"subproc{self.subproc_id}.channel{channel_id}" - if topo_count > 1: - prefix += f".subchan{topo_index}" - discrete_sym, discrete_flavor = self.build_discrete( - len(chan_permutations), len(self.meta["flavors"]), prefix - ) - channels.append(Channel( - phasespace_mapping = mapping, - adaptive_mapping = self.build_vegas(mapping, prefix), - discrete_sym = discrete_sym, - discrete_flavor = discrete_flavor, - channel_weight_indices = indices, - name = f"{channel_id}", - active_flavors = active_flavors, - )) - - remapped_chan_count = sum( - len(indices) for indices in mcdata.channel_indices - ) - if self.process.run_card["phasespace"]["sde_strategy"] == "denominators": - prop_chan_weights = ms.PropagatorChannelWeights( - [topo[0] for topo in mcdata.topologies], mcdata.permutations, - mcdata.channel_indices - ) - chan_weight_remap = [] - else: - prop_chan_weights = None - chan_weight_remap = [ - [ - len(mcdata.symfact) if remap == -1 else remap - for remap in amp2_remap - ] - for amp2_remap in mcdata.amp2_remaps - ] - - if any(len(topos) > 1 for topos in mcdata.topologies): - subchan_weights = ms.SubchannelWeights( - mcdata.topologies, mcdata.permutations, mcdata.channel_indices - ) - else: - subchan_weights = None - - return PhaseSpace( - mode="multichannel", - channels=channels, - first_chan_weight_remap=chan_weight_remap, - first_remapped_chan_count=remapped_chan_count, - symfact=mcdata.symfact, - prop_chan_weights=prop_chan_weights, - subchan_weights=subchan_weights, - ) - - def build_flat_phasespace(self) -> PhaseSpace: - mapping = ms.PhaseSpaceMapping( - self.incoming_masses + self.outgoing_masses, - self.process.e_cm, - mode=self.t_channel_mode(self.process.run_card["phasespace"]["flat_mode"]), - cuts=self.cuts, - leptonic=self.process.leptonic, - ) - prefix = f"subproc{self.subproc_id}.flat" - discrete_sym, discrete_flavor = self.build_discrete( - 1, len(self.meta["flavors"]), prefix - ) - channel = Channel( - phasespace_mapping = mapping, - adaptive_mapping = self.build_vegas(mapping, prefix), - discrete_sym = discrete_sym, - discrete_flavor = discrete_flavor, - channel_weight_indices = [0], - name = "F", - active_flavors = [], - ) - if self.unmerged_meta is None: - remap = [list(range(self.meta["diagram_count"]))] - else: - remap = [ - list(range(self.unmerged_meta[subproc]["diagram_count"])) - for subproc in self.meta["subprocesses"] - ] - return PhaseSpace( - mode="flat", - channels=[channel], - first_chan_weight_remap=remap, - first_remapped_chan_count=1, - symfact=[None], - ) - - def simplify_phasespace( - self, - multi_phasespace: PhaseSpace, - cross_sections: list[float] - ) -> PhaseSpace | None: - assert multi_phasespace.mode == "multichannel" - - threshold = 1 - self.process.run_card["phasespace"]["combine_channel_threshold"] - kept_channels = [] - tot_cs = sum(cross_sections) - cum_cs = 0. - seen_active_flavors = set() - #seen_resonances = set() - for index, (cs, chan) in sorted( - enumerate(zip(cross_sections, multi_phasespace.channels)), - key=lambda pair: pair[1][0], - reverse=True - ): - cum_cs += cs - has_unseen_flavors = False - has_unseen_resonances = False - for flavs in chan.active_flavors: - for flav in flavs: - if flav not in seen_active_flavors: - has_unseen_flavors = True - seen_active_flavors.add(flav) - #for resonance in chan.resonances: - # if resonance not in seen_resonances: - # has_unseen_resonances = True - # seen_resonances.add(flav) - if has_unseen_flavors or has_unseen_resonances or not tot_cs \ - or cum_cs / tot_cs < threshold: - kept_channels.append(index) - if len(kept_channels) >= len(cross_sections) - 1: - return multi_phasespace - - channels = [] - channel_map = {} - symfact = [] - for old_chan_index in kept_channels: - channel = multi_phasespace.channels[old_chan_index] - perm_count = max(1, channel.phasespace_mapping.channel_count()) - channel_index = len(symfact) - symfact.append(None) - symfact.extend([channel_index] * (perm_count - 1)) - channel_map.update({ - old_index: new_index - for new_index, old_index in enumerate( - channel.channel_weight_indices, start=channel_index - ) - }) - channels.append(Channel( - phasespace_mapping = channel.phasespace_mapping, - adaptive_mapping = channel.adaptive_mapping, - discrete_sym = channel.discrete_sym, - discrete_flavor = channel.discrete_flavor, - channel_weight_indices = list(range( - channel_index, channel_index + perm_count - )), - name = channel.name, - active_flavors = channel.active_flavors, - event_generator = channel.event_generator, - )) - - flat_phasespace = self.build_flat_phasespace() - flat_channel = flat_phasespace.channels[0] - channels.append(Channel( - phasespace_mapping = flat_channel.phasespace_mapping, - adaptive_mapping = flat_channel.adaptive_mapping, - discrete_sym = flat_channel.discrete_sym, - discrete_flavor = flat_channel.discrete_flavor, - channel_weight_indices = [len(symfact)], - name = flat_channel.name, - active_flavors = flat_channel.active_flavors, - )) - flat_index = len(symfact) - symfact.append(None) - channel_map[len(multi_phasespace.symfact)] = len(symfact) - if multi_phasespace.subchan_weights is None and len(multi_phasespace.first_chan_weight_remap) > 0: - first_chan_weight_remap = [ - [ - channel_map.get(remap, flat_index) - for remap in cw_remap - ] - for cw_remap in multi_phasespace.first_chan_weight_remap - ] - first_remapped_chan_count = len(symfact) - second_chan_weight_remap = [] - second_remapped_chan_count = 0 - else: - first_chan_weight_remap = multi_phasespace.first_chan_weight_remap - first_remapped_chan_count = multi_phasespace.first_remapped_chan_count - chan_count = multi_phasespace.first_remapped_chan_count if multi_phasespace.subchan_weights is None else multi_phasespace.subchan_weights.channel_count() - second_chan_weight_remap = [ - channel_map.get(i, flat_index) - for i in range(chan_count) - ] - second_remapped_chan_count = len(symfact) - - return PhaseSpace( - mode="both", - channels=channels, - first_chan_weight_remap=first_chan_weight_remap, - first_remapped_chan_count=first_remapped_chan_count, - second_chan_weight_remap=second_chan_weight_remap, - second_remapped_chan_count=second_remapped_chan_count, - symfact=symfact, - prop_chan_weights=multi_phasespace.prop_chan_weights, - subchan_weights=multi_phasespace.subchan_weights, - ) - - def drop_qcd_s_channels(self, mcdata: MultiChannelData) -> MultiChannelData: - """Drop channels with non-resonant QCD s-channel propagators, to reduce the - channel count for processes with many diagrams. Channels are grouped by the - number of QCD s-channels, allowing for more if necessary to map out all flavor - indices. Channel weights belonging to a dropped channel are left unmapped.""" - groups_by_s_count = {} - for index, count in enumerate(mcdata.qcd_s_channel_count): - groups_by_s_count.setdefault(count, []).append(index) - - covered_flavors = set() - kept_groups = set() - for s_count in sorted(groups_by_s_count): - s_count_groups = groups_by_s_count[s_count] - if s_count == 0: - selected = s_count_groups - else: - selected = [ - index - for index in s_count_groups - if any( - flav not in covered_flavors - for flavs in mcdata.active_flavors[index] - for flav in flavs - ) - ] - kept_groups.update(selected) - for index in selected: - for flavs in mcdata.active_flavors[index]: - covered_flavors.update(flavs) - - return self.keep_channel_groups(mcdata, sorted(kept_groups)) - - def drop_empty_channels(self, mcdata: MultiChannelData) -> MultiChannelData: - """Drop the channels whose phase space the cuts exclude entirely: a pair - mass cut whose floor lies at or above the top of an on-shell window of - the channel (ms.PhaseSpaceMapping.empty). Every point such a channel - generates fails the cuts, so it could only ever fail to find one. - - If that leaves a flavour with no channel at all, its cross section is - zero -- typically a cut that excludes the window of a decay chain every - diagram shares -- and the run cannot sample it: that is reported as an - error naming the cause, rather than as the survey finding no point - passing the cuts.""" - if self.cuts is None: - return mcdata - t_channel_mode = self.t_channel_mode( - self.process.run_card["phasespace"]["t_channel"] - ) - empty = [ - ms.PhaseSpaceMapping( - chan_topologies[0], - self.process.e_cm, - t_channel_mode=t_channel_mode, - cuts=self.cuts, - invariant_power=self.process.run_card["phasespace"]["invariant_power"], - permutations=chan_permutations, - leptonic=self.process.leptonic, - ).empty() - for chan_topologies, chan_permutations in zip( - mcdata.topologies, mcdata.permutations - ) - ] - if not any(empty): - return mcdata - kept_groups = [index for index, is_empty in enumerate(empty) if not is_empty] - covered_flavors = { - flav - for index in kept_groups - for flavs in mcdata.active_flavors[index] - for flav in flavs - } - uncovered = sorted({ - flav - for index, is_empty in enumerate(empty) if is_empty - for flavs in mcdata.active_flavors[index] - for flav in flavs - } - covered_flavors) - if uncovered: - raise ValueError( - f"subprocess {self.meta['incoming']} > {self.meta['outgoing']}: " - "the pair mass cuts exclude the whole on-shell window " - f"(bw_cutoff = {self.process.run_card['phasespace']['bw_cutoff']}) " - f"of every channel of flavour(s) {uncovered}, so their cross " - "section is zero; loosen the cut or remove that decay" - ) - logger.info( - "subprocess %d: %d channel(s) dropped, the cuts exclude the " - "on-shell window they sample", self.subproc_id, len(empty) - len(kept_groups) - ) - return self.keep_channel_groups(mcdata, kept_groups) - - def keep_channel_groups( - self, mcdata: MultiChannelData, kept_groups: list[int] - ) -> MultiChannelData: - """Rebuild the multi-channel data with only the channel groups (indices - into mcdata.topologies) in kept_groups. Channel weights belonging to a - dropped channel are left unmapped.""" - if len(kept_groups) == len(mcdata.topologies): - return mcdata - - amp2_remaps = [[-1] * len(remap) for remap in mcdata.amp2_remaps] - symfact = [] - topologies = [] - permutations = [] - channel_indices = [] - channel_weight_indices = [] - diagram_indices = [] - diagram_color_indices = [] - diagram_propagator_pdgs = [] - active_flavors = [] - qcd_s_channel_count = [] - channel_index = 0 - - for group in kept_groups: - chan_topologies = mcdata.topologies[group] - chan_permutations = mcdata.permutations[group] - chan_diagram_indices = mcdata.diagram_indices[group] - topo_count = len(chan_topologies) - - channel_index_first = channel_index - symfact_index_first = len(symfact) - for i, diag in enumerate(chan_diagram_indices): - if self.unmerged_meta is None: - amp2_remaps[0][diag] = channel_index - else: - for amp2_remap, d in zip(amp2_remaps, diag): - if d != -1: - amp2_remap[d] = channel_index - if i == 0: - symfact.extend([None] * topo_count) - else: - symfact.extend(range(symfact_index_first, symfact_index_first + topo_count)) - channel_index += 1 - - topologies.append(chan_topologies) - permutations.append(chan_permutations) - channel_indices.append(list(range(channel_index_first, channel_index))) - channel_weight_indices.append([ - [ - symfact_index_first + topo_index + i * topo_count - for i in range(len(chan_permutations)) - ] - for topo_index in range(topo_count) - ]) - diagram_indices.append(chan_diagram_indices) - if mcdata.diagram_color_indices: - diagram_color_indices.append(mcdata.diagram_color_indices[group]) - if mcdata.diagram_propagator_pdgs: - diagram_propagator_pdgs.append( - mcdata.diagram_propagator_pdgs[group] - ) - active_flavors.append(mcdata.active_flavors[group]) - qcd_s_channel_count.append(mcdata.qcd_s_channel_count[group]) - - return MultiChannelData( - amp2_remaps, - symfact, - topologies, - permutations, - channel_indices, - channel_weight_indices, - diagram_indices, - diagram_color_indices, - diagram_propagator_pdgs, - active_flavors, - qcd_s_channel_count, - ) - - def set_madnis_auto_settings(self, rsd: float): - madnis_args = self.process.run_card["madnis"] - n_out = len(self.meta["outgoing"]) - n_events = self.process.run_card["generation"]["events"] - is_gridpack = self.process.run_card["gridpack"]["save_gridpack"] - train_batches = madnis_args["train_batches"] - hidden_dim = min(max(int((7 * rsd) / 32) * 32 + 64, 64), 256) - flow_layers = 3 if rsd < 32 else 4 - lr = min(max((10000 - train_batches) / 8000 * 7e-4 + 3e-4, 3e-4), 1e-3) - if n_out <= 2: - enable = False - elif n_out == 3: - enable = rsd > 10. or n_events > 1000000 or is_gridpack - else: - enable = True - fixed_cwnet_fraction = max(0.33, 1.0 - 10000. / train_batches) - batch_size_per_channel = min(max(int((7 * rsd) / 32) * 32 + 64, 128), 512) - - self.madnis_settings = { - "enable": enable, - "flow_layers": flow_layers, - "flow_hidden_dim": hidden_dim, - "discrete_hidden_dim": hidden_dim, - "cwnet_hidden_dim": hidden_dim, - "train_batches": train_batches, - "lr": lr, - "fixed_cwnet_fraction": fixed_cwnet_fraction, - "batch_size_per_channel": batch_size_per_channel, - } - for key, value in self.madnis_settings.items(): - run_card_value = madnis_args[key] - if run_card_value != "auto": - self.madnis_settings[key] = run_card_value - - def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: - madnis_args = self.process.run_card["madnis"] - # Shared across all networks below: initialize_globals() derives an - # independent, non-colliding stream per tensor from this one base seed. - seed = self.process.run_seed - channels = [] - for channel_id, channel in enumerate(phasespace.channels): - prefix = f"subproc{self.subproc_id}.channel{channel_id}" - cond_dim = 0 - - flow_dim = channel.phasespace_mapping.random_dim() - flow = ms.Flow( - input_dim=flow_dim, - condition_dim=cond_dim, - prefix=prefix, - bin_count=madnis_args["flow_spline_bins"], - subnet_hidden_dim=self.madnis_settings["flow_hidden_dim"], - subnet_layers=self.madnis_settings["flow_layers"], - subnet_activation=self.activation(madnis_args["flow_activation"]), - invert_spline=madnis_args["flow_invert_spline"], - ) - if channel.adaptive_mapping is None: - flow.initialize_globals(self.process.contexts[0], seed) - else: - flow.initialize_from_vegas( - self.process.contexts[0], channel.adaptive_mapping.grid_name(), - seed - ) - cond_dim += flow_dim - - # discrete_sym runs after the adaptive map, so it can condition on its latent. - discrete_sym = channel.discrete_sym - if discrete_sym is not None: - perm_count = channel.phasespace_mapping.channel_count() - discrete_sym = ms.DiscreteFlow( - option_counts=[perm_count], - prefix=f"{prefix}.discrete_flow_sym", - dims_with_prior=[], - condition_dim=cond_dim, - subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], - subnet_layers=madnis_args["discrete_layers"], - subnet_activation=self.activation(madnis_args["discrete_activation"]), - ) - discrete_sym.initialize_globals(self.process.contexts[0], seed) - cond_dim += perm_count - - discrete_flavor = channel.discrete_flavor - if discrete_flavor is not None: - discrete_flavor = ms.DiscreteFlow( - option_counts=[len(self.meta["flavors"])], - prefix=f"{prefix}.discrete_flow_flavor", - dims_with_prior=[0], - condition_dim=cond_dim, - subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], - subnet_layers=madnis_args["discrete_layers"], - subnet_activation=self.activation(madnis_args["discrete_activation"]), - ) - discrete_flavor.initialize_globals(self.process.contexts[0], seed) - - channels.append(Channel( - phasespace_mapping = channel.phasespace_mapping, - adaptive_mapping = flow, - discrete_sym = discrete_sym, - discrete_flavor = discrete_flavor, - channel_weight_indices = channel.channel_weight_indices, - name = channel.name, - active_flavors = channel.active_flavors, - )) - - return PhaseSpace( - mode="both", - channels=channels, - first_chan_weight_remap=phasespace.first_chan_weight_remap, - first_remapped_chan_count=phasespace.first_remapped_chan_count, - second_chan_weight_remap=phasespace.second_chan_weight_remap, - second_remapped_chan_count=phasespace.second_remapped_chan_count, - symfact=phasespace.symfact, - cwnet=self.build_cwnet(len(phasespace.symfact)), - prop_chan_weights=phasespace.prop_chan_weights, - subchan_weights=phasespace.subchan_weights, - ) - - def build_vegas(self, mapping: ms.PhaseSpaceMapping, prefix: str) -> ms.VegasMapping: - if not self.process.run_card["vegas"]["enable"]: - return None - - vegas = ms.VegasMapping( - mapping.random_dim(), - self.process.run_card["vegas"]["bins"], - prefix, - ) - for context in self.process.contexts: - vegas.initialize_globals(context) - return vegas - - def build_discrete( - self, permutation_count: int, flavor_count: int, prefix: str - ) -> tuple[ms.DiscreteSampler | None, ms.DiscreteSampler | None]: - is_adaptive = self.process.run_card["phasespace"]["adaptive_symmetry_sampling"] - if is_adaptive and permutation_count > 1: - discrete_sym = ms.DiscreteSampler( - [permutation_count], f"{prefix}.discrete_sym" - ) - for context in self.process.contexts: - discrete_sym.initialize_globals(context) - else: - discrete_sym = None - - if flavor_count > 1: - # the flavor probabilities are conditioned on the PDF prior; a - # leptonic process (or a decay) has none, and the Integrand then - # passes no condition: declaring one anyway fails to build - # ("keys and values must have the same size"), e.g. z > q q~ with - # both u- and d-type rows - discrete_flavor = ms.DiscreteSampler( - [flavor_count], f"{prefix}.discrete_flavor", - [] if self.process.leptonic else [0] - ) - for context in self.process.contexts: - discrete_flavor.initialize_globals(context) - else: - discrete_flavor = None - - return discrete_sym, discrete_flavor - - def build_cwnet(self, channel_count: int) -> ms.ChannelWeightNetwork: - #if channel_count == 1: - # return None - madnis_args = self.process.run_card["madnis"] - cwnet = ms.ChannelWeightNetwork( - channel_count=channel_count, - particle_count=self.particle_count, - hidden_dim=self.madnis_settings["cwnet_hidden_dim"], - layers=madnis_args["cwnet_layers"], - activation=self.activation(madnis_args["cwnet_activation"]), - prefix=f"subproc{self.subproc_id}.cwnet", - ) - cwnet.initialize_globals( - self.process.contexts[0], self.process.run_seed - ) - return cwnet - - def t_channel_mode(self, name: str) -> ms.PhaseSpaceMapping.TChannelMode: - modes = { - "propagator": ms.PhaseSpaceMapping.propagator, - "rambo": ms.PhaseSpaceMapping.rambo, - "chili": ms.PhaseSpaceMapping.chili, - } - if name in modes: - return modes[name] - else: - raise ValueError(f"Invalid t-channel mode '{name}'") - - def activation(self, name: str) -> ms.MLP.Activation: - activations = { - "relu": ms.MLP.relu, - "leaky_relu": ms.MLP.leaky_relu, - "elu": ms.MLP.elu, - "gelu": ms.MLP.gelu, - "sigmoid": ms.MLP.sigmoid, - "softplus": ms.MLP.softplus, - } - if name in activations: - return activations[name] - else: - raise ValueError(f"Invalid activation function '{name}'") - - def build_integrands( - self, - phasespace: PhaseSpace, - madnis_training: bool = False, - drop_cuts_and_rescale: bool = False - ) -> list[ms.Integrand]: - flavors = [] - flavor_remap = [] - flavor_factors = [] - flavor_mirror = [] - flavor_diff_xs_indices = [] - flavor_subproc_indices = [] - flavor_per_subproc_remap = [] - - for flav in self.meta["flavors"]: - if self.unmerged_meta is not None: - diff_xs_index = flav["subprocess"] - subproc_index = self.meta["subprocesses"][diff_xs_index] - ps_flavor = flav["flavor"] - flavor_diff_xs_indices.append(diff_xs_index) - flavor_subproc_indices.append(subproc_index) - flavor_per_subproc_remap.append(ps_flavor) - flav = self.unmerged_meta[subproc_index]["flavors"][ps_flavor] - flavors.append(flav["options"][0]) - flavor_remap.append(flav["index"]) - flavor_factors.append(len(flav["options"])) - flavor_mirror.append(flav["mirror"]) - - cross_sections = [] - self.step_matrix_elements = [] - for me_index, matrix_element in enumerate(self.matrix_elements): - if self.surrogates: - # the surrogate in the integrand, the full matrix element in - # the second pass, with the random numbers passed in - mat = ms.MatrixElement( - self.surrogates[me_index].index(), - self.particle_count, - [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, - ms.MatrixElement.flavor_in], - [ms.MatrixElement.matrix_element_out], - 1, - False, - self.process.me_frame, - self.process.incoming_count, - ) - self.step_matrix_elements.append(ms.MatrixElement( - matrix_element, - [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, - ms.MatrixElement.flavor_in, ms.MatrixElement.random_color_in, - ms.MatrixElement.random_helicity_in, - ms.MatrixElement.random_diagram_in], - [ms.MatrixElement.matrix_element_out, - ms.MatrixElement.color_index_out, - ms.MatrixElement.helicity_index_out, - ms.MatrixElement.diagram_index_out], - False, - self.process.me_frame, - self.process.incoming_count, - )) - elif matrix_element: - mat = ms.MatrixElement( - matrix_element, - ms.Integrand.matrix_element_inputs, - ms.Integrand.matrix_element_outputs, - True, - self.process.me_frame, - self.process.incoming_count, - ) - else: - #TODO: not working in merged mode - mat = ms.MatrixElement( - 0xBADCAFE, - self.particle_count, - ms.Integrand.matrix_element_inputs, - ms.Integrand.matrix_element_outputs, - self.meta["diagram_count"], - True, - self.process.me_frame, - self.process.incoming_count, - ) - pdf_grid = None if self.process.leptonic else self.process.pdf_grid - pdf_arg = None if self.process.leptonic else ms.CachedPdf() - cross_sections.append( - ms.DifferentialCrossSection( - matrix_element=mat, - # For a decay this is the decaying particle's mass, and the - # flux becomes 1/(2M): the result is a partial width in - # GeV, not a cross section in pb. - cm_energy=self.process.e_cm, - running_coupling=None, - energy_scale=ms.CachedScale(), - pid_options=[], - pdf1=pdf_arg, - pdf2=pdf_arg, - input_momentum_fraction=True, - decay=self.process.is_decay, - return_matrix_element=bool(self.surrogates), - ) - ) - partial_weights = self.process.systematics_enabled() - madnis_args = self.process.run_card["madnis"] - custom_bias = self.process.custom_function("bias", self.particle_count, flavors) - integrands = [] - for channel in phasespace.channels: - integrands.append(ms.Integrand( - channel.phasespace_mapping, - cross_sections, - channel.adaptive_mapping, - channel.discrete_sym, - channel.discrete_flavor, - flavors, - pdf_grid, - self.process.running_coupling, - self.scale, - phasespace.prop_chan_weights, - phasespace.subchan_weights, - phasespace.cwnet, - phasespace.first_chan_weight_remap, - phasespace.first_remapped_chan_count, - phasespace.second_chan_weight_remap, - phasespace.second_remapped_chan_count, - madnis_training, - drop_cuts_and_rescale, - partial_weights, - channel.channel_weight_indices, - channel.active_flavors, - flavor_remap, - flavor_factors, - flavor_mirror, - flavor_diff_xs_indices, - flavor_subproc_indices, - flavor_per_subproc_remap, - madnis_args["compressed_channel_weight_count"], - custom_bias, - )) - #print(integrands[1].function()) - #for i in integrands: print(i.function()) - return integrands - - -_ROOTED_OPTIONS = ('lhapdf', 'lhapdf_py3', 'lhapdf_py2', 'heptools_install_dir') - -def load_mg5_options(me_dir=None) -> dict: - """Read the tool paths from the MadGraph configuration, so the launcher - knows where LHAPDF and the optional programs (Pythia8/Delphes/MadSpin/ - reweight/analysis) live. - - Files are read least- to most-specific, each overriding the previous, which - is the same layering CommonRunCmd.set_configuration uses: the card inside - this process directory has the last word, then this installation, then the - per-user file. Relative values are resolved against the root of the file - they came from. - - ``me_dir`` is the process directory whose Cards/me5_configuration.txt has - the last word. It defaults to the working directory, which is right for - bin/generate_events (it chdirs there first) but not for an in-process - caller, which has not moved yet. - """ - - import madgraph - mg5dir = os.path.dirname(os.path.dirname(os.path.abspath(madgraph.__file__))) - me_dir = os.getcwd() if me_dir is None else me_dir - - options = { - 'pythia-pgs_path': None, 'pythia8_path': None, - 'madanalysis5_path': None, 'exrootanalysis_path': None, 'delphes_path': None, - 'rivet_path': None, 'contur_path': None, 'f2py_compiler': None, - 'lhapdf': None, 'lhapdf_py3': None, 'lhapdf_py2': None, 'timeout': 0, - 'heptools_install_dir': None, 'hepmc3_path': None, - } - config_files = [] - base_config = misc.base_config_file() - if base_config: - config_files.append((base_config, mg5dir)) - user_config = misc.user_config_file() - if user_config: - config_files.append((user_config, mg5dir)) - config_files.append((misc.install_config_file(mg5dir), mg5dir)) - config_files.append((os.path.join(me_dir, 'Cards', 'me5_configuration.txt'), - me_dir)) - for cfg, root in config_files: - if not os.path.exists(cfg): - continue - with open(cfg) as fsock: - for line in fsock: - line = line.split('#', 1)[0] - if '=' not in line: - continue - name, value = (x.strip() for x in line.split('=', 1)) - if name not in options or value in ('', 'None'): - continue - if (name.endswith('_path') or name in _ROOTED_OPTIONS) and \ - not os.path.isabs(value) and \ - (os.sep in value or value.startswith('.')): - value = os.path.normpath(os.path.join(root, value)) - options[name] = value - options['mg5_path'] = mg5dir # enables MadSpin/reweight - return options - - -_LHAPDF = None - -def lhapdf_paths(refresh: bool = False) -> misc.LhapdfPaths: - """The LHAPDF installation this run should use, resolved once. - - Uses the same helper as the 'launch' command, so driving a process - directory straight through bin/generate_events behaves identically to - launching it from MadGraph. - """ - - global _LHAPDF - if _LHAPDF is None or refresh: - _LHAPDF = misc.resolve_lhapdf(load_mg5_options()) - # Export what we found: the real LHAPDF library, systematics and the - # MadSpin/reweight subprocesses read these from the environment. - if _LHAPDF.data_paths: - os.environ["LHAPDF_DATA_PATH"] = os.pathsep.join(_LHAPDF.data_paths) - if _LHAPDF.config: - os.environ["MADGRAPH_LHAPDF_CONFIG"] = _LHAPDF.config - return _LHAPDF - - -class MG7Cmd(Cmd): - """The mg7 run interface: the command object that drives a launch. - - It plays two roles, and they are the same object on purpose. - - 1. It is the *mother* of the merged switch/card question (the thing - ``AskRunEditCard`` calls back into for ``keep_cards``/``do_open``/ - ``compute_widths``). - 2. It is the child command interface MG5 registers via - ``define_child_cmd_interface`` when you type ``launch`` on an mg7 output. - - Role 2 is what makes role 1 behave. ``Cmd.ask`` resolves a scripted answer - through ``self.check_answer_in_input_file``, where ``self`` is the mother -- - so once MG5 hands us its ``inputfile`` and registers itself as our - ``mother``, the question reads the very lines that follow ``launch`` in the - user's command file, and any line it does not recognise goes back up to MG5 - through ``store_line`` instead of being silently consumed. This is exactly - how a madevent output already behaves; before this class existed the mg7 - branch shelled out to ``bin/generate_events`` and had to guess which - following lines belonged to the run, which quietly swallowed every ``set``. - """ - - prompt = 'MG7> ' - - def __init__(self, me_dir='.', options=None, *args, **opts): - super().__init__(*args, **opts) - self.me_dir = os.path.abspath(me_dir) if me_dir != '.' else '.' - # The rest of the launcher is written against the process directory as - # the working directory, so do_generate_events chdirs there; but at - # construction time MG5 has not moved, hence the explicit me_dir. - self.options = load_mg5_options( - self.me_dir if self.me_dir != '.' else None) - if options: - # MG5's own resolved tool paths win: they are the ones the user - # configured in the session that is doing the launching. - self.options.update({k: v for k, v in options.items() - if k in self.options and v is not None}) - self.plugin_path = [] - self.proc_characteristics = {'grouped_matrix': False, 'limitations': []} - self._model = None # lazily imported by get_model() - - # ------------------------------------------------------------------ - # the contract the card question expects from its mother - # ------------------------------------------------------------------ - def keep_cards(self, need_card=[], ignore=[]): - from madgraph.interface.common_run_interface import CommonRunCmd - return CommonRunCmd.keep_cards(self, need_card, ignore) - - def do_open(self, line): - from madgraph.interface.common_run_interface import CommonRunCmd - CommonRunCmd.do_open(self, line) - - def check_open(self, args): - from madgraph.interface.common_run_interface import CommonRunCmd - CommonRunCmd.check_open(self, args) - - def get_model(self): - """The model this process was generated with (an ``import_ufo`` model). - - This is the hook the card question goes through for everything that - needs the model rather than the card: 'update dependent' (which resets - the masses/widths the model computes from the free parameters) and the - generic auto-width handler both call ``mother_interface.get_model()``. - Without it the launch question could only warn that it had failed to - update the dependent parameters, and MadSpin/the shower then read - whatever inconsistent values the param_card still held. - - None when the process recorded no model (an output made before - SubProcesses/model.txt existed) -- the caller warns in that case. - """ - if self._model is None: - me_dir = self.me_dir if self.me_dir != '.' else None - # False (not None) so a model that cannot be imported is not - # retried at every question. - self._model = load_process_model(me_dir) or False - return self._model or None - - def do_compute_widths(self, line): - # The interactive card editor delegates 'auto' width computation to - # the mother interface. Reuse the runtime helper (madgraph subprocess - # + the model stored at output time). ``line`` looks like - # " --path= [--nlo]"; we only need the card path. - m = re.search(r'--path=(\S+)', line or "") - path = m.group(1) if m else os.path.join("Cards", "param_card.dat") - compute_auto_widths(path) - # return an empty mapping: the caller iterates out.items() for the - # small-width treatment, which mg7 does not apply. - return {} - - # ------------------------------------------------------------------ - # running - # ------------------------------------------------------------------ - def do_generate_events(self, line): - """generate_events [-f] [-n NAME] [--laststep=parton] - - Run the mg7 generation and the selected post-processing tools.""" - - opts = self._parse_run_options(line) - cwd = os.getcwd() - target = self.me_dir if self.me_dir != '.' else cwd - try: - os.chdir(target) - _setup_logging() - if opts['name']: - self._set_run_name(opts['name']) - switch = {} - if not opts['force']: - switch = ask_edit_cards(mother=self) - switch = self._apply_laststep(switch, opts) - force_lhe_output_if_needed(switch) - _raise_open_file_limit() - run_generation(switch) - finally: - os.chdir(cwd) - - # ``launch`` is the name MG5 users type; keep it working here too. - do_launch = do_generate_events - - def do_npy_to_lhe(self, line): - """npy_to_lhe [RUN_NAME] [-o OUTPUT] [--seed N] [--no-gzip] - - Write the LHE file (events.lhe.gz) of a run that produced npy events, - e.g. to shower it. Without a run name, the most recent such run.""" - import shlex - me_dir = self.me_dir if self.me_dir != '.' else os.getcwd() - try: - npy_to_lhe.main(shlex.split(line), me_dir=me_dir) - except SystemExit as error: - if error.code not in (None, 0): - logger.error(str(error.code)) - - def do_quit(self, line): - """Leave the mg7 run interface.""" - return super().do_quit(line) - - do_exit = do_quit - - # ------------------------------------------------------------------ - # helpers - # ------------------------------------------------------------------ - def _parse_run_options(self, line): - """Understand the subset of MG5's `launch` options that mean something - for an mg7 run, and say so out loud for the ones that do not.""" - out = {'force': False, 'name': '', 'laststep': ''} - args = self.split_arg(line) - index = 0 - while index < len(args): - arg = args[index] - index += 1 - if arg in ('-f', '--force'): - out['force'] = True - elif arg in ('-n', '--name') and index < len(args): - out['name'] = args[index] - index += 1 - elif arg.startswith(('-n=', '--name=')): - out['name'] = arg.split('=', 1)[1] - elif arg in ('-s', '--laststep') and index < len(args): - out['laststep'] = args[index] - index += 1 - elif arg.startswith(('-s=', '--laststep=')): - out['laststep'] = arg.split('=', 1)[1] - elif arg.startswith('-'): - logger.warning( - "'%s' is not supported by the mg7 output and is ignored; " - "the equivalent settings live in Cards/run_card.toml", arg) - elif arg and not out['name']: - out['name'] = arg - return out - - def _set_run_name(self, name): - """Honour `launch -n NAME`: the mg7 run directory is - Events/_NN, so the name is a run_card setting.""" - path = os.path.join("Cards", "run_card.toml") - run_card = RunCardMG7(path, consistency=False) - if run_card["run"]["run_name"] != name: - run_card["run"]["run_name"] = name - run_card.write(path) - logger.info("run name set to '%s'", name) - - def _apply_laststep(self, switch, opts): - """`--laststep=parton` means: generate the events, run nothing on them. - - The mg7 tool selection is the switch dict the question returns, so - stopping at parton level is switching those entries off. - """ - if opts['laststep'] not in ('parton', 'auto', ''): - logger.warning( - "--laststep=%s is not supported by the mg7 output; " - "select the programs to run in the launch question instead", - opts['laststep']) - return switch - if opts['laststep'] != 'parton' or not switch: - return switch - switch = dict(switch) - for key in ('shower', 'detector', 'analysis', 'madspin', 'reweight'): - if not _off(switch.get(key, 'OFF')): - logger.info("--laststep=parton: not running %s", key) - switch[key] = 'OFF' - return switch - - -def build_selector_cmd(mother=None): - """Build the (monkey-patched) mother command + merged switch/card selector - used by the mg7 output. Returns the selector *class* and a mother instance - understood by AskRun/AskforEditCard. - - ``mother`` lets an in-process caller supply its own :class:`MG7Cmd` -- the - one MG5 registered as its child -- so that the question resolves scripted - answers against MG5's own command file (see :class:`MG7Cmd`). When it is - omitted a standalone instance is built, which is what ``bin/generate_events`` - gets. - """ - - from madgraph.interface.madevent_interface import AskRunEditCard - - from madgraph.various import banner as _banner_mod - from madgraph.various import misc as _misc - - class MG7Selector(AskRunEditCard): - """Merged switch/card question for the mg7 output. - - The mg7 run_card is a TOML file, so define_paths/init_run/do_set are - overridden *on this class* (not globally on AskforEditCard) to treat it - as such. Scoping them here is important: MadSpin/reweight spawn their own - legacy madevent runs in the same process, and a global patch would break - their (run_card.dat based) card editing.""" - - # param_card + the TOML run_card are always offered - always_cards = ['param_card.dat', 'run_card.toml'] - optional_cards = [] - - def define_paths(self, **opt): - super().define_paths(**opt) - self.paths["run"] = os.path.join(self.me_dir, "Cards", "run_card.toml") - self.paths["run_card.toml"] = os.path.join(self.me_dir, "Cards", "run_card.toml") - # the TOML run_card uses its own default file (concrete defaults - # written at output time); this powers "set default". - self.paths["run_default"] = os.path.join(self.me_dir, "Cards", "run_card_default.toml") - - def init_run(self, cards): - # Make sure the run_card is loaded as a RunCardMG7 (the generic - # editor does not recognise run_card.toml), else "set " fails. - out = super().init_run(cards) - if not isinstance(getattr(self, "run_card", None), RunCardMG7): - toml_path = self.paths.get("run") or os.path.join( - self.me_dir, "Cards", "run_card.toml") - if os.path.exists(toml_path): - try: - # allow_scan so a run_card holding scan:[...] values loads - with _misc.TMP_variable(_banner_mod.RunCard, "allow_scan", True): - self.run_card = RunCardMG7(toml_path, consistency="warning") - self.run_set = list(self.run_card.keys()) - except Exception as err: - logger.warning("could not load %s: %s", toml_path, err) - if isinstance(getattr(self, "run_card", None), RunCardMG7): - self.run_card.allow_scan = True - return getattr(self, "run_set", out) - - def do_set(self, line, *args, **kwargs): - # madevent-style shortcuts (lhc/lep/fixed_scale/no_parton_cut), cut - # editing, energy units and arithmetic/mass expressions on the TOML - # run_card, intercepted before delegating to the generic editor. - targs = self.split_arg(line) - run_card = getattr(self, "run_card", None) - if isinstance(run_card, RunCardMG7) and targs: - start = 1 if targs[0] == "run_card" else 0 - if len(targs) > start: - name = targs[start] - nlow = name.lower() - rest = " ".join(targs[start + 1:]).split("#")[0].strip() - masses = run_card.get_mass_shortcuts(getattr(self, "param_card", None)) - - if nlow in ("no_parton_cut", "nocut", "no_cut"): - run_card.remove_all_cut() - logger.info("removing all cuts from the run_card.toml") - self.modified_card.add("run") - return - # [histograms] is a list of observables, not a parameter, - # so "set histograms ..." can only mean the whole section. - if nlow in ("histograms", "histogram", "plots"): - value = rest.lower() - if value in ("off", "none", "no", "false", "0"): - run_card.remove_all_histograms() - logger.info( - "removing all histograms from the run_card.toml") - self.modified_card.add("run") - return - if value in ("default", "on"): - if run_card.restore_default_histograms( - self.paths.get("run_default")): - logger.info( - "restored the %d histograms written at " - "output time", - len(run_card["histograms"])) - self.modified_card.add("run") - else: - logger.warning( - "no run_card_default.toml to restore the " - "histograms from") - return - logger.warning( - "'set histograms %s': only OFF (remove them all) " - "and default (put the generated ones back) are " - "understood; edit the [histograms] section of " - "Cards/run_card.toml for anything else", - rest or "(no value)") - return - if nlow in ("lhc", "lep", "ilc", "lcc") and rest: - ecm = run_card.set_collider(nlow, rest, masses) - logger.info("set %s collider: e_cm = %s GeV", nlow, ecm) - self.modified_card.add("run") - return - if nlow == "fixed_scale" and rest: - val = run_card.set_fixed_scale(rest, masses) - logger.info("set fixed scales to %s GeV", val) - self.modified_card.add("run") - return - - if rest and run_card.is_cut_name(name): - cut, bound, val = run_card.set_cut(name, run_card.evaluate(rest, masses)) - logger.info("modify cut %s.%s of the run_card.toml to %s", cut, bound, val) - self.modified_card.add("run") - return - - # 'iseed' is the madevent spelling of the mg7 run.seed, - # but the two disagree on how to ask for a random seed: - # madevent uses 0, mg7 uses -1 (0 being a perfectly good - # fixed seed there). Translate here rather than in - # _LO_SCALAR_MAP, which also drives the LO -> MG7 run_card - # conversion and would silently pin those runs to seed 0. - if rest and nlow == 'iseed': - value = run_card.evaluate(rest, masses) - try: - value = int(value) - except (TypeError, ValueError): - value = None - if value is None: - logger.warning("ignoring 'set iseed %s': not an integer", rest) - return - seed = -1 if value == 0 else value - run_card.set('run.seed', seed, user=True) - logger.info("set iseed (mg7 run.seed) of the run_card.toml to %s%s", - seed, " (random)" if seed == -1 else "") - self.modified_card.add("run") - return - - # legacy madevent run_card names -> mg7 "section.key", so a - # madevent-style launch script ("set nevents 500", "set - # use_syst F", "set bwcutoff 10", ...) edits the mg7 - # run_card.toml verbatim. RunCardMG7._LO_SCALAR_MAP is the - # same rename table used by the LO->MG7 run_card conversion. - if rest and nlow in run_card._LO_SCALAR_MAP \ - and nlow not in [k.lower() for k in run_card.keys()]: - target = run_card._LO_SCALAR_MAP[nlow] - run_card.set(target, rest, user=True) - logger.info("set %s (mg7 %s) of the run_card.toml to %s", - nlow, target, run_card[target]) - self.modified_card.add("run") - return - - if rest and nlow in [k.lower() for k in run_card.keys()]: - current = run_card[nlow] - if isinstance(current, (int, float)) and not isinstance(current, bool): - resolved = run_card.evaluate(rest, masses) - if not isinstance(resolved, str): - prefix = "run_card " if start == 1 else "" - line = "%s%s %s" % (prefix, name, resolved) - return super().do_set(line, *args, **kwargs) - - return MG7Selector, (mother if mother is not None else MG7Cmd()) - - -def ask_edit_cards(mother=None) -> dict: - """Single (MadDM-style) question letting the user both pick which programs - to run after generation and edit the associated cards. Returns the switch - dict describing the selected tools. - - ``mother`` is the :class:`MG7Cmd` asking the question. Passing the instance - MG5 registered as its child is what lets a scripted ``launch`` answer this - question from MG5's own command file. - """ - - selector_class, mother = build_selector_cmd(mother) - # path_msg is what makes Cmd.check_answer_in_input_file accept a bare path as - # an answer (its "elif path:" branch), so that a scripted launch can hand the - # question a card/banner path -- as the question itself advertises -- and have - # it replace the corresponding card. Without it the path is rejected ("This - # answer is not valid for current question") and the default is used instead. - try: - switch, question = mother.ask('', '0', [], path_msg='enter path', - ask_class=selector_class, - mode='auto', line_args=[], force=False, - return_instance=True) - except BaseException: - # interrupted (e.g. ctrl-C) before the prune below could run: prune - # here too, using the instance ask() stashed pre-interrupt, or the - # eagerly-materialised tool cards get misread as "left ON" next time. - instance = getattr(mother, '_last_ask_instance', None) - if instance is not None and hasattr(instance, 'active_cards'): - try: - mother.keep_cards(instance.active_cards()) - except Exception as error: - logger.debug('could not prune tool cards after aborted launch question: %s', error) - raise - switch = dict(switch) - prune_unselected_tool_cards(mother, question, switch) - return switch - - -def prune_unselected_tool_cards(mother, question, switch) -> None: - """Hide the cards of the tools that were NOT selected in the question. - - The question materialises every candidate tool card (copying the - *_default.dat) so that it can offer them all for edition; madevent's - ask_run_configuration prunes them again afterwards (keep_cards on the - selected cards), but the mg7 launcher had no such cleanup. The materialised - cards therefore persisted, and a later re-launch of the same output saw them - all present and defaulted every switch to ON (set_default_ keys off - card presence). Prune here too, using the question's own switch->card map, so - a re-launch reflects the previous selection instead of turning everything on. - """ - try: - keep = list(question.always_cards) - for spec in question.switch_cards: - if spec['on'](switch): - keep.append(spec['card']) - mother.keep_cards(keep) - except Exception as error: - logger.debug('could not prune unselected tool cards: %s', error) - - -def _find_event_file(run_path): - """Locate the LHE event file produced by generate_events (if any).""" - for name in ("events.lhe", "events.lhe.gz"): - path = os.path.join(run_path, name) - if os.path.exists(path): - return path - return None - - -def _find_or_convert_event_file(run_path, log): - """The run's LHE event file, converting its npy events when that is all - the run wrote (see npy_to_lhe).""" - lhe_path = _find_event_file(run_path) - if lhe_path is None and npy_to_lhe.can_convert(run_path): - log.info("No LHE event file in %s: converting the npy events.", run_path) - try: - lhe_path = npy_to_lhe.convert(run_path) - except Exception as error: - _report_failure(log, "npy->LHE conversion", error, run_path) - return lhe_path - - -def _report_failure(log, what, error, directory=None): - """Log a post-processing failure and write the full traceback to a file - (whose path is printed) so the problem can be investigated.""" - import traceback - if not directory or not os.path.isdir(directory): - directory = os.getcwd() - path = os.path.join(directory, "%s_crash.log" % what.replace(" ", "_")) - try: - with open(path, "w") as fsock: - fsock.write(traceback.format_exc()) - except Exception: - path = None - log.warning("%s failed: %s", what, error) - if path: - log.warning("full traceback written to: %s", path) - - -def _off(value) -> bool: - """True when a switch value means the tool was not selected.""" - return value in (None, "OFF", "Not Avail.", "Not Avail. (numpy missing)") - - -def _has_effective_handler(lg) -> bool: - """True when ``lg`` (or an ancestor it propagates to) already has a handler, - i.e. some other party is already displaying these records.""" - while lg: - if lg.handlers: - return True - if not lg.propagate: - return False - lg = lg.parent - return False - - -_TOOL_LOGGING_READY = False - - -def _setup_logging(): - """Make the reused madevent tool drivers' progress visible on screen. - - The drivers already narrate what they do through logger.info / - update_status ("Running Pythia8 [arXiv:...]", "Splitting .lhe event - file...", the live Idle/Running/Completed job counters, "Running - MadSpin", ...), but the standalone mg7 launcher never attaches a handler - to the 'madgraph'/'madevent' loggers, so all of it is swallowed. Attach a - colored INFO console handler to them (idempotent).""" - global _TOOL_LOGGING_READY - if _TOOL_LOGGING_READY: - return - if _has_effective_handler(logging.getLogger("madgraph")): - # Someone already configured logging for us -- in practice MG5, which - # now runs the launch in its own process. Its configuration is - # authoritative (it owns stdout_level and the tutorial logger); adding - # a second handler here would print every line twice. - _TOOL_LOGGING_READY = True - return - try: - import madgraph.interface.coloring_logging # registers ColorFormatter - formatter = logging.ColorFormatter("%(message)s") - except Exception: - formatter = logging.Formatter("%(message)s") - handler = logging.StreamHandler(sys.stdout) - handler.setFormatter(formatter) - handler._mg7_tool_handler = True - for name in ("madgraph", "madevent", "cmdprint", "madgraph7"): - lg = logging.getLogger(name) - if not any(getattr(h, "_mg7_tool_handler", False) for h in lg.handlers): - lg.addHandler(handler) - lg.setLevel(logging.INFO) - lg.propagate = False - _TOOL_LOGGING_READY = True - - -def selected_tools(switch) -> list: - """The post-processing drivers run_selected_tools runs for this switch.""" - # MadAnalysis5 hadron level analyses the shower/detector output, so it only - # makes sense when a shower ran (mirrors madevent's card gating: - # analysis == 'MadAnalysis5' and shower != 'OFF'). - ma5 = switch.get("analysis") == "MadAnalysis5" - showered = not _off(switch.get("shower")) - return [t for t, on in ( - ("reweighting", not _off(switch.get("reweight"))), - ("MadSpin", not _off(switch.get("madspin"))), - ("MadAnalysis5 (parton level)", ma5), - ("Pythia8 shower", switch.get("shower") == "Pythia8"), - ("Delphes", switch.get("detector") == "Delphes"), - ("MadAnalysis5 (hadron level)", ma5 and showered), - ("Rivet", switch.get("analysis") == "Rivet"), - ) if on] - - -def run_selected_tools(switch, process) -> None: - """Run the optional post-processing programs selected in the merged - question on the generated events. - - All the tools are driven through :class:`MG7RunCmd`, a thin adapter around - the standard madevent run interface (``MadEventCmd``). This means the mg7 - output reuses the *exact* madevent tool drivers -- in particular Pythia8 is - showered in parallel over the cluster/multicore backend, identically to a - madevent run -- instead of re-implementing each tool here. The command - sequence mirrors ``MadEventCmd.do_launch`` (reweight, MadSpin, MA5 parton, - shower, Delphes, MA5 hadron, Rivet); each tool checks its own card, so a - tool is invoked only when its switch is on. - """ - log = logging.getLogger("madevent") - - tools = selected_tools(switch) - if not tools: - # A switch that is not off does not necessarily select a driver here: - # "Not Avail." is not off, analysis = ExRoot has no mg7 driver, and a - # shower switch set to something other than Pythia8 selects nothing - # either. Without this a plain generate/output/launch announced a - # post-processing step with nothing after the colon, and paid for - # building the run interface (or an npy->LHE conversion) to do nothing. - return - - lhe_path = _find_or_convert_event_file(process.run_path, log) - if lhe_path is None: - log.warning("No LHE event file in %s; cannot run %s.", - process.run_path, ", ".join(tools)) - return - - run_name = os.path.basename(os.path.dirname(os.path.abspath(lhe_path))) - run_dir = os.path.dirname(os.path.abspath(lhe_path)) - - ma5 = switch.get("analysis") == "MadAnalysis5" - showered = not _off(switch.get("shower")) - - log.info("") - log.info("Post-processing the generated events with: %s", ", ".join(tools)) - - try: - from madgraph.iolibs.template_files.mg7.run_interface import MG7RunCmd - cmd = MG7RunCmd(os.getcwd(), load_mg5_options(), run_name, lhe_path) - cmd.set_run_name(run_name, None, "parton") - except Exception as error: - _report_failure(log, "post-processing setup", error, run_dir) - return - - def run(tool, command): - bar = "=" * 60 - log.info("") - log.info(bar) - log.info(" post-processing step: %s", tool) - log.info(bar) - start = time.time() - try: - cmd.exec_cmd(command, postcmd=False, printcmd=False) - except Exception as error: - _report_failure(log, tool, error, run_dir) - else: - log.info(" -> %s done (%.1fs)", tool, time.time() - start) - - # Same order as MadEventCmd.do_launch. The commands take no run name and so - # act on cmd.run_name: reweight runs on the generated events, then - # decay_events (MadSpin) repoints cmd.run_name to the "_decayed_i" run, - # so the shower and the analyses that follow act on the decayed events -- - # exactly as a madevent run does. Each driver still guards on its own card. - if not _off(switch.get("reweight")): - run("reweight", "reweight -from_cards") - if not _off(switch.get("madspin")): - run("MadSpin", "decay_events -from_cards") - if ma5: - run("MadAnalysis5 (parton)", "madanalysis5_parton --no_default") - if switch.get("shower") == "Pythia8": - run("Pythia8 shower", "shower --no_default") - if switch.get("detector") == "Delphes": - run("Delphes", "delphes --no_default") - # hadron-level MA5 needs the shower/detector output: only run it if a shower - # was requested (otherwise there is no hadron-level event file to analyse). - if ma5 and showered: - run("MadAnalysis5 (hadron)", "madanalysis5_hadron --no_default") - if switch.get("analysis") == "Rivet": - run("Rivet", "rivet --no_default") - - # Finalize the run exactly as MadEventCmd.do_launch does after the - # shower/detector/analysis tools: store_result() processes the deferred - # "to_store" actions -- in particular it gzips the Pythia8 HepMC to - # _pythia8_events.hepmc.gz (the pythia8_card default is HEPMCoutput:file - # = hepmc.gz). The deferred Rivet job reads exactly that .gz path, so this - # must run before the Rivet/Contur postprocessor below. - try: - cmd.store_result() - except Exception as error: - _report_failure(log, "store_result", error, run_dir) - - if switch.get("analysis") == "Rivet": - # do_rivet only *prepared* the run (wrote Events//run_rivet.sh) and - # deferred execution to the postprocessor -- the rivet_card default has - # run_rivet_later = True, so with --no_default it logged "Skipping Rivet - # for now, passing it to postprocessor" and appended the run to - # cmd.postprocessing_dirs. MadEventCmd.do_launch runs that postprocessor - # at the end of a run; do the same here so the .yoda (and, when - # run_contur = True, the Contur limits) are actually produced. - bar = "=" * 60 - log.info("") - log.info(bar) - log.info(" post-processing step: Rivet/Contur postprocessing") - log.info(bar) - start = time.time() - try: - cmd.postprocessing() - except Exception as error: - _report_failure(log, "Rivet-Contur postprocessing", error, run_dir) - else: - log.info(" -> Rivet/Contur postprocessing done (%.1fs)", - time.time() - start) - - -def _add_time_of_flight(lhe_path, threshold, param_card_path, log): - """Add invariant-lifetime (vtim) information to the LHE events, drawing each - unstable particle's decay length from its width (in mm). Mirrors - common_run_interface.CommonRunCmd.do_add_time_of_flight but reads the widths - from the local param_card instead of the LHE banner.""" - import random - try: - import madgraph.various.lhe_parser as lhe_parser - except ImportError: - import internal.lhe_parser as lhe_parser - from models import check_param_card as param_card_mod - from madgraph.various import misc as _misc - import madgraph.iolibs.files as files - - need_zip = lhe_path.endswith('.gz') - if need_zip: - _misc.gunzip(lhe_path) - lhe_path = lhe_path[:-3] - - param_card = param_card_mod.ParamCard(param_card_path) - cst = 6.58211915e-25 # hbar in GeV s - c = 299792458000 # speed of light in mm/s - log.info('Adding time of flight information on %s', lhe_path) - lhe = lhe_parser.EventFile(lhe_path) - out = open('%s_2vertex.lhe' % lhe_path, 'w') - out.write(lhe.banner) - for event in lhe: - for particle in event: - # default=0 -> particles without a decay entry are treated as stable - width = param_card['decay'].get((abs(particle.pid),), 0.).value - if width: - vtim = c * random.expovariate(width / cst) - if vtim > threshold: - particle.vtim = vtim - out.write(str(event)) - out.write('\n') - out.close() - lhe.close() - files.mv('%s_2vertex.lhe' % lhe_path, lhe_path) - if need_zip: - _misc.gzip(lhe_path) - - -def _lhapdf_config_path(): - """Path to lhapdf-config so systematics can import the python lhapdf - module (required to compute PDF/scale variations).""" - return lhapdf_paths().config - - -def _run_systematics(lhe_path, cfg, log): - """Run systematics.py (scale/PDF variations) on the LHE file, reusing the - exact worker madevent uses (systematics.call_systematics).""" - try: - import madgraph.various.systematics as systematics - except ImportError: - import internal.systematics as systematics - - def fmt(vals): - return ','.join(str(v) for v in vals) - - opts = [] - if cfg.get('systematics_mur'): - opts.append('--mur=%s' % fmt(cfg['systematics_mur'])) - if cfg.get('systematics_muf'): - opts.append('--muf=%s' % fmt(cfg['systematics_muf'])) - if cfg.get('systematics_pdf'): - opts.append('--pdf=%s' % fmt(cfg['systematics_pdf'])) - extra = cfg.get('systematics_str_options', '') - if extra: - opts.extend(extra.split()) - # The mg7 LHE has no block, so systematics cannot read the per-event - # LO reweighting info. For a process with a single alpha_s power (stored in - # SubProcesses/proc_characteristics at output time) it can be reconstructed - # from the events: pass that power as --lo_nqcd. -1 means the QCD power is - # not uniform, so the reconstruction is not applicable. - if not any(o.startswith('--lo_nqcd') for o in opts): - try: - from madgraph.various import banner as _banner_mod - except ImportError: - import internal.banner as _banner_mod - pc_path = os.path.join('SubProcesses', 'proc_characteristics') - if os.path.exists(pc_path): - try: - nqcd = int(_banner_mod.ProcCharacteristic(pc_path)['single_qcd_order']) - except Exception: - nqcd = -1 - if nqcd >= 0: - opts.append('--lo_nqcd=%d' % nqcd) - # tell systematics where to find lhapdf (so it can link the python module) - if not any(o.startswith('--lhapdf_config') for o in opts): - lhapdf_config = _lhapdf_config_path() - if lhapdf_config: - opts.append('--lhapdf_config=%s' % lhapdf_config) - - log.info('Running systematics on %s %s', lhe_path, ' '.join(opts)) - systematics.call_systematics([lhe_path, lhe_path] + opts, - log=lambda x: log.info(str(x))) - - -def run_lhe_postprocessing(process) -> None: - """Run the LHE-level post-processings configured in the run_card - [postprocessing] section on the generated event file (displaced-vertex - time-of-flight and systematics). Only applies when an LHE file exists.""" - lhe_path = _find_event_file(process.run_path) - if lhe_path is None: - return - # process.run_card is a RunCardMG7; the section view exposes .get(key, def) - try: - cfg = process.run_card["postprocessing"] - except Exception: - return - log = logging.getLogger('madevent') - - if cfg.get('systematics') and getattr(process, 'is_decay', False): - # Scale and PDF variations are a beam quantity. A decay has neither, so - # systematics can only fail here ("not supported for pdlabel=none") -- - # and it is not free: MadSpin reruns the launcher for every pool refill. - log.info("decay mode: skipping systematics (no beams to vary)") - elif cfg.get('systematics'): - native = getattr(process, 'systematics', None) - if native is not None and native.weight_count > 0: - log.warning("[postprocessing] systematics is ignored: the scale/PDF " - "weights were already computed by madspace ([systematics] " - "section); set systematics.enable = false to use the legacy " - "systematics.py path instead.") - else: - try: - _run_systematics(lhe_path, cfg, log) - except Exception as error: - _report_failure(log, "systematics computation", error, - os.path.dirname(lhe_path)) - - tof = cfg.get('time_of_flight', -1.0) - try: - if tof is not None and float(tof) >= 0: - _add_time_of_flight(lhe_path, float(tof), process.param_card_path, log) - except Exception as error: - _report_failure(log, "add_time_of_flight", error, - os.path.dirname(lhe_path)) - - -_model_hash_warned = set() -_model_cache = {} - - -def read_stored_model(me_dir=None) -> "str | None": - """The model this process was generated with, as recorded at output time in - ``SubProcesses/model.txt``: a UFO path carrying its restriction (e.g. - ``.../models/sm-no_b_mass``) or a model name -- in either case something - ``import model`` understands. Returns None when the process was written - without that information. - - Warns (once per directory) when the model on disk no longer matches the - hash stored on the second line, i.e. it changed since the output was made. - """ - base = me_dir or os.getcwd() - model_file = os.path.join(base, "SubProcesses", "model.txt") - if not os.path.exists(model_file): - return None - with open(model_file) as f: - lines = f.read().splitlines() - model = lines[0].strip() if lines else "" - stored_hash = lines[1].strip() if len(lines) > 1 else "" - if not model: - logger.warning("SubProcesses/model.txt is empty; the model of this " - "process is unknown.") - return None - - # verify the model on disk still matches the one used at output time. The - # reference may carry a restriction suffix ('sm-no_b_mass'); the hash was - # taken on the UFO directory itself ('sm'). - model_dir = model if os.path.isdir(model) else model.rsplit('-', 1)[0] - if stored_hash and os.path.isdir(model_dir) \ - and model_file not in _model_hash_warned: - current_hash = misc.hash_model_files(model_dir) - if current_hash and current_hash != stored_hash: - _model_hash_warned.add(model_file) - logger.warning( - "The model at %s has changed since this process was generated " - "(hash mismatch); anything computed from it now (the 'auto' " - "widths, the dependent masses/widths of the param_card) may be " - "inconsistent with the matrix element.", model_dir) - return model - - -def _proc_characteristic(me_dir=None) -> dict: - """SubProcesses/proc_characteristics as a plain dict (empty when the file is - missing or unreadable). Plain, because ProcCharacteristic is a ConfigFile, - whose get() is the parameter accessor and takes no default.""" - from madgraph.various import banner as _banner_mod - path = os.path.join(me_dir or os.getcwd(), 'SubProcesses', - 'proc_characteristics') - if os.path.exists(path): - try: - return dict(_banner_mod.ProcCharacteristic(path)) - except Exception as error: - logger.debug("could not read %s: %s", path, error) - return {} - - -def load_process_model(me_dir=None): - """Import the UFO model this process was generated with, as - :func:`models.import_ufo.import_model` builds it -- the object every - model-level consumer of the param_card expects (see - :meth:`MG7Cmd.get_model`). None when the process recorded no model, or when - that model no longer imports. - - Cached per (reference, complex-mass-scheme): importing a model is slow and - the card question asks for it again at every 'update dependent'.""" - ref = read_stored_model(me_dir) - if not ref: - return None - cms = _proc_characteristic(me_dir).get('complex_mass_scheme', False) - key = (ref, bool(cms)) - if key not in _model_cache: - import models.import_ufo as import_ufo - try: - with misc.MuteLogger(['madgraph.model'], [50]): - _model_cache[key] = import_ufo.import_model( - ref, complex_mass_scheme=bool(cms)) - except Exception as error: - # 'update dependent' calls this under a SIGALRM whose handler - # raises a TimeOutError of its own (a class defined inside the - # caller, so it can only be recognised by name). Let it through: - # the caller then says the model took too long to load and how to - # force it, instead of reporting a model that does not import. - if type(error).__name__ == 'TimeOutError': - raise - logger.debug("could not import the model %s: %s", ref, error) - _model_cache[key] = None - return _model_cache[key] - - -def compute_auto_widths(param_card_path=os.path.join("Cards", "param_card.dat")) -> None: - """Fill any width set to ``auto`` in the param_card, using madgraph and the - model stored at output time (``SubProcesses/model.txt``), and write the - result back into the card. A no-op when the card has no ``auto`` width. - - Called before every generation, so a single run and each scan point (whose - param_card is rewritten just before ``run_single``) both get model-computed - widths.""" - try: - with open(param_card_path) as f: - text = f.read() - except OSError: - return - # matches "DECAY auto" (optionally "auto@NLO"), as in - # common_run_interface.static_check_param_card - pdgs = re.findall(r"(?im)^\s*decay\s+([+-]?\d+)\s+auto", text) - if not pdgs: - return - pdgs = list(dict.fromkeys(pdgs)) # de-duplicate, keep order - - model = read_stored_model() - if not model: - logger.warning( - "The param_card requests 'auto' width(s) for %s but the model of " - "this process is unknown; leaving them as-is.", " ".join(pdgs)) - return - - mg5 = str(_MG_ROOT / "bin" / "madgraph") - if not os.path.exists(mg5): - logger.warning("Cannot find madgraph at %s; 'auto' widths not computed.", mg5) - return - - import tempfile - cmds = "import model %s\ncompute_widths %s --path=%s\n" % ( - model, " ".join(pdgs), os.path.abspath(param_card_path)) - with tempfile.NamedTemporaryFile("w", suffix=".mg5", delete=False) as fh: - fh.write(cmds) - cmdfile = fh.name - try: - logger.info("Computing 'auto' width(s) for %s ...", " ".join(pdgs)) - proc = subprocess.run([mg5, cmdfile]) - if proc.returncode != 0: - logger.warning( - "compute_widths returned a non-zero exit code; the param_card " - "may still contain 'auto' entries.") - finally: - try: - os.remove(cmdfile) - except OSError: - pass - - -def _release_channel_generators(process) -> None: - """Drop the per-channel event/weight generators so their intermediate - .npy files are deleted right away. madspace only frees them when the - owning Python objects are collected, and process.event_generator / - process.phasespaces sit in a reference cycle that plain refcounting - never breaks -- only an explicit gc.collect() does, promptly, instead of - leaving it to whenever the cyclic GC next runs on its own.""" - process.event_generator = None - process.phasespaces = None - gc.collect() - - -def run_single(switch=None) -> "MadgraphProcess": - """Run a single generation and return the process (for its result).""" - compute_auto_widths() - process = MadgraphProcess() - process.survey() - process.train_madnis() - process.generate_events() - if process.gridpack_run_mode != "regular": - return process - # run_card-driven LHE post-processing (displaced vertex + systematics) - run_lhe_postprocessing(process) - # run the post-processing tools (Pythia8/Delphes/MadSpin/reweight/analysis) - # selected in the merged question above on the generated events. - if switch: - run_selected_tools(switch, process) - return process - - -def detect_run_scan(run_card_path): - """Return a banner.RunCardIterator if the run_card contains scan:[...] - values, else None.""" - from madgraph.various import banner as banner_mod - from madgraph.various import misc as _misc - with _misc.TMP_variable(banner_mod.RunCard, 'allow_scan', True): - rc = banner_mod.RunCard(run_card_path, consistency=False) - if getattr(rc, 'scan_set', None): - return banner_mod.RunCardIterator(run_card_path) - return None - - -def detect_param_scan(param_card_path): - """Return a ParamCardIterator if the param_card contains scan:[...] values, - else None.""" - if not os.path.exists(param_card_path): - return None - from models import check_param_card as param_card_mod - it = param_card_mod.ParamCardIterator(param_card_path) - for block in it.order: - for param in block: - if isinstance(param.value, str) and param.value.strip().lower().startswith('scan'): - return it - return None - - -def run_scan(iterator, card_path, switch=None) -> None: - """Iterate over all scan points, running a full generation for each and - accumulating the results, then write the scan summary. Works for both the - run_card (RunCardIterator) and the param_card (ParamCardIterator); their - interface (__iter__/write/store_entry/get_next_name/write_summary) is the - same. The scan card is restored afterwards.""" - import tomllib - with open(os.path.join("Cards", "run_card.toml"), "rb") as f: - run_name = tomllib.load(f).get("run", {}).get("run_name", "run") - - from models import check_param_card as param_card_mod - is_param_scan = isinstance(iterator, param_card_mod.ParamCardIterator) - - backup = card_path + ".scan_bak" - shutil.copy(card_path, backup) - try: - for i, point in enumerate(iterator): - point.write(card_path) - logger.info("=== scan point %d ===", i + 1) - process = run_single(switch) - # use the run directory the process actually created, so the - # per-point params.dat written by write_summary has a home - name = os.path.basename(process.run_path) - if is_param_scan: - # pass the (possibly auto-width-updated) card so the summary - # records the model-computed width for each 'auto' entry - iterator.store_entry(name, process.get_result(), - param_card_path=card_path) - else: - iterator.store_entry(name, process.get_result()) - _release_channel_generators(process) - os.makedirs("Events", exist_ok=True) - summary = os.path.join("Events", "scan_%s.txt" % run_name) - iterator.write_summary(summary) - logger.info("scan results written to %s", summary) - finally: - shutil.move(backup, card_path) - - -def run_generation(switch=None) -> None: - """Run the generation, expanding a scan over the run_card or the param_card - when one is present (scanning both simultaneously is not allowed).""" - run_card_path = os.path.join("Cards", "run_card.toml") - param_card_path = os.path.join("Cards", "param_card.dat") - run_iter = detect_run_scan(run_card_path) - param_iter = detect_param_scan(param_card_path) - if run_iter and param_iter: - raise RuntimeError( - "Scanning simultaneously over the run_card and the param_card is " - "not allowed. Please keep the scan:[...] entries in only one card.") - if run_iter: - run_scan(run_iter, run_card_path, switch) - elif param_iter: - run_scan(param_iter, param_card_path, switch) - else: - _release_channel_generators(run_single(switch)) - - -def _lhe_needed_by(switch, card) -> "str | None": - """Why the run needs an LHE event file, or None when it does not: every - post-processing tool run_selected_tools drives and the run_card - [postprocessing] steps read the LHE file. ``card`` is the raw TOML content - of the run_card.""" - tools = selected_tools(switch) if switch else [] - if tools: - return tools[0] - cfg = card.get("postprocessing", {}) - try: - if float(cfg.get("time_of_flight", -1.0)) >= 0: - return "time_of_flight" - except (TypeError, ValueError): - pass - # the legacy systematics.py path only runs when madspace did not compute - # the weights itself (see run_lhe_postprocessing) - if cfg.get("systematics") \ - and not card.get("systematics", {}).get("enable", True): - return "systematics" - return None - - -def force_lhe_output_if_needed(switch) -> None: - """The events are written as npy by default; switch the run_card to the - LHE format when a selected post-processing needs an LHE file. The card is - rewritten rather than refused, so that asking for a shower in the launch - question is enough -- the user does not have to know that it constrains - output_format. - - The card is read as raw TOML and only its output_format line is edited: a - run_card scan ("scan:[...]" values) is not a valid RunCardMG7 and must - survive untouched for run_scan.""" - import tomllib - path = os.path.join("Cards", "run_card.toml") - if not os.path.exists(path): - return - with open(path, "rb") as f: - card = tomllib.load(f) - output_format = card.get("run", {}).get( - "output_format", RunCardMG7()["run"]["output_format"]) - if output_format == "lhe": - return - reason = _lhe_needed_by(switch, card) - if reason is None: - return - with open(path) as f: - text = f.read() - text, count = re.subn(r'(?m)^(\s*output_format\s*=\s*)("[^"]*"|\'[^\']*\'|\S+)', - r'\g<1>"lhe"', text, count=1) - if not count: - text, count = re.subn(r'(?m)^\[run\][^\n]*\n', '\\g<0>output_format = "lhe"\n', - text, count=1) - if not count: - text += '\n[run]\noutput_format = "lhe"\n' - with open(path, "w") as f: - f.write(text) - logging.getLogger("madevent").info( - "output_format set to 'lhe' (required by %s).", reason) - - -def _raise_open_file_limit() -> None: - """Remove the soft limit on the number of open files: it can be quite low - on some systems and the event-combination step opens one file per channel.""" - soft_lim, hard_lim = resource.getrlimit(resource.RLIMIT_NOFILE) - try: - resource.setrlimit(resource.RLIMIT_NOFILE, (hard_lim, hard_lim)) - except (ValueError, OSError) as error: - logger.debug("could not raise the open-file limit: %s", error) - - -def main() -> None: - """``bin/generate_events`` entry point. - - It runs the same :class:`MG7Cmd` command MG5's in-process ``launch`` uses, - so the two entry points cannot drift apart; the only difference is that - nothing here supplies a ``mother``/``inputfile``, so the question falls back - to reading this process's own stdin. - """ - parser = argparse.ArgumentParser() - parser.add_argument("-f", action="store_false", dest="ask_edit_cards") - args = parser.parse_args() - - MG7Cmd().run_cmd("generate_events" if args.ask_edit_cards else "generate_events -f") diff --git a/madgraph/iolibs/template_files/mg7/madgraph_process.py b/madgraph/iolibs/template_files/mg7/madgraph_process.py new file mode 100644 index 0000000000..7321b86638 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/madgraph_process.py @@ -0,0 +1,1966 @@ +"""The run of one process directory: cards, integration, events and output.""" + +import glob +import json +import os +import shutil +import subprocess +import time + +from madgraph.iolibs.template_files.mg7.common import backend_of, format_time, fptype_of, get_start_time, interference_helicities_summed, lhe_helicities, lhe_weight_info, logger, me_parameters, ms, resolve_auto_backend, resolve_seed, resolve_verbosity +from madgraph.various import misc +from madgraph.various.banner import RunCardMG7 +from models.check_param_card import ParamCard +from madgraph.iolibs.template_files.mg7 import hwu_output +from madgraph.iolibs.template_files.mg7 import npy_to_lhe +from madgraph.iolibs.template_files.mg7 import options +from madgraph.iolibs.template_files.mg7 import plots +from madgraph.iolibs.template_files.mg7 import systematics_summary +from madgraph.iolibs.template_files.mg7.bootstrap import INSTALL_DIR as _INSTALL_DIR +from madgraph.iolibs.template_files.mg7.bootstrap import MADSPACE_DIR as _MADSPACE_DIR +from madgraph.iolibs.template_files.mg7.custom_function_loader import load_custom_function +from madgraph.iolibs.template_files.mg7.observables import CutItem, HistItem, build_observable, clean_pids, observable_pids +from madgraph.iolibs.template_files.mg7.options import lhapdf_paths +from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import MadgraphSubprocess, PhaseSpace, build_multi_channel_data + +class MadgraphProcess: + def __init__(self): + self.madnis_trained = False + self.lhe_completer = None + self.load_cards() + self.init_backend() + self.init_event_dir() + self.init_context() + self.init_custom_functions() + self.init_surrogate() + self.init_cuts() + self.init_histograms() + self.init_generator_config() + self.init_beam() + self.init_subprocesses() + + def load_cards(self) -> None: + self.run_card = RunCardMG7(os.path.join("Cards", "run_card.toml")) + # Resolved once so every generator built during this run shares the same + # seed; the concrete value (even if randomly drawn) is recorded in each + # generator's info.json. + self.run_seed = resolve_seed(self.run_card["run"]["seed"]) + self.param_card_path = os.path.join("Cards", "param_card.dat") + self.param_card = ParamCard(self.param_card_path) + with open(os.path.join("SubProcesses", "subprocesses.json")) as f: + self.subprocess_data = json.load(f) + if self.run_card["phasespace"]["merge_subprocesses"]: + with open(os.path.join("SubProcesses", "merged_subprocesses.json")) as f: + self.merged_subprocess_data = json.load(f) + else: + self.merged_subprocess_data = None + + self.init_decay_mode() + self.init_me_frame() + if interference_helicities_summed(self.run_card): + count = sum(1 for meta in self.subprocess_data if meta.get("interference")) + if count: + logger.info( + "interference_helicity = summed: the weight of the %d interference " + "subprocess(es) is the helicity sum, and their helicities are " + "written as 9 in the LHE", count) + + def init_me_frame(self) -> None: + """Resolve the run card's me_frame into the list of external particles + whose momentum sum defines the frame the matrix element is evaluated in. + + Empty (the default) means no boost: the matrix element sees the momenta + in the frame they are generated in. Unlike madevent, which hands its + matrix element momenta that are already in the partonic centre of mass, + madspace generates them in the lab frame, so asking for madevent's own + frame ([1, 2]) does mean an actual boost here. It changes nothing for a + Lorentz invariant matrix element (only the rounding), but it is what + defines the polarisation axes of a polarised one. + + The particle numbers are checked here rather than only in madspace, + which would not see them until the integrands are built -- long after + the matrix-element libraries have been compiled. Every subprocess is + checked against its own particles: an output can mix multiplicities + (p p > z j and p p > z j j, say, for MLM merging), and madspace builds + one matrix element per subprocess, each refusing a particle it does not + have. That is stricter than madevent, which ignores the particles a + process does not have and so changes the frame from one multiplicity + to the next without a word. + """ + me_frame = list(self.run_card["run"]["me_frame"]) + pdgs = [clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) + for meta in self.subprocess_data] + counts = sorted({len(pdg_list) for pdg_list in pdgs}) + for meta, pdg_list in zip(self.subprocess_data, pdgs): + for index in me_frame: + if 1 <= index <= len(pdg_list): + continue + process = "%s > %s" % ( + " ".join(str(pid) for pid in meta["incoming"]), + " ".join(str(pid) for pid in meta["outgoing"]), + ) + mixed = ( + f"; this output mixes subprocesses with " + f"{', '.join(str(count) for count in counts)} external " + "particles, and me_frame has to name particles every one " + "of them has" + if len(counts) > 1 else "" + ) + raise ValueError( + f"me_frame particle {index} out of range: subprocess " + f"{process} has {len(pdg_list)} external particles{mixed}" + ) + if len(me_frame) == 1: + # A massless particle has no rest frame: its momentum squares to + # zero, so the boost divides by it and every momentum comes back as + # a NaN, silently. Name the mistake instead. + massless = [pdg_list[me_frame[0] - 1] for pdg_list in pdgs + if self.get_mass(pdg_list[me_frame[0] - 1]) == 0.] + if massless: + raise ValueError( + f"me_frame = {me_frame} asks for the rest frame of particle " + f"{me_frame[0]}, which is massless (pdg " + f"{sorted(set(massless))}) and has none. Name a massive " + "particle, or several particles whose sum is massive." + ) + self.me_frame = me_frame + self.incoming_count = 1 if self.is_decay else 2 + + def init_decay_mode(self) -> None: + """Decide whether this directory is a decay (1 -> n) or a collision. + + Read off the exported process rather than the run card: the two cannot + then disagree. MadSpin generates its decay matrix elements this way. + """ + incoming_counts = { + len(clean_pids(meta["incoming"])) for meta in self.subprocess_data + } + if incoming_counts - {1, 2}: + raise ValueError( + f"processes with {sorted(incoming_counts)} incoming particles " + "are not supported" + ) + if len(incoming_counts) > 1: + raise ValueError( + "cannot mix decays and collisions in one output directory" + ) + self.is_decay = incoming_counts == {1} + if not self.is_decay: + self.decaying_mass = None + return + + masses = { + self.get_mass(clean_pids(meta["incoming"])[0]) + for meta in self.subprocess_data + } + if len(masses) > 1: + raise ValueError( + f"decaying particles have different masses: {sorted(masses)}" + ) + self.decaying_mass = masses.pop() + if self.decaying_mass <= 0.0: + raise ValueError("the decaying particle must have a non-zero mass") + self.drop_closed_channels() + + def drop_closed_channels(self) -> None: + """Remove subprocesses the decaying particle is too light to produce. + + A multiparticle decay definition enumerates every vertex the model + allows, closed ones included: "t > b w+, w+ > all all" yields + t > b t b~ (and b W+ Z, b W+ h, ...), which need more mass than the top + has. Their partial width is exactly zero, but the phase-space mapping + has no physical point to hand back -- the invariant's lower bound ends + up above its upper bound -- so it produces NaN momenta and poisons the + whole integral. Drop them here instead. + """ + kept, dropped = [], [] + for meta in self.subprocess_data: + total = sum( + self.get_mass(pid) for pid in clean_pids(meta["outgoing"]) + ) + if total < self.decaying_mass: + kept.append(meta) + else: + dropped.append((meta["outgoing"], total)) + if dropped and self.merged_subprocess_data is not None: + # merged_subprocesses.json indexes the *unfiltered* subprocess + # list, so dropping entries here would silently shift every index + # it holds. Refuse rather than mis-map. + raise ValueError( + "merge_subprocesses is not supported for a decay directory " + "with kinematically closed channels" + ) + if dropped: + for outgoing, total in dropped: + logger.info( + "skipping closed decay channel -> %s (needs %.4g GeV, " + "the decaying particle has %.4g GeV)", + outgoing, total, self.decaying_mass, + ) + if not kept: + raise ValueError( + "every decay channel is kinematically closed: the decaying " + f"particle's mass is {self.decaying_mass} GeV" + ) + self.subprocess_data = kept + + def init_backend(self) -> None: + ms.set_simd_vector_size(self.run_card["run"]["simd_vector_size"]) + + def init_event_dir(self) -> None: + output_dir = self.run_card["gridpack"]["output_dir"] + temp_dir = self.run_card["gridpack"]["temp_output_dir"] + if output_dir: + self.run_path = os.path.abspath(os.path.expanduser(output_dir)) + os.makedirs(self.run_path, exist_ok=True) + else: + run_name = self.run_card["run"]["run_name"] + os.makedirs("Events", exist_ok=True) + run_dir_prefix = os.path.join("Events", f"{run_name}_") + existing_run_dirs = glob.glob(f"{run_dir_prefix}*") + run_index = 1 + for run_dir in existing_run_dirs: + run_index_str = run_dir[len(run_dir_prefix):] + if run_index_str.isnumeric(): + run_index = max(run_index, int(run_index_str) + 1) + while True: + try: + self.run_path = f"{run_dir_prefix}{run_index:02d}" + os.mkdir(self.run_path) + break + except FileExistsError: + run_index += 1 + # temporary npy files of the channel generators + if temp_dir: + self.temp_path = os.path.abspath(os.path.expanduser(temp_dir)) + os.makedirs(self.temp_path, exist_ok=True) + else: + self.temp_path = self.run_path + # Absolute on purpose. StatusFile is a C++ object that finishes its + # write from its destructor (rename info.json.tmp -> info.json), and + # that destructor runs whenever Python gets round to collecting the + # process object. When the launcher runs in MG5's process, an + # interrupted run is collected *after* the launch command has restored + # the working directory, and a relative path would then resolve + # somewhere else -- the rename throws a C++ filesystem_error out of a + # destructor, which is an immediate abort. + self.status_file = ms.StatusFile( + os.path.abspath(os.path.join(self.run_path, "info.json"))) + + def init_context(self) -> None: + device_names = self.run_card["run"]["device"] + self.contexts = [] + self.device_types = [] + self.devices = [] + self.pool_sizes = [] + for device_name in device_names: + if ":" in device_name: + device_type, device_index_str = device_name.split(":") + device_index = int(device_index_str) + else: + device_type = device_name + device_index = 0 + self.device_types.append(device_type) + if device_type == "cuda": + device = ms.cuda_device(device_index) + pool_size = self.run_card["run"]["gpu_thread_pool_size"] + elif device_type == "hip": + device = ms.hip_device(device_index) + pool_size = self.run_card["run"]["gpu_thread_pool_size"] + else: + device = ms.cpu_device() + pool_size = self.run_card["run"]["cpu_thread_pool_size"] + self.devices.append(device) + self.pool_sizes.append(pool_size) + self.contexts.append(ms.Context(device=device, thread_count=pool_size)) + + CUSTOM_FUNCTION_KINDS = ("cut", "scale", "pdf", "bias") + + def init_custom_functions(self) -> None: + """Load the functions of [custom_functions], building the libraries + first. They are loaded before the matrix elements, so that they have the + same index in every context and in a gridpack made from this run.""" + self.custom_functions = {} + self.custom_function_paths = [] + for kind in self.CUSTOM_FUNCTION_KINDS: + path = self.run_card["custom_functions"][kind] + if not path: + continue + build_dir = f"custom_{kind}" + if path.endswith(".so") and os.path.isfile(os.path.join(build_dir, "Makefile")): + self.build_custom_function(build_dir) + if not os.path.isfile(path): + raise RuntimeError( + f"[custom_functions] {kind} = {path!r}: file not found. " + "Create a .so library or a .py file with bin/init_custom_func." + ) + apis = load_custom_function(ms, self.contexts, self.device_types, kind, path) + if len({api.index() for api in apis}) != 1: + raise RuntimeError(f"custom {kind} function loaded at different indices") + self.custom_functions[kind] = apis[0] + self.custom_function_paths.append([kind, path]) + logger.info("Using the custom %s function %s", kind, path) + + def init_surrogate(self) -> None: + """Check the [surrogate] settings. The surrogates are loaded per + subprocess, by MadgraphSubprocess.""" + settings = self.run_card["surrogate"] + self.surrogate = None + if not settings["function"]: + return + path = settings["function"] + if path.endswith(".py") and not os.path.isfile(path): + raise RuntimeError(f"[surrogate] function = {path!r}: file not found") + if self.run_card["phasespace"]["sde_strategy"] != "denominators": + raise RuntimeError( + "[surrogate] needs [phasespace] sde_strategy = \"denominators\", " + "as the surrogate does not return channel weights" + ) + if self.run_card["run"]["dummy_matrix_element"]: + raise RuntimeError("[surrogate] cannot be used with dummy_matrix_element") + if self.run_card["gridpack"]["save_gridpack"]: + raise RuntimeError("[surrogate] does not support gridpacks yet") + self.surrogate = dict(settings) + if not self.surrogate["param_card"]: + self.surrogate["param_card"] = self.param_card_path + logger.info( + "Using the surrogate %s, %s with the full matrix element", + path, + "unweighting" if settings["mode"] == "unweight" else "reweighting", + ) + + def weighting_groups(self, subprocess_channels: list[list[int]]) -> list: + """One ms.WeightingGroup per subprocess with a [surrogate]: the second + pass with the full matrix element over its channels, given by their + indices in the event generator.""" + if self.surrogate is None: + return [] + gen_args = self.run_card["generation"] + return [ + ms.WeightingGroup( + channel_indices, + [ms.WeightingStep( + subproc.step_matrix_elements, + mode=self.surrogate["mode"], + max_overweight_truncation=gen_args["max_overweight_truncation"], + freeze_max_weight_after=gen_args["freeze_max_weight_after"], + )], + ) + for subproc, channel_indices in zip(self.subprocesses, subprocess_channels) + if channel_indices + ] + + @staticmethod + def build_custom_function(build_dir: str) -> None: + """Run make in a custom_ directory, logging to compile.log.""" + log_path = os.path.join(build_dir, "compile.log") + with open(log_path, "w") as log: + returncode = subprocess.call( + ["make"], cwd=build_dir, stdout=log, stderr=subprocess.STDOUT + ) + if returncode != 0: + raise RuntimeError( + f"Compilation of {build_dir} failed with exit status {returncode}; " + f"see {log_path} for details" + ) + + def custom_function(self, kind: str, particle_count: int, pid_options: list): + """The ms.CustomFunction of [custom_functions] , or None.""" + api = self.custom_functions.get(kind) + if api is None: + return None + function_type = getattr(ms.CustomFunction, kind) + return ms.CustomFunction(api, function_type, particle_count, pid_options) + + def parse_observable(self, name: str, order_observable: str) -> dict: + parts = name.split("-") + sum_momenta = False + sum_observable = False + ordered = False + multiparticles = self.run_card["multiparticles"] + + if len(parts) == 0: + raise ValueError("Invalid observable name") + elif len(parts) == 1: + # event-level observables + obs_name = parts[0] + select_pids = [] + else: + if parts[-1] == "sum": + sum_observable = True + obs_name = parts[-2] + selection = parts[:-2] + elif parts[-2] == "sum": + sum_momenta = True + obs_name = parts[-1] + selection = parts[:-2] + else: + obs_name = parts[-1] + selection = parts[:-1] + select_pids = [] + order_indices = [] + for mp_name in selection: + mp_parts = mp_name.split("_") + if mp_parts[-1].isnumeric(): + order_indices.append(int(mp_parts[-1])) + select_pids.append(multiparticles["_".join(mp_parts[:-1])]) + ordered = True + else: + order_indices.append(0) + select_pids.append(multiparticles[mp_name]) + + return dict( + observable=obs_name, + select_pids=select_pids, + sum_momenta=sum_momenta, + sum_observable=sum_observable, + order_observable=order_observable if ordered else None, + order_indices=order_indices if ordered else [], + ignore_incoming=True, + name=name, + ) + + def init_cuts(self) -> None: + inf = float("inf") + order_observable = self.run_card["cuts"].get("order_by", "pt") + self.cut_data = [ + CutItem( + observable_kwargs=self.parse_observable(key, order_observable), + min=values.get("min", -inf), + max=values.get("max", inf), + mode=values.get("mode", "all"), + ) + for key, values in self.run_card["cuts"].items() + if key != "order_by" + ] + + # [histograms] key that means "the distribution of the event weight", + # normalised to the unit weight sigma_abs, rather than an observable of the + # momenta: an unweighted sample is a spike at 1 (and at -1 for its negative + # weights) and a partially unweighted one shows its spread. + weight_histogram_key = "weight" + + def init_histograms(self) -> None: + order_observable = self.run_card["histograms"].get("order_by", "pt") + self.hist_data = [] + for key, values in self.run_card["histograms"].items(): + if key == "order_by": + continue + from_weight = key == self.weight_histogram_key + self.hist_data.append(HistItem( + observable_kwargs=( + None if from_weight + else self.parse_observable(key, order_observable)), + min=values["min"], + max=values["max"], + bin_count=values["bin_count"], + name=key, + from_weight=from_weight, + )) + + def ensure_pdf_set(self, pdf_set: str) -> misc.LhapdfPaths: + """Make sure the requested LHAPDF set is available, downloading it into + the first writable PDF directory if needed, and return the (possibly + updated) LHAPDF resolution. Nothing is downloaded when the set is + already present.""" + paths = lhapdf_paths() + if paths.find_set(pdf_set): + return paths + if not paths.config: + logger.debug("no usable lhapdf-config: cannot download %s", pdf_set) + return paths + if not paths.download_path: + logger.debug("no writable PDF directory: cannot download %s", pdf_set) + return paths + try: + from madgraph.interface.common_run_interface import CommonRunCmd + os.makedirs(paths.download_path, exist_ok=True) + logger.info("PDF set %s not found; downloading into %s", + pdf_set, paths.download_path) + CommonRunCmd.install_lhapdf_pdfset_static( + paths.config, paths.download_path, pdf_set) + except Exception as err: + logger.warning("Could not download PDF set %s: %s", pdf_set, err) + return paths + options._LHAPDF = paths.with_data_path(paths.download_path) + return options._LHAPDF + + def _no_pdf_message(self, pdf_set: str) -> str: + """Explain a missing PDF set in terms of the settings that fix it.""" + searched = os.pathsep.join(self.lhapdf.data_paths) or \ + "(no PDF set directory could be located)" + if self.lhapdf.config: + found = "using lhapdf-config: %s" % self.lhapdf.config + else: + found = ("no usable lhapdf-config was found, so the set could not " + "be downloaded automatically") + return ( + "LHAPDF set %r, requested by [beam] pdf in Cards/run_card.toml, " + "was not found.\n" + " searched: %s\n" + " %s\n" + "Fix this in one of the following ways:\n" + " * install LHAPDF for this MadGraph: MG7> install lhapdf6\n" + " * point MadGraph at an existing one:\n" + " MG7> set lhapdf /path/to/lhapdf-config\n" + " MG7> save options\n" + " * or set it for this directory only, by adding\n" + " lhapdf = /path/to/lhapdf-config\n" + " to %s\n" + " * or set $LHAPDF_DATA_PATH to a directory containing %r." + % (pdf_set, searched, found, + os.path.join("Cards", "me5_configuration.txt"), pdf_set)) + + def init_beam(self) -> None: + beam_args = self.run_card["beam"] + + # built lazily by build_systematics (needs the subprocess data) + self.systematics = None + self.systematics_data = None + self.systematics_context = None + self.event_histograms = None + self.event_histograms_context = None + + if self.is_decay: + # No beams: the total energy is the decaying particle's mass, and + # "leptonic" is what the mappings call "no parton luminosity". + self.e_cm = self.decaying_mass + self.leptonic = True + else: + self.e_cm = beam_args["e_cm"] + self.leptonic = beam_args["leptonic"] + + dynamical_scales = { + "transverse_energy": ms.EnergyScale.transverse_energy, + "transverse_mass": ms.EnergyScale.transverse_mass, + "half_transverse_mass": ms.EnergyScale.half_transverse_mass, + "partonic_energy": ms.EnergyScale.partonic_energy, + } + if beam_args["dynamical_scale_choice"] in dynamical_scales: + dynamical_scale_type = dynamical_scales[beam_args["dynamical_scale_choice"]] + else: + raise ValueError("Unknown dynamical scale choice") + self.scale_kwargs = dict( + dynamical_scale_type=dynamical_scale_type, + ren_scale_fixed=beam_args["fixed_ren_scale"], + fact_scale_fixed=beam_args["fixed_fact_scale"], + ren_scale=beam_args["ren_scale"], + fact_scale1=beam_args["fact_scale1"], + fact_scale2=beam_args["fact_scale2"], + scale_factor=float(beam_args.get("scale_factor", 1.)), + ) + if self.is_decay and self.custom_functions.pop("scale", None) is not None: + logger.warning("custom scale ignored: a decay uses the decaying mass") + if self.is_decay and self.custom_functions.pop("pdf", None) is not None: + logger.warning("custom pdf ignored: a decay has no parton densities") + if self.is_decay: + # One scale is available for a decay -- the decaying mass -- so use + # it, fixed, whatever the card asks for. + self.scale_kwargs.update( + ren_scale_fixed=True, + fact_scale_fixed=True, + ren_scale=self.decaying_mass, + fact_scale1=self.decaying_mass, + fact_scale2=self.decaying_mass, + scale_factor=1., + ) + self.pdf_grid = None + self.lhapdf = None + self.pdf_dir = None + self.fixed_alphas_info = self.write_fixed_alphas_info() + self.alphas_grid = ms.AlphaSGrid(self.fixed_alphas_info) + for context in self.contexts: + self.alphas_grid.initialize_globals(context) + self.running_coupling = ms.RunningCoupling(self.alphas_grid) + return + + pdf_set = beam_args["pdf"] + self.pdf_set = pdf_set + self.fixed_alphas_info = None + self.lhapdf = self.ensure_pdf_set(pdf_set) + self.pdf_dir = self.lhapdf.find_set(pdf_set) + if self.pdf_dir is None: + raise RuntimeError(self._no_pdf_message(pdf_set)) + self.alphas_grid = ms.AlphaSGrid(os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}.info")) + if "pdf" in self.custom_functions: + # alpha_s still comes from the PDF set of the run card + self.pdf_grid = ms.CustomPdf(self.custom_functions["pdf"]) + else: + self.pdf_grid = ms.PdfGrid( + os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}_0000.dat") + ) + for context in self.contexts: + if not isinstance(self.pdf_grid, ms.CustomPdf): + self.pdf_grid.initialize_globals(context) + self.alphas_grid.initialize_globals(context) + self.running_coupling = ms.RunningCoupling(self.alphas_grid) + + # ------------------------------------------------------------------ + # scale / PDF systematics + # ------------------------------------------------------------------ + def systematics_enabled(self) -> bool: + """The native scale/PDF weights are requested ([systematics] enable, or + the legacy [generation] systematics alias).""" + try: + enabled = bool(self.run_card["systematics"]["enable"]) + except Exception: + enabled = False + if not enabled: + try: + enabled = bool(self.run_card["generation"]["systematics"]) + except Exception: + enabled = False + if enabled and "pdf" in getattr(self, "custom_functions", {}): + # the variations re-evaluate the PDF set, which a custom PDF replaces + if not getattr(self, "_warned_custom_pdf_systematics", False): + logger.warning("systematics are disabled with a custom pdf") + self._warned_custom_pdf_systematics = True + return False + return enabled + + @staticmethod + def pdf_set_info(info_path: str) -> dict: + """Read SetIndex, NumMembers, ErrorType and SetDesc from an LHAPDF + .info file (missing keys get neutral defaults).""" + info = {"SetIndex": 0, "NumMembers": 1, "ErrorType": "", "SetDesc": ""} + try: + with open(info_path) as f: + for line in f: + key, sep, value = line.partition(":") + key = key.strip() + if not sep or key not in info: + continue + value = value.strip().strip('"').strip("'") + if key in ("SetIndex", "NumMembers"): + info[key] = int(value) + else: + info[key] = value + except OSError as err: + logger.warning("could not read %s: %s", info_path, err) + return info + + def resolve_pdf_set_name(self, entry: str) -> tuple[str, int | None]: + """Turn a [systematics] pdf entry ("", "", "@", + "@") into (set name, member or None) using the LHAPDF + pdfsets.index when an id is given.""" + member = None + if "@" in entry: + entry, member_str = entry.split("@", 1) + member = int(member_str) + entry = entry.strip() + if not entry.isdigit(): + return entry, member + lhaid = int(entry) + index_path = os.path.join(self.pdf_dir or "", "pdfsets.index") + sets = [] + try: + with open(index_path) as f: + for line in f: + parts = line.split() + if len(parts) >= 2 and parts[0].isdigit(): + sets.append((int(parts[0]), parts[1])) + except OSError: + raise ValueError( + "cannot resolve LHAPDF id %s: no %s" % (lhaid, index_path)) + for set_id, name in sets: + if set_id == lhaid: + return name, member + # an id inside a set's member range + sets.sort() + for (set_id, name), (next_id, _) in zip(sets, sets[1:] + [(None, None)]): + if set_id < lhaid and (next_id is None or lhaid < next_id): + return name, lhaid - set_id if member is None else member + raise ValueError("cannot resolve LHAPDF id %s" % lhaid) + + def resolve_pdf_variations(self) -> list: + """Expand the [systematics] pdf entries into ms.PdfMemberSpec objects + (downloading the sets when needed). The nominal member itself is never + listed: the calculator adds it as the central weight of its group.""" + entries = self.run_card["systematics"]["pdf"] + if isinstance(entries, str): + entries = [entries] + specs = [] + seen = set() + for raw in entries: + entry = str(raw).strip() + if not entry or entry.lower() in ("none", "no", "false"): + continue + if entry.lower() == "central": + continue + if entry.lower() == "errorset": + set_name, member = self.pdf_set, None + else: + set_name, member = self.resolve_pdf_set_name(entry) + variation_lhapdf = self.ensure_pdf_set(set_name) + variation_dir = variation_lhapdf.find_set(set_name) + if variation_dir is None: + logger.warning("PDF set %s not available; its variation is skipped", set_name) + continue + info_path = os.path.join(variation_dir, set_name, f"{set_name}.info") + if not os.path.exists(info_path): + logger.warning("PDF set %s not available; its variation is skipped", set_name) + continue + info = self.pdf_set_info(info_path) + is_nominal = set_name == self.pdf_set + if member is not None: + members = [member] + else: + members = list(range(info["NumMembers"])) + if is_nominal: + members = members[1:] + if not members: + logger.warning( + "PDF set %s has a single member: no PDF uncertainty can be " + "computed from it. Use a set with error members in " + "[systematics] pdf (e.g. a PDF4LHC or NNPDF error set).", + set_name) + for m in members: + if is_nominal and m == 0: + continue + if (set_name, m) in seen: + continue + seen.add((set_name, m)) + grid_file = os.path.join(variation_dir, set_name, f"{set_name}_{m:04d}.dat") + if not os.path.exists(grid_file): + logger.warning("PDF member file %s not found; skipped", grid_file) + continue + specs.append(ms.PdfMemberSpec( + set_name=set_name, set_lhaid=info["SetIndex"], member=m, + grid_file=grid_file, info_file=info_path, + error_type=info["ErrorType"], description=info["SetDesc"], + )) + return specs + + def build_systematics_args(self) -> list: + """One ms.SubprocessSystArgs per (unmerged) subprocess: the alpha_s power + of |M|^2 and the beam parton ids of every flavor index, indexed like the + LHECompleter (subprocess_index / flavor_index of the combined events).""" + args = [] + for meta in self.subprocess_data: + beam_pdgs = [ + [int(flavor["options"][0][0]), int(flavor["options"][0][1])] + for flavor in meta["flavors"] + ] + args.append(ms.SubprocessSystArgs( + qcd_power=int(meta.get("qcd_power", -1)), beam_pdgs=beam_pdgs)) + return args + + def build_systematics(self): + """Create the ms.SystematicsCalculator from the [systematics] section, or + None when disabled. Also keeps the JSON description needed by gridpacks.""" + if not self.systematics_enabled(): + self.systematics = None + return None + syst = self.run_card["systematics"] + config = ms.SystematicsConfig() + config.mur = [float(v) for v in syst["mur"]] + config.muf = [float(v) for v in syst["muf"]] + config.together = bool(syst["together"]) + config.dyn_scales = self.resolve_dynamical_scales() + # alternative dynamical scales get the same factor as the nominal one + config.scale_factor = float(self.run_card["beam"].get("scale_factor", 1.)) + config.write_inputs = bool(syst["write_inputs"]) + config.has_pdf = not self.leptonic + # Without parton luminosity -- a decay, or leptonic beams -- there is no + # nominal PDF *set* to describe and no PDF variation to compute, so + # has_pdf gates the nominal-set metadata and the member list below. + # The alpha_s grid is a separate matter: the mu_R variations always need + # one, so its .info file has to be recorded whatever has_pdf says. + # init_beam() skips the PDF lookup for a decay only -- leptonic beams + # still resolve a set and take alpha_s from it -- so the path follows + # self.pdf_dir, and falls back to the constant grid written for a decay. + if self.pdf_dir is not None: + info_path = os.path.join(self.pdf_dir, self.pdf_set, f"{self.pdf_set}.info") + else: + info_path = self.fixed_alphas_info + if config.has_pdf: + info = self.pdf_set_info(info_path) + config.nominal_set_name = self.pdf_set + config.nominal_lhaid = info["SetIndex"] + config.nominal_error_type = info["ErrorType"] + config.nominal_description = info["SetDesc"] + config.pdf_members = self.resolve_pdf_variations() + else: + config.pdf_members = [] + args = self.build_systematics_args() + # the PDFs, alpha_s and (for mixed-order subprocesses) the matrix + # elements are evaluated with the batched madspace functions on this + # CPU context + self.systematics_context = ms.Context(device=ms.cpu_device(), thread_count=1) + matrix_elements, flavor_remap, me_backend = self.build_systematics_matrix_elements( + args, self.systematics_context) + self.systematics = ms.SystematicsCalculator( + config, args, self.pdf_grid, self.alphas_grid, + self.systematics_context, matrix_elements, flavor_remap) + for warning in self.systematics.warnings: + logger.warning("systematics: %s", warning) + self.systematics_data = { + "config": json.loads(config.to_json()), + "subproc_args": [json.loads(a.to_json()) for a in args], + "nominal_grid_file": os.path.join( + self.pdf_dir, self.pdf_set, f"{self.pdf_set}_0000.dat") + if config.has_pdf else None, + "nominal_info_file": info_path, + # matrix elements re-evaluated for the mixed-order subprocesses + "me_paths": [meta["me_path"] for meta in self.subprocess_data], + "me_backend": me_backend, + "flavor_remap": flavor_remap, + "me_frame": self.me_frame, + "incoming_count": self.incoming_count, + } + return self.systematics + + _DYN_SCALE_CODES = { + "transverse_energy": 1, "transverse_mass": 2, + "half_transverse_mass": 3, "partonic_energy": 4, + } + + def resolve_dynamical_scales(self) -> list: + """[systematics] dynamical_scale entries (names or systematics.py codes + 1-4) as codes, without the choice the events were generated with.""" + try: + entries = self.run_card["systematics"]["dynamical_scale"] + except Exception: + return [] + if isinstance(entries, str): + entries = [entries] + beam = self.run_card["beam"] + generated = None + if not (beam["fixed_ren_scale"] and beam["fixed_fact_scale"]): + generated = self._DYN_SCALE_CODES.get(beam["dynamical_scale_choice"]) + codes = [] + for raw in entries: + entry = str(raw).strip().lower() + if not entry: + continue + if entry.isdigit(): + code = int(entry) + elif entry in self._DYN_SCALE_CODES: + code = self._DYN_SCALE_CODES[entry] + else: + raise ValueError("unknown dynamical scale choice %r in [systematics] " + "dynamical_scale" % raw) + if code not in (1, 2, 3, 4): + raise ValueError("dynamical scale code %s must be 1-4" % code) + if code == generated or code in codes: + continue + codes.append(code) + return codes + + def build_systematics_matrix_elements(self, args, context): + """For the subprocesses whose |M|^2 mixes several powers of alpha_s + (qcd_power = -1), load the matrix element into the systematics CPU + context so that the renormalisation scale variations can re-evaluate it + at the varied alpha_s. Returns ([MatrixElement or None per subprocess], + flavor remap, backend name); empty when nothing is needed or no CPU + library exists.""" + need = [i for i, a in enumerate(args) if a.qcd_power < 0] + empty = ([], [], None) + if not need or self.run_card["run"]["dummy_matrix_element"]: + return empty + cpu_backends = [b for b in self.backends if not b.startswith(("cuda", "hip"))] + if not cpu_backends: + logger.warning("systematics: no CPU matrix-element library available, the " + "renormalisation scale variations of the mixed-order " + "subprocesses cannot be computed") + return empty + backend = cpu_backends[0] + matrix_elements = [] + flavor_remap = [] + for i, meta in enumerate(self.subprocess_data): + flavor_remap.append([int(f["index"]) for f in meta["flavors"]]) + if i not in need: + matrix_elements.append(None) + continue + api = context.load_matrix_element( + meta["me_path"].format(device=backend), self.param_card_path, + me_parameters(self.run_card)) + matrix_elements.append(ms.MatrixElement( + api, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in], + [ms.MatrixElement.matrix_element_out], + False, + self.me_frame, + self.incoming_count, + )) + logger.info("systematics: re-evaluating the matrix element of %d mixed-order " + "subprocess(es) for the renormalisation scale variations", len(need)) + return matrix_elements, flavor_remap, backend + + def build_event_histograms(self): + """ms.EventHistograms for the [histograms] observables, filled with the + written events and all their variation weights (info.json + "event_histograms"); None without histograms.""" + if not self.hist_data \ + or not self.run_card["run"]["postprocessing_histograms"]: + self.event_histograms = None + return None + context = ms.Context(device=ms.cpu_device(), thread_count=1) + specs = [ + ms.EventHistogramSpec( + name=item.name, min=item.min, max=item.max, + bin_count=item.bin_count, from_weight=item.from_weight) + for item in self.hist_data + ] + from_momenta = [item for item in self.hist_data if not item.from_weight] + observables = [] + for meta in self.subprocess_data: + if not from_momenta: + # nothing to evaluate on the momenta (only the weight is + # histogrammed): no observable runtime to build + observables.append(None) + continue + all_pids = clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) + values = ms.ObservableValues([ + build_observable( + observable_pids(meta, item.observable_kwargs), + len(meta["incoming"]), item.observable_kwargs + ) + for item in from_momenta + ]) + observables.append(ms.SubprocessObservables(values, len(all_pids))) + # the weight histograms are drawn in units of the unit weight sigma_abs + # (the cross section of a sample without negative weights) + self.event_histograms_context = context + self.event_histograms = ms.EventHistograms( + context, specs, observables, + reference_weight=self._mean_event_weight()) + return self.event_histograms + + def _mean_event_weight(self) -> float: + """The unit weight of the unweighted events, sigma_abs, used as the + unit of the weight histograms: the cross section of a positive sample, + and for one with negative weights (an interference) what puts them at + +-1 rather than dividing by a signed, possibly small, total. 0 (the raw + weight) when it is not available.""" + try: + return float(self.event_generator.status().mean_abs) + except Exception as error: + logger.debug("no cross section for the weight histograms: %s", error) + return 0. + + def log_systematics_summary(self) -> None: + """Print the scale/PDF uncertainties on the total cross section (the + per-variation cross sections are in events.weights.json / info.json).""" + if self.systematics is None or self.systematics.weight_count == 0: + return + summary = json.loads(self.systematics.summary()) + # the percentages come from systematics_summary, which get_result() + # reads too: the box and the scan row must not disagree, so neither + # computes them itself + xsec = systematics_summary.nominal_cross_section(summary) + if xsec is None: + return + def format_variation(up, down): + # right-justify the signed numbers (not the sign alone) so the % + # values line up across rows, with extra spacing between the two + return f"{'+%.3g' % up:>6}% {'-%.3g' % down:>6}%" + + scale_count = sum(1 for v in self.systematics.variations if v.is_scale) + member_count = len(self.systematics.members) + rows = [("Variations per event:", + f"{self.systematics.weight_count} ({scale_count} scale, {member_count} PDF members)")] + for pdf in summary.get("pdf", []): + rows.append(("PDF set:", f"{pdf['pdf_set']}, {pdf['error_type']}")) + rows.append(("Original cross-section:", f"{xsec} pb")) + scale = systematics_summary.scale_percentages(summary) + if scale is not None: + rows.append(("Scale variation:", format_variation(*scale))) + for _entry, up, down in systematics_summary.pdf_percentages(summary): + rows.append(("PDF variation:", format_variation(up, down))) + if self.event_generator_config.verbosity == ms.Verbosity.pretty: + box = ms.PrettyBox("Systematics", len(rows), [24, 0]) + box.set_column(0, [label for label, _ in rows]) + box.set_column(1, [value for _, value in rows]) + box.print_first() + else: + for label, value in rows: + logger.info("systematics, %s %s", label, value) + + def write_fixed_alphas_info(self) -> str: + """Write a minimal LHAPDF ``.info`` holding a constant alpha_s. + + A decay has no beams, so there is no PDF set to take alpha_s from -- + and demanding one (possibly downloading it) just to evaluate a coupling + would be absurd. The renormalisation scale of a decay is fixed at the + decaying mass anyway, so a constant alpha_s is the right answer, not an + approximation: take it from the param card, exactly as the Fortran + decay matrix elements do. + """ + # SMINPUTS entry 3 is alpha_s(m_Z), the same value the Fortran + # parameter setup feeds to G. + alpha_s = float(self.param_card.get_value("sminputs", 3)) + path = os.path.join(self.run_path, "fixed_alphas.info") + # AlphaSGrid interpolates in log(q^2) across the grid, so give it a + # comfortable number of nodes rather than the bare minimum of three. + q_values = [10 ** (i / 4.0) for i in range(-4, 21)] + with open(path, "w") as f: + f.write("SetDesc: constant alpha_s for a decay (no beams)\n") + f.write("AlphaS_Qs: [%s]\n" % ", ".join(f"{q:g}" for q in q_values)) + f.write( + "AlphaS_Vals: [%s]\n" + % ", ".join(f"{alpha_s:g}" for _ in q_values) + ) + logger.info( + "decay mode: fixed alpha_s = %g at mu = %g GeV", alpha_s, self.e_cm + ) + return path + + def needs_systematics_matrix_elements(self) -> bool: + """The systematics will re-evaluate a matrix element per scale point. + + True when the native weights are on and some subprocess mixes alpha_s + powers (qcd_power < 0), i.e. when build_systematics_matrix_elements() + loads a matrix element. Conservative: it does not check that a CPU + backend exists, so it can be True where that call ends up loading + nothing. + """ + if not self.systematics_enabled(): + return False + if self.run_card["run"]["dummy_matrix_element"]: + return False + return any(int(meta.get("qcd_power", -1)) < 0 + for meta in self.subprocess_data) + + def init_generator_config(self) -> None: + run_args = self.run_card["run"] + gen_args = self.run_card["generation"] + vegas_args = self.run_card["vegas"] + cfg = ms.GeneratorConfig() + cfg.target_count = gen_args["events"] + cfg.vegas_damping = vegas_args["damping"] + cfg.max_overweight_truncation = gen_args["max_overweight_truncation"] + cfg.freeze_max_weight_after = gen_args["freeze_max_weight_after"] + if self.surrogate is not None and self.surrogate["mode"] == "unweight": + # [generation] applies to the second pass, which gives the events + cfg.max_overweight_truncation = ( + self.surrogate["first_pass_max_overweight_truncation"] + ) + cfg.freeze_max_weight_after = ( + self.surrogate["first_pass_freeze_max_weight_after"] + ) + cfg.start_batch_size = vegas_args["start_batch_size"] + cfg.max_batch_size = vegas_args["max_batch_size"] + cfg.survey_min_iters = gen_args["survey_min_iters"] + cfg.survey_max_iters = gen_args["survey_max_iters"] + cfg.survey_target_precision = gen_args["survey_target_precision"] + cfg.optimization_patience = vegas_args["optimization_patience"] + cfg.optimization_threshold = vegas_args["optimization_threshold"] + cfg.cpu_batch_size = gen_args["cpu_batch_size"] + cfg.gpu_batch_size = gen_args["gpu_batch_size"] + cfg.verbosity = resolve_verbosity(run_args["verbosity"]) + cfg.combine_thread_count = run_args["combine_thread_pool_size"] + if self.needs_systematics_matrix_elements() and cfg.combine_thread_count != 1: + # SystematicsCalculator::compute() runs inside the combine workers, + # and a matrix element's per-thread process instances live in a + # ThreadResource indexed by ThreadPool::thread_index() -- the + # *calling* thread's index. A parallel combine pool would therefore + # need one process instance per worker; building N of them to index + # them safely is not worth it, the more so as + # SystematicsCalculator::matrix_elements() already serialises every + # evaluation on its own mutex, so that parallelism was not being + # used anyway. Run the combine stage single-threaded instead. + logger.info( + "systematics: the combine stage runs single-threaded because a " + "subprocess needs its matrix element re-evaluated for the " + "renormalisation scale variations") + cfg.combine_thread_count = 1 + cfg.cut_efficiency_threshold = gen_args["cut_efficiency_threshold"] + cfg.max_cut_repetitions = gen_args["max_cut_repetitions"] + self.event_generator_config = cfg + self.event_generator = None + + def init_subprocesses(self) -> None: + self.backends = self.compile_matrix_elements() + self.subprocesses = [] + if self.merged_subprocess_data is None: + for subproc_id, meta in enumerate(self.subprocess_data): + self.subprocesses.append(MadgraphSubprocess(self, meta, subproc_id)) + else: + for subproc_id, meta in enumerate(self.merged_subprocess_data): + self.subprocesses.append( + MadgraphSubprocess(self, meta, subproc_id, self.subprocess_data) + ) + + def compile_matrix_elements(self) -> list[str]: + """Build the matrix-element library of every subprocess, and return the + build backend of each requested device (in the same order). + + A 'cpu' device builds the backend named by the 'cpu_mode' entry; the + cuda/hip devices ignore cpu_mode and build their own backend. A + cpu_mode of 'auto' is resolved to the actual SIMD backend chosen on + this machine. + + SubProcesses/makefile is a dispatcher over the P* directories, so a + single 'make -j N' there builds all the subprocesses at once with one + shared pool of N jobs: no subprocess is built with N jobs while the + others wait, and none is limited to N/#subprocesses jobs either. + """ + device_names = self.run_card["run"]["device"] + if not isinstance(device_names, list): + device_names = [device_names] + cpu_mode = self.run_card["run"]["cpu_mode"] + backends = [backend_of(name, cpu_mode) for name in device_names] + if not self.subprocess_data: + return backends + + first_proc_path = self.subprocess_data[0]["path"] + subproc_path = os.path.dirname(first_proc_path) + + fptype = fptype_of(self.run_card["run"]["precision"]) + + # Resolve cpu_mode='auto' once (the build rules pick the best SIMD + # backend available here), so all subprocesses agree on the library names. + auto_backend = None + if "auto" in backends: + auto_backend = resolve_auto_backend(first_proc_path) + logger.info("cpu_mode 'auto' resolved as '%s'", auto_backend) + resolved = [ + auto_backend if backend == "auto" else backend + for backend in backends + ] + + nb_core = self.run_card["run"]["cpu_thread_pool_size"] + if not nb_core or nb_core < 0: + nb_core = os.cpu_count() or 1 + + log_path = os.path.join(self.run_path, "compile_subprocesses.log") + for backend in resolved: + missing = [ + meta for meta in self.subprocess_data + if not os.path.isfile(meta["me_path"].format(device=backend)) + ] + if not missing: + continue + logger.info( + f"Start compilation of SubProcesses for device '{backend}' " + f"({len(missing)} subprocess(es), {nb_core} parallel job(s)), " + f"see log detail in {log_path}" + ) + start_time = time.time() + self.make_subprocesses( + subproc_path, + [f"BACKEND={backend}", f"FPTYPE={fptype}", "USEBUILDDIR=1"], + nb_core, log_path, + ) + logger.info( + f"Compilation of SubProcesses done in {time.time() - start_time:.1f} s" + ) + return resolved + + @staticmethod + def make_subprocesses( + subproc_path: str, args: list[str], nb_core: int, log_path: str + ) -> None: + """Run 'make -j nb_core' in SubProcesses/, appending the full build + output to log_path (one run per device, so the file is appended to).""" + command = ["make", f"-j{nb_core}"] + args + with open(log_path, "a") as log: + log.write(f"\n$ cd {subproc_path} && {' '.join(command)}\n") + log.flush() + returncode = subprocess.call( + command, cwd=subproc_path, stdout=log, stderr=subprocess.STDOUT + ) + if returncode != 0: + raise RuntimeError( + f"Compilation of the SubProcesses failed with exit status " + f"{returncode}; see {log_path} for details" + ) + + def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenerator: + channel_generators = [] + subprocess_channels = [[] for _ in self.subprocesses] + for i, (subproc, phasespace) in enumerate(zip(self.subprocesses, phasespaces)): + for integrand, channel in zip( + subproc.build_integrands(phasespace), phasespace.channels + ): + if channel.event_generator is None: + channel.event_generator = ms.ChannelEventGenerator( + contexts=self.contexts, + integrand=integrand, + event_file=os.path.join(self.temp_path, f"events.{i}.{channel.name}.npy"), + weight_file=os.path.join(self.temp_path, f"weights.{i}.{channel.name}.npy"), + config=self.event_generator_config, + subprocess_index=i, + name=f"{i}.{channel.name}", + histograms=subproc.histograms + ) + subprocess_channels[i].append(len(channel_generators)) + channel_generators.append(channel.event_generator) + + event_generator = ms.EventGenerator( + contexts=self.contexts, + channels=channel_generators, + status_file=self.status_file, + config=self.event_generator_config, + seed=self.run_seed, + groups=self.weighting_groups(subprocess_channels), + ) + unused_globals = ( + set(self.contexts[0].global_names()) - event_generator.used_globals() + ) + for context in self.contexts: + for global_name in unused_globals: + context.delete_global(global_name) + return event_generator + + def survey_phasespaces( + self, phasespaces: list[PhaseSpace | None], survey_pass: int = 0 + ) -> ms.EventGenerator | None: + ps_filtered = [ps for ps in phasespaces if ps is not None] + if len(ps_filtered) == 0: + return None + event_generator = self.build_event_generator(ps_filtered) + event_generator.survey(survey_pass) + return event_generator + + def survey(self) -> None: + # survey_pass distinguishes the survey() calls below: "both" mode can + # re-survey a channel carried over unchanged from the multichannel pass + # into the final (simplified) pass, and both passes schedule jobs on the + # same underlying ChannelEventGenerator. The explicit pass index keeps each + # pass's job seeds independent of the other passes' job counts, rather than + # depending on call history. + phasespace_mode = self.run_card["phasespace"]["mode"] + if self.is_decay and phasespace_mode != "multichannel": + # The flat mapping is built from a synthetic two-incoming diagram, + # which a decay has no counterpart for -- and 'both'/'auto' end up + # simplifying towards it. Decays have few channels, so the + # multichannel phase space is the right one regardless. + logger.info( + "decay mode: using the multichannel phase space instead of " + "'%s'", phasespace_mode + ) + phasespace_mode = "multichannel" + if phasespace_mode in ["multichannel", "both", "auto"]: + self.phasespaces = [ + subproc.build_multichannel_phasespace() + for subproc in self.subprocesses + ] + self.event_generator = self.survey_phasespaces(self.phasespaces, 0) + elif phasespace_mode == "flat": + self.phasespaces = [ + subproc.build_flat_phasespace() + for subproc in self.subprocesses + ] + self.event_generator = self.survey_phasespaces(self.phasespaces, 0) + else: + raise ValueError("Unknown phasespace mode") + + channel_status = self.event_generator.channel_status() + chan_offset = 0 + madnis_enabled = False + for subproc, ps in zip(self.subprocesses, self.phasespaces): + mean = 0. + variance = 0. + count_opt = 0 + for status in channel_status[chan_offset:chan_offset + len(ps.channels)]: + mean += status.mean_abs + variance += status.error_abs**2 + count_opt += status.count_opt + # A subprocess whose matrix element is identically zero (an FCNC + # channel with every Wilson coefficient at zero, say) integrates to + # exactly zero, so there is no relative spread to speak of. Report 0 + # rather than dividing by it: that asks for the smallest networks and + # leaves MadNIS off, which is what an empty subprocess wants. + rsd = (variance * count_opt)**0.5 / mean if mean else 0. + subproc.set_madnis_auto_settings(rsd) + chan_offset += len(ps.channels) + if subproc.madnis_settings["enable"]: + madnis_enabled = True + if not ( + phasespace_mode == "both" or (phasespace_mode == "auto" and madnis_enabled) + ): + return + + phasespaces_multi = self.phasespaces + cross_sections = [] + index = 0 + for phasespace in phasespaces_multi: + channel_count = len(phasespace.channels) + # rank the channels by the integral of |w|: a channel whose signed + # integral cancels (an interference) can still carry most of the + # variance, and folding it into the flat channel would blow it up + cross_sections.append([ + status.mean_abs + for status in channel_status[index:index + channel_count] + ]) + index += channel_count + + self.phasespaces = [ + subproc.simplify_phasespace(ps_multi, cross_secs) + for subproc, ps_multi, cross_secs in zip( + self.subprocesses, phasespaces_multi, cross_sections + ) + ] + if any( + ps_multi is not ps_both + for ps_multi, ps_both in zip(phasespaces_multi, self.phasespaces) + ): + # distinct survey_pass: a channel carried over unchanged from the + # multichannel pass (pass 0) into this resurvey must not share its + # seed stream with that earlier pass. + self.event_generator = self.survey_phasespaces(self.phasespaces, 1) + + def train_madnis(self) -> None: + madnis_args = self.run_card["madnis"] + if not any(subproc.madnis_settings["enable"] for subproc in self.subprocesses): + return + + gen_args = self.run_card["generation"] + run_args = self.run_card["run"] + + verbosity = resolve_verbosity(run_args["verbosity"]) + madnis_phasespaces = [] + training_args = [] + self.event_generator = None + for subproc, phasespace in zip(self.subprocesses, self.phasespaces): + for channel in phasespace.channels: + channel.event_generator = None + + config = ms.MadnisConfig() + config.learning_rate = subproc.madnis_settings["lr"] + config.batches = subproc.madnis_settings["train_batches"] + config.log_interval = madnis_args["log_interval"] + config.integration_history_length = madnis_args["integration_history_length"] + config.channel_dropping_interval = madnis_args["channel_dropping_interval"] + config.channel_dropping_threshold = madnis_args["channel_dropping_threshold"] + config.cpu_generator_batch_size = gen_args["cpu_batch_size"] + config.gpu_generator_batch_size = gen_args["gpu_batch_size"] + config.gpu_generator_batch_granularity = madnis_args["gpu_generator_batch_granularity"] + config.generator_target_size_factor = madnis_args["generator_target_size_factor"] + config.batch_size_offset = madnis_args["batch_size_offset"] + config.batch_size_per_channel = subproc.madnis_settings["batch_size_per_channel"] + config.uniform_channel_ratio = madnis_args["uniform_channel_ratio"] + config.lr_schedule = madnis_args["lr_scheduler"] + config.adam_beta1 = madnis_args["adam_beta1"] + config.adam_beta2 = madnis_args["adam_beta2"] + config.adam_eps = madnis_args["adam_eps"] + config.adam_weight_decay = madnis_args["adam_weight_decay"] + config.grad_clip_threshold = madnis_args["grad_clip_threshold"] + config.buffer_capacity = madnis_args["buffer_capacity"] + config.minimum_buffer_size = madnis_args["minimum_buffer_size"] + config.buffered_steps_fraction = madnis_args["buffered_steps_fraction"] + config.buffer_skip_batches = madnis_args["buffer_skip_batches"] + config.buffer_unweighting_quantile = madnis_args["buffer_unweighting_quantile"] + config.fixed_cwnet_fraction = subproc.madnis_settings["fixed_cwnet_fraction"] + config.softclip_threshold = madnis_args["softclip_threshold"] + config.compressed_channel_weight_count = madnis_args["compressed_channel_weight_count"] + phasespace = subproc.build_madnis(phasespace) + madnis_phasespaces.append(phasespace) + training_args.append( + ms.TrainingArgs( + config=config, + integrands=subproc.build_integrands( + phasespace, + madnis_training=True, + drop_cuts_and_rescale=madnis_args["drop_zero_integrands"] + ), + cwnet=phasespace.cwnet, + ) + ) + + gen_context = self.contexts[0] + opt_context = ms.Context( + device=self.devices[0], thread_count=self.pool_sizes[0] + ) + opt_context.copy_globals_from(gen_context) + + madnis_training = ms.MultiMadnisTraining( + generator_context=gen_context, + optimizer_context=opt_context, + training_args=training_args, + verbosity=verbosity, + status_file=self.status_file, + # Reuses the run's resolved seed (also used by build_event_generator()). + # Only the single-channel CPU sample-generation path is currently seeded + # -- buffered training and GPU multi-channel batches are still + # non-deterministic. + seed=self.run_seed, + ) + madnis_training.train() + for phasespace, active_channels in zip( + madnis_phasespaces, madnis_training.active_channels() + ): + phasespace.channels = [ + phasespace.channels[index] for index in active_channels + ] + del madnis_training + del opt_context + self.phasespaces = madnis_phasespaces + for context in self.contexts[1:]: + context.copy_globals_from(self.contexts[0]) + self.event_generator = self.build_event_generator(madnis_phasespaces) + self.madnis_trained = True + + def update_madnis_status_single( + self, batch: int, batch_target: int, loss: float, lr: float, channel_count: int + ) -> None: + now = time.time() + if batch + 1 < batch_target: + if now - self.last_update_time < 0.1: + return + self.last_update_time = now + progress_bar = ms.format_progress((batch + 1) / batch_target, 52) + time_diff = now - self.madnis_wall_time + time_str = f"{format_time(time_diff)}" + else: + progress_bar = "" + wall_diff = now - self.madnis_wall_time + cpu_diff = time.process_time() - self.madnis_cpu_time + time_str = ( + f"{format_time(wall_diff, centi=True)} wall, " + f"{format_time(cpu_diff, centi=True)} cpu" + ) + batch_str = f"{batch + 1} / {batch_target}" + self.madnis_box.set_column(1, [ + f"{batch_str:<15} {progress_bar}", + f"{loss:>.4f}", + f"{channel_count}", + time_str + ]) + self.madnis_box.print_update() + + def update_madnis_status_multi( + self, + subproc_id: int, + batch: int, + batch_target: int, + loss: float, + lr: float, + channel_count: int + ) -> None: + now = time.time() + subproc_count = len(self.subprocesses) + if batch + 1 < batch_target: + if now - self.last_update_time < 0.1: + return + self.last_update_time = now + progress_bar = ms.format_progress((batch + 1) / batch_target, 34) + progress_bar_all = ms.format_progress( + (subproc_id * batch_target + batch + 1) / (subproc_count * batch_target), + 52 + ) + time_str = f"{format_time(now - self.madnis_wall_time)}" + subproc_str = f"{subproc_id} / {subproc_count}" + elif subproc_id < subproc_count - 1: + progress_bar = "" + progress_bar_all = ms.format_progress( + ((subproc_id + 1) * batch_target + 1) / (subproc_count * batch_target), + 52 + ) + time_str = f"{format_time(now - self.madnis_wall_time)}" + subproc_str = f"{subproc_id} / {subproc_count}" + else: + progress_bar = "" + progress_bar_all = "" + wall_diff = now - self.madnis_wall_time + cpu_diff = time.process_time() - self.madnis_cpu_time + time_str = ( + f"{format_time(wall_diff, centi=True)} wall, " + f"{format_time(cpu_diff, centi=True)} cpu" + ) + subproc_str = f"{subproc_count} / {subproc_count}" + batch_str = f"{batch + 1} / {batch_target}" + self.madnis_upper_box.set_column(1, [ + f"{subproc_str:<15} {progress_bar_all}", + time_str, + ]) + self.madnis_lower_box.set_row(subproc_id + 1, [ + f"{subproc_id}", + f"{loss:>.4f}", + f"{channel_count}", + f"{batch_str:<15} {progress_bar}", + ]) + self.madnis_upper_box.print_update() + self.madnis_lower_box.print_update() + + @property + def gridpack_run_mode(self) -> str: + """How this run ends: a gridpack run may skip the events (see the + [gridpack] run_mode run card option); any other run is a regular one.""" + if not self.run_card["gridpack"]["save_gridpack"]: + return "regular" + return self.run_card["gridpack"]["run_mode"] + + def finish_without_events(self) -> None: + """The end of a gridpack run that writes no events: only what the + gridpack needs is prepared, then it is stored.""" + if self.gridpack_run_mode == "minimal": + # a trained MadNIS needs nothing more, a VEGAS grid has to converge + if not self.madnis_trained: + self.event_generator.optimize() + else: + self.event_generator.fix_max_weights() + self.build_systematics() + self.save_gridpack() + + def generate_events(self) -> None: + if self.gridpack_run_mode != "regular": + self.finish_without_events() + return + start_time = get_start_time() + self.event_generator.generate() + output_format = self.run_card["run"]["output_format"] + systematics = self.build_systematics() + histograms = self.build_event_histograms() + if output_format == "compact_npy": + self.lhe_completer = None + self.event_generator.combine_to_compact_npy( + os.path.join(self.run_path, "events.npy"), systematics, histograms + ) + self.write_lhe_header(systematics) + self.save_lhe_completer() + elif output_format == "lhe_npy": + self.lhe_completer = self.build_lhe_completer() + self.event_generator.combine_to_lhe_npy( + os.path.join(self.run_path, "events.npy"), self.lhe_completer, + systematics, histograms + ) + self.write_lhe_header(systematics) + elif output_format == "lhe": + self.lhe_completer = self.build_lhe_completer() + lhe_path = os.path.join(self.run_path, "events.lhe") + self.event_generator.combine_to_lhe( + lhe_path, self.lhe_completer, + self.build_lhe_meta(), systematics, histograms + ) + # Ship the LHE compressed by default. These files are large and + # very compressible, madevent has always stored its events + # gzipped, and every consumer here already accepts either form + # (see _find_event_file). misc.gzip replaces events.lhe with + # events.lhe.gz, and switches to an external multithreaded tool + # above 256 MB. + misc.gzip(lhe_path) + else: + raise ValueError("Unknown output format") + self.write_hwu(histograms) + self.make_plots(histograms) + if systematics is not None: + self.write_systematics_sidecar() + self.log_systematics_summary() + self.save_gridpack() + + def save_lhe_completer(self) -> None: + """Keep the LHE completer next to compact_npy events, so npy_to_lhe + can complete them into an LHE file later (a shower or MadSpin requested + after the run). Never fatal to the run.""" + try: + self.build_lhe_completer().save( + os.path.join(self.run_path, npy_to_lhe.COMPLETER_FILE)) + except Exception as err: + logger.warning("could not save the LHE completer; these events " + "cannot be converted to LHE later: %s", err) + + hwu_file_name = "MADatLO.HwU" + plot_dir_name = "plots" + + def make_plots(self, histograms) -> None: + """Draw the event-sample histograms into Events//plots, with the + scale and PDF bands info.json carries, and say where they landed. + + This does not go through the HwU file: the bands are already computed + per bin by madspace, so drawing them directly is both fewer steps and + more than histograms.py can do without the LHAPDF python module. + """ + if histograms is None or not self.run_card["run"]["make_plots"]: + return + out_dir = os.path.join(self.run_path, self.plot_dir_name) + try: + data = json.loads(histograms.to_json(self.systematics)) + written = plots.render(data, out_dir) + except plots.BackendMissing as error: + logger.info("matplotlib is not available (%s): no plots. The " + "distributions are in %s", + error, os.path.join(self.run_path, "info.json")) + return + except Exception as error: + logger.warning("could not draw the histograms: %s", error) + return + if not written: + return + logger.info("%d plot(s) stored in %s", len(written), out_dir) + + def write_hwu(self, histograms) -> None: + """Write the event-sample histograms next to the events in the HwU + format (MADatLO.HwU), the one an aMC@NLO run writes as MADatNLO.HwU + and madgraph/various/histograms.py reads: an mg7 distribution can then + be overlaid on an NLO one without converting anything in between. The + numbers are the ones info.json carries, bands included, so this is off + by default -- it is a second copy in another format, for the runs that + are going to be plotted.""" + if histograms is None or not self.run_card["run"]["write_hwu"]: + return + try: + data = json.loads(histograms.to_json(self.systematics)) + summary = (json.loads(self.systematics.summary()) + if self.systematics is not None else None) + text = hwu_output.to_hwu(data, summary) + except Exception as error: + logger.warning("could not write the HwU histograms: %s", error) + return + if not text: + return + with open(os.path.join(self.run_path, self.hwu_file_name), "w") as f: + f.write(text) + + def write_systematics_sidecar(self) -> None: + """Describe the variation weights next to the event file + (events.weights.json): ids, scale factors, PDF set/member and the + text, so npy consumers get the same metadata as the LHE + header, plus the per-variation cross sections.""" + if self.systematics is None: + return + data = json.loads(self.systematics.summary()) + data["initrwgt"] = self.systematics.initrwgt() + data["columns"] = ["rwgt_%d" % i for i in self.systematics.weight_ids] + with open(os.path.join(self.run_path, "events.weights.json"), "w") as f: + json.dump(data, f, indent=1) + + @staticmethod + def _histogram_mean(hist_min, hist_max, bin_values): + """Cross-section-weighted mean of a histogrammed observable over its + range. ``bin_values`` carries the underflow and overflow bins first and + last, as both madspace histogram kinds do; they have no position and + are left out.""" + values = list(bin_values)[1:-1] + total = sum(values) + if not values or total == 0: + return None + width = (hist_max - hist_min) / len(values) + return sum(v * (hist_min + (i + 0.5) * width) + for i, v in enumerate(values)) / total + + def _histogram_means(self) -> dict: + """name -> mean of every [histograms] entry, from the post-processing + histograms (the final event sample) when they were filled, otherwise + from the weighted integration histograms.""" + if self.event_histograms is not None: + data = json.loads(self.event_histograms.to_json(self.systematics)) + hists = [(h["name"], h["min"], h["max"], h["bin_values"]) + for h in data] + else: + hists = [(h.name, h.min, h.max, h.bin_values) + for h in self.event_generator.histograms()] + means = {} + for name, hist_min, hist_max, values in hists: + mean = self._histogram_mean(hist_min, hist_max, values) + if mean is not None: + means[name] = mean + return means + + def get_result(self) -> dict: + """Return the run result: cross-section (pb) with MC error, the number + of (unweighted) events, and the mean of every observable declared in the + [histograms] section. Used to build the scan summary.""" + status = self.event_generator.status() + result = {'cross(pb)': status.mean, 'error(pb)': status.error, + 'nb_event': status.count_unweighted} + try: + if self.systematics is not None and self.systematics.weight_count: + summary = json.loads(self.systematics.summary()) + scale = systematics_summary.scale_percentages(summary) + if scale is not None: + result['scale_up(%)'], result['scale_down(%)'] = scale + pdf = systematics_summary.pdf_percentages(summary) + if pdf: + # one column pair; the nominal set's entry comes first + _entry, up, down = pdf[0] + result['pdf_up(%)'], result['pdf_down(%)'] = up, down + except Exception as err: + logger.warning("could not extract the systematics summary: %s", err) + try: + for name, mean in self._histogram_means().items(): + result['<%s>' % name] = mean + except Exception as err: + logger.warning("could not extract observable means: %s", err) + return result + + def _beam_info(self): + """Return (beam_pdg_ids, beam_energies) for the LHE block. + + `incoming` holds the *partonic* initial state (e.g. gluons), not the + beam particle, so hadronic beams are protons (2212); leptonic beams are + the incoming leptons themselves.""" + half_e = float(self.e_cm) / 2. + if self.is_decay: + # No beams. LHE has no way to say that, so report the decaying + # particle at rest as a single "beam"; the second slot is empty. + data = self.subprocess_data[0] + pdg = clean_pids(data["incoming"])[0] + return [pdg, 0], [float(self.e_cm), 0.0] + if not self.leptonic: + # hadronic collider: proton beams (p-pbar is not distinguished) + return [2212, 2212], [half_e, half_e] + data = self.subprocess_data[0] + incoming = data["incoming"] + flavor0 = data["flavors"][0]["options"][0] + init_pdgs = flavor0[:len(incoming)] + beam_pdgs = [pdg if abs(code) in (81, 82) else code + for code, pdg in zip(incoming, init_pdgs)] + return beam_pdgs, [half_e, half_e] + + def _lhapdf_id(self): + """Central LHAPDF id of the beam PDF set (read from its .info SetIndex), + or -1 for a leptonic beam (no PDF).""" + if self.leptonic or not self.pdf_dir: + return -1 + pdf_set = self.run_card["beam"]["pdf"] + info = os.path.join(self.pdf_dir, pdf_set, "%s.info" % pdf_set) + try: + for line in open(info): + if line.strip().startswith("SetIndex:"): + return int(line.split(":", 1)[1].strip()) + except Exception as err: + logger.warning("could not read LHAPDF id from %s: %s", info, err) + return -1 + + def build_lhe_meta(self, systematics=None): + """Build the LHE header/ metadata: the param_card () and the + run_card.toml () headers plus the beam/PDF/cross-section info + needed by downstream tools (systematics, MadSpin, ...). + + `combine_to_lhe` injects the header itself, so callers that + go through it should leave `systematics` unset; pass it only when + building meta for a writer that bypasses that injection (e.g. + write_lhe_header).""" + beam_pdgs, energies = self._beam_info() + lhaid = self._lhapdf_id() + pdf_group = -1 if self.leptonic else 0 + status = self.event_generator.status() + xsec, err = status.mean, status.error + weight_mode, max_weight = lhe_weight_info(status) + with open(self.param_card_path) as f: + param_text = f.read() + with open(os.path.join("Cards", "run_card.toml")) as f: + run_text = f.read() + headers = [] + # generation commands (model + process): read by MadSpin/reweight/... + proc_card = os.path.join("Cards", "proc_card_mg5.dat") + if os.path.exists(proc_card): + with open(proc_card) as f: + headers.append(ms.LHEHeader(name="MG5ProcCard", content=f.read())) + headers.append(ms.LHEHeader(name="slha", content=param_text)) + headers.append(ms.LHEHeader(name="MG7RunCard", content=run_text)) + # The resolved seed (even when the run_card requested a random one via + # seed = -1), so the run can be reproduced from the LHE file alone. + headers.append(ms.LHEHeader(name="MG7Seed", content=str(self.run_seed))) + if systematics is not None and systematics.weight_ids: + headers.append(ms.LHEHeader( + name="initrwgt", content=systematics.initrwgt(), escape_content=False + )) + return ms.LHEMeta( + beam1_pdg_id=beam_pdgs[0], beam2_pdg_id=beam_pdgs[1], + beam1_energy=energies[0], beam2_energy=energies[1], + beam1_pdf_authors=pdf_group, beam2_pdf_authors=pdf_group, + beam1_pdf_id=lhaid, beam2_pdf_id=lhaid, + weight_mode=weight_mode, + # positional: the pybind arg name for max_weight is non-kwarg-safe + processes=[ms.LHEProcess(xsec, err, max_weight, 1)], + headers=headers, + ) + + def write_lhe_header(self, systematics) -> None: + """Write header.lhe next to events.npy: the
/ blocks + (run card, param card, beam/PDF, cross section) that the npy formats + otherwise drop, with no events.""" + header_path = os.path.join(self.run_path, "header.lhe") + writer = ms.LHEFileWriter(header_path, self.build_lhe_meta(systematics)) + del writer # closes the file (writes the closing tag) + + def build_lhe_completer(self): + all_mcdata = ( + [subproc.build_multi_channel_data() for subproc in self.subprocesses] + if self.merged_subprocess_data is None else + [build_multi_channel_data(meta, self) for meta in self.subprocess_data] + ) + subproc_args = [ + ms.SubprocArgs( + topologies = [topo[0] for topo in mcdata.topologies], + permutations = mcdata.permutations, + diagram_indices = mcdata.diagram_indices, + diagram_color_indices = mcdata.diagram_color_indices, + diagram_propagator_pdgs = mcdata.diagram_propagator_pdgs, + color_flows = meta["color_flows"], + pdg_color_types = { + int(key): value + for key, value in meta["pdg_color_types"].items() + }, + helicities = lhe_helicities(meta, self.run_card), + pdg_ids = [flavor["options"] for flavor in meta["flavors"]], + ) + for mcdata, meta in zip(all_mcdata, self.subprocess_data) + ] + return ms.LHECompleter( + subproc_args=subproc_args, + bw_cutoff=self.run_card["phasespace"]["bw_cutoff"] + ) + + def save_gridpack(self) -> None: + if not self.run_card["gridpack"]["save_gridpack"]: + return + + gridpack_path = os.path.join(self.run_path, "gridpack") + data_path = os.path.join(gridpack_path, "data") + events_path = os.path.join(gridpack_path, "Events") + os.mkdir(gridpack_path) + os.mkdir(data_path) + os.mkdir(events_path) + self.contexts[0].save_globals(os.path.join(data_path, "globals.tar")) + + # one json object per channel, joined by hand so that the channel + # json is not parsed and dumped again + channel_entries = [ + f"{json.dumps(channel.status().name)}:{channel.to_json(self.gridpack_run_mode == "fix_max_weight")}" + for channel in self.event_generator.channels() + ] + with open(os.path.join(data_path, "channels.json"), "w") as f: + f.write("{" + ",".join(channel_entries) + "}") + + lib_path = os.path.join(gridpack_path, "lib") + if self.run_card["gridpack"]["include_source"]: + os.mkdir(lib_path) + shutil.copytree("src", os.path.join(gridpack_path, "src")) + shutil.copytree("SubProcesses", os.path.join(gridpack_path, "SubProcesses")) + shutil.copytree("backend", os.path.join(gridpack_path, "backend")) + else: + shutil.copytree("lib", lib_path) + for _, path in self.custom_function_paths: + target = path if os.path.isabs(path) else os.path.join(gridpack_path, path) + if not os.path.exists(target): + os.makedirs(os.path.dirname(target), exist_ok=True) + shutil.copy(path, target) + + if self.run_card["gridpack"]["include_madspace_source"]: + shutil.copytree( + _MADSPACE_DIR, + os.path.join(gridpack_path, "madspace"), + ignore=shutil.ignore_patterns("build", "install"), + ) + + if self.run_card["gridpack"]["include_madspace"]: + shutil.copytree( + _INSTALL_DIR / "madspace", + os.path.join(gridpack_path, "madspace", "install", "madspace"), + ignore=shutil.ignore_patterns("include"), + ) + + matrix_elements = [] + for subproc in self.subprocess_data: + me_path = subproc["me_path"] + matrix_elements.append(me_path) + + cards_path = os.path.join(gridpack_path, "Cards") + os.mkdir(cards_path) + shutil.copy(os.path.join("Cards", "param_card.dat"), cards_path) + # Full run card with a header noting it is read-only in gridpack context. + import io as _io + _buf = _io.StringIO() + self.run_card.write(_buf) + _header = ( + "# This is the run card used to generate this gridpack.\n" + "# Modifying this file will have no effect on gridpack execution.\n" + "# To change event-generation settings, edit grid_run_card.toml.\n\n" + ) + with open(os.path.join(cards_path, "run_card.toml"), 'w') as _f: + _f.write(_header + _buf.getvalue()) + # Minimal card containing only the settings used by generate_events. + # The gridpack ships the libraries that were actually built, and their + # names carry the resolved backend, so the card it runs from has to name + # that backend as well: a cpu_mode of 'auto' would send the gridpack + # looking for a ..._auto.so that was never built. The full run_card.toml + # written just above keeps 'auto', since that is what was asked for. + device_names = self.run_card["run"]["device"] + if not isinstance(device_names, list): + device_names = [device_names] + resolved_cpu_mode = None + for name, backend in zip(device_names, getattr(self, "backends", [])): + if name.split(":")[0] == "cpu": + resolved_cpu_mode = backend + break + previous_cpu_mode = self.run_card["run"]["cpu_mode"] + if resolved_cpu_mode is not None: + self.run_card["run"]["cpu_mode"] = resolved_cpu_mode + # A gridpack is run to feed a shower or detector chain, which reads + # LHE, so it writes LHE by default whatever format this run used + # (bin/generate_events --output_format still selects an npy output). + previous_output_format = self.run_card["run"]["output_format"] + self.run_card["run"]["output_format"] = "lhe" + # the directories of this run mean nothing on the machine running the gridpack + previous_dirs = (self.run_card["gridpack"]["temp_output_dir"], + self.run_card["gridpack"]["output_dir"]) + self.run_card["gridpack"]["temp_output_dir"] = "" + self.run_card["gridpack"]["output_dir"] = "" + try: + self.run_card.write_gridpack_card( + os.path.join(cards_path, "grid_run_card.toml")) + finally: + self.run_card["run"]["cpu_mode"] = previous_cpu_mode + self.run_card["run"]["output_format"] = previous_output_format + (self.run_card["gridpack"]["temp_output_dir"], + self.run_card["gridpack"]["output_dir"]) = previous_dirs + + bin_path = os.path.join(gridpack_path, "bin") + os.mkdir(bin_path) + gen_events_file = os.path.join(bin_path, "generate_events") + template_dir = os.path.dirname(__file__) + shutil.copy(os.path.join(template_dir, "gridpack.py"), gen_events_file) + shutil.copy(os.path.join(template_dir, "custom_function_loader.py"), bin_path) + os.chmod(gen_events_file, 0o755) + shutil.copy(os.path.join(template_dir, "gridpack_setup.py"), bin_path) + # streaming needs the fixed maximum weights of this mode + if self.gridpack_run_mode == "fix_max_weight": + shutil.copy(os.path.join(template_dir, "event_stream.py"), bin_path) + + data = { + "matrix_elements": matrix_elements, + # loaded before the matrix elements, as in the run + "custom_functions": self.custom_function_paths, + "me_parameters": me_parameters(self.run_card), + "source_hash": ms.SOURCE_HASH, + } + with open(os.path.join(data_path, "data.json"), "w") as f: + json.dump(data, f) + + if self.lhe_completer is None: + self.lhe_completer = self.build_lhe_completer() + self.lhe_completer.save(os.path.join(data_path, "lhe.json")) + self.save_gridpack_lhe_meta(os.path.join(data_path, "lhe_meta.json")) + if self.systematics_data is not None: + systematics_data = dict(self.systematics_data) + # A decay takes alpha_s from a constant grid written into the run + # directory rather than from an LHAPDF set, so the gridpack cannot + # look that file up: ship it and point at the shipped copy. + if self.fixed_alphas_info and \ + systematics_data["nominal_info_file"] == self.fixed_alphas_info: + name = os.path.basename(self.fixed_alphas_info) + shutil.copy(self.fixed_alphas_info, os.path.join(data_path, name)) + systematics_data["nominal_info_file"] = os.path.join("data", name) + with open(os.path.join(data_path, "systematics.json"), "w") as f: + json.dump(systematics_data, f) + if self.run_card["gridpack"]["compress"]: + shutil.make_archive(gridpack_path, "gztar", self.run_path, "gridpack") + shutil.rmtree(gridpack_path) + + def save_gridpack_lhe_meta(self, path) -> None: + """The run-independent part of the LHE header (cards, beams, PDF), for + the gridpack to write complete LHE headers: it adds the cross section + and seed of its own run, which is why those two are left out here.""" + meta = self.build_lhe_meta() + data = {key: getattr(meta, key) for key in ( + "beam1_pdg_id", "beam2_pdg_id", "beam1_energy", "beam2_energy", + "beam1_pdf_authors", "beam2_pdf_authors", "beam1_pdf_id", + "beam2_pdf_id", "weight_mode")} + data["headers"] = [ + {"name": h.name, "content": h.content, + "escape_content": h.escape_content} + for h in meta.headers if h.name != "MG7Seed"] + with open(path, "w") as f: + json.dump(data, f) + + def get_mass(self, pid: int) -> float: + return self.param_card.get_value("mass", pid) + + def get_width(self, pid: int) -> float: + return self.param_card.get_value("width", pid) diff --git a/madgraph/iolibs/template_files/mg7/model_widths.py b/madgraph/iolibs/template_files/mg7/model_widths.py new file mode 100644 index 0000000000..f007473a2a --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/model_widths.py @@ -0,0 +1,155 @@ +"""The model of a process directory and its auto widths. Does not import madspace.""" + +import logging +import os +import re +import subprocess + +from madgraph.various import misc +from madgraph.iolibs.template_files.mg7.bootstrap import MG_ROOT as _MG_ROOT + +logger = logging.getLogger("madgraph7") + + +_model_hash_warned = set() +_model_cache = {} + + +def read_stored_model(me_dir=None) -> "str | None": + """The model this process was generated with, as recorded at output time in + ``SubProcesses/model.txt``: a UFO path carrying its restriction (e.g. + ``.../models/sm-no_b_mass``) or a model name -- in either case something + ``import model`` understands. Returns None when the process was written + without that information. + + Warns (once per directory) when the model on disk no longer matches the + hash stored on the second line, i.e. it changed since the output was made. + """ + base = me_dir or os.getcwd() + model_file = os.path.join(base, "SubProcesses", "model.txt") + if not os.path.exists(model_file): + return None + with open(model_file) as f: + lines = f.read().splitlines() + model = lines[0].strip() if lines else "" + stored_hash = lines[1].strip() if len(lines) > 1 else "" + if not model: + logger.warning("SubProcesses/model.txt is empty; the model of this " + "process is unknown.") + return None + + # verify the model on disk still matches the one used at output time. The + # reference may carry a restriction suffix ('sm-no_b_mass'); the hash was + # taken on the UFO directory itself ('sm'). + model_dir = model if os.path.isdir(model) else model.rsplit('-', 1)[0] + if stored_hash and os.path.isdir(model_dir) \ + and model_file not in _model_hash_warned: + current_hash = misc.hash_model_files(model_dir) + if current_hash and current_hash != stored_hash: + _model_hash_warned.add(model_file) + logger.warning( + "The model at %s has changed since this process was generated " + "(hash mismatch); anything computed from it now (the 'auto' " + "widths, the dependent masses/widths of the param_card) may be " + "inconsistent with the matrix element.", model_dir) + return model + + +def _proc_characteristic(me_dir=None) -> dict: + """SubProcesses/proc_characteristics as a plain dict (empty when the file is + missing or unreadable). Plain, because ProcCharacteristic is a ConfigFile, + whose get() is the parameter accessor and takes no default.""" + from madgraph.various import banner as _banner_mod + path = os.path.join(me_dir or os.getcwd(), 'SubProcesses', + 'proc_characteristics') + if os.path.exists(path): + try: + return dict(_banner_mod.ProcCharacteristic(path)) + except Exception as error: + logger.debug("could not read %s: %s", path, error) + return {} + + +def load_process_model(me_dir=None): + """Import the UFO model this process was generated with, as + :func:`models.import_ufo.import_model` builds it -- the object every + model-level consumer of the param_card expects (see + :meth:`MG7Cmd.get_model`). None when the process recorded no model, or when + that model no longer imports. + + Cached per (reference, complex-mass-scheme): importing a model is slow and + the card question asks for it again at every 'update dependent'.""" + ref = read_stored_model(me_dir) + if not ref: + return None + cms = _proc_characteristic(me_dir).get('complex_mass_scheme', False) + key = (ref, bool(cms)) + if key not in _model_cache: + import models.import_ufo as import_ufo + try: + with misc.MuteLogger(['madgraph.model'], [50]): + _model_cache[key] = import_ufo.import_model( + ref, complex_mass_scheme=bool(cms)) + except Exception as error: + # 'update dependent' calls this under a SIGALRM whose handler + # raises a TimeOutError of its own (a class defined inside the + # caller, so it can only be recognised by name). Let it through: + # the caller then says the model took too long to load and how to + # force it, instead of reporting a model that does not import. + if type(error).__name__ == 'TimeOutError': + raise + logger.debug("could not import the model %s: %s", ref, error) + _model_cache[key] = None + return _model_cache[key] + + +def compute_auto_widths(param_card_path=os.path.join("Cards", "param_card.dat")) -> None: + """Fill any width set to ``auto`` in the param_card, using madgraph and the + model stored at output time (``SubProcesses/model.txt``), and write the + result back into the card. A no-op when the card has no ``auto`` width. + + Called before every generation, so a single run and each scan point (whose + param_card is rewritten just before ``run_single``) both get model-computed + widths.""" + try: + with open(param_card_path) as f: + text = f.read() + except OSError: + return + # matches "DECAY auto" (optionally "auto@NLO"), as in + # common_run_interface.static_check_param_card + pdgs = re.findall(r"(?im)^\s*decay\s+([+-]?\d+)\s+auto", text) + if not pdgs: + return + pdgs = list(dict.fromkeys(pdgs)) # de-duplicate, keep order + + model = read_stored_model() + if not model: + logger.warning( + "The param_card requests 'auto' width(s) for %s but the model of " + "this process is unknown; leaving them as-is.", " ".join(pdgs)) + return + + mg5 = str(_MG_ROOT / "bin" / "madgraph") + if not os.path.exists(mg5): + logger.warning("Cannot find madgraph at %s; 'auto' widths not computed.", mg5) + return + + import tempfile + cmds = "import model %s\ncompute_widths %s --path=%s\n" % ( + model, " ".join(pdgs), os.path.abspath(param_card_path)) + with tempfile.NamedTemporaryFile("w", suffix=".mg5", delete=False) as fh: + fh.write(cmds) + cmdfile = fh.name + try: + logger.info("Computing 'auto' width(s) for %s ...", " ".join(pdgs)) + proc = subprocess.run([mg5, cmdfile]) + if proc.returncode != 0: + logger.warning( + "compute_widths returned a non-zero exit code; the param_card " + "may still contain 'auto' entries.") + finally: + try: + os.remove(cmdfile) + except OSError: + pass diff --git a/madgraph/iolibs/template_files/mg7/observables.py b/madgraph/iolibs/template_files/mg7/observables.py new file mode 100644 index 0000000000..27523e7d73 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/observables.py @@ -0,0 +1,178 @@ +"""Observables, cuts and histograms of the run card, and PDG id helpers.""" + +from dataclasses import dataclass + +from madgraph.iolibs.template_files.mg7.common import ms + +@dataclass +class CutItem: + observable_kwargs: dict + min: float + max: float + mode: str + + +@dataclass +class HistItem: + observable_kwargs: dict + min: float + max: float + bin_count: int + # the distribution of the event weight itself (the reserved "weight" key + # of [histograms]); it is read off the event record, not computed from + # the momenta, so it has no observable + name: str = "" + from_weight: bool = False + + +# Flavor-merged legs carry a group id instead of a pdg. Every member of a group +# shares the same mass -- that is what makes them mergeable -- so any member is +# a valid representative for the kinematics. +_MERGED_PID_REPRESENTATIVE = { + 81: 1, # light quarks d u s c + 82: 11, # charged leptons e mu + 83: 12, # neutrinos ve vm vt +} + + +def clean_pids(pids: list[int]) -> list[int]: + pids_out = [] + for pid in pids: + pid = abs(pid) + if pid in _MERGED_PID_REPRESENTATIVE: + pid = _MERGED_PID_REPRESENTATIVE[pid] + elif 81 <= pid <= 99: + # Reserved for flavor merging. Anything outside this window is a + # real pdg (BSM models use codes in the millions), so let it pass. + raise ValueError( + f"unknown flavor-merged particle id {pid}; add its " + "representative to _MERGED_PID_REPRESENTATIVE" + ) + pids_out.append(pid) + return pids_out + + +# observables that tell a particle from its antiparticle, and so need the signed +# pdg ids of the actual flavours rather than clean_pids' unsigned representatives +_SIGNED_OBSERVABLES = {"sfos_pair_mass"} + + +def _flavor_options(meta: dict, unmerged_meta: dict | None = None) -> list[list[int]]: + """Every flavour assignment (signed pdg ids, incoming then outgoing) a + subprocess evaluates with the same momenta.""" + if "flavors" in meta: + return [ + option for flavor in meta["flavors"] for option in flavor["options"] + ] + return [ + option + for subproc in meta.get("subprocesses", []) + for option in _flavor_options(unmerged_meta[subproc]) + ] + + +def observable_pids( + meta: dict, observable_kwargs: dict, unmerged_meta: dict | None = None +) -> list[int]: + """The pdg ids an observable of this subprocess is built with. + + Most observables only select particles by group, which the unsigned + representatives of clean_pids do. One that pairs a particle with its + antiparticle needs the real signed flavours: those of the first flavour + assignment, provided every assignment of the subprocess forms the same + pairs -- the cut is applied once to momenta that all of them share. + """ + all_pids = clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) + if observable_kwargs.get("observable") not in _SIGNED_OBSERVABLES: + return all_pids + options = _flavor_options(meta, unmerged_meta) + if not options: + raise ValueError( + f"{observable_kwargs.get('name')}: no flavour information for " + f"subprocess {meta['incoming']} > {meta['outgoing']}" + ) + n_in = len(meta["incoming"]) + groups = [set(group) for group in observable_kwargs["select_pids"]] + if len(groups) == 1: + groups = groups * 2 + + def sfos_pairs(pids): + out = range(n_in, len(pids)) + return frozenset( + frozenset((i, j)) + for i in out if pids[i] in groups[0] + for j in out if pids[j] in groups[1] + if i != j and pids[i] != 0 and pids[i] == -pids[j] + ) + + reference = sfos_pairs(options[0]) + if any(sfos_pairs(option) != reference for option in options[1:]): + raise ValueError( + f"{observable_kwargs.get('name')}: the flavours grouped into subprocess " + f"{meta['incoming']} > {meta['outgoing']} pair the particles " + "differently, so a same-flavour opposite-sign cut cannot be applied " + "to their common momenta; regenerate the process after " + "'set apply_flavor_grouping False'" + ) + return list(options[0]) + + +def build_observable(pids: list[int], n_in: int, observable_kwargs: dict): + """ms.Observable of a subprocess with `n_in` incoming particles. + + madspace takes the first two momenta to be the beams (ignore_incoming), so + for a 1 -> N decay it would drop the first decay product from every + selection. There the selection runs over all the momenta instead, with the + incoming particle set to pdg id 0, which is in no group. + """ + if n_in == 2: + return ms.Observable(pids, **observable_kwargs) + pids = [0] * n_in + list(pids[n_in:]) + return ms.Observable(pids, **dict(observable_kwargs, ignore_incoming=False)) + + +def decay_products(meta: dict, unmerged_meta: dict | None = None) -> set[int]: + """Outgoing positions (0-based, beams excluded) of the particles that come + from an on-shell (decay-chain) propagator, as MadEvent's check_decay marks + them. Every channel of a subprocess shares its decay chains, so the first + channel tells, read through its first diagram's permutation.""" + if not meta["channels"]: + return set() + channel = meta["channels"][0] + topo_channel = channel + if "vertices" not in channel: + topo_channel = unmerged_meta[channel["subprocess"]]["channels"][channel["channel"]] + on_shell = topo_channel.get("on_shell_propagators", []) + if not on_shell: + return set() + vertices = topo_channel["vertices"] + # the closing vertex creates no propagator + created = {vertex[-1]: vertex[:-1] for vertex in vertices[:-1]} + + def legs(name): + if name in created: + return set().union(*(legs(child) for child in created[name])) + return {name} + + n_in = len(meta["incoming"]) + n_out = len(meta["outgoing"]) + products = set() + for prop in on_shell: + external = legs(f"p{prop}") + incoming = {leg for leg in external if leg.startswith("i")} + outgoing = {int(leg[1:]) for leg in external if leg.startswith("o")} + if not incoming: + products |= outgoing + elif len(incoming) == n_in: + # built from the beams: the decay is everything else + products |= set(range(n_out)) - outgoing + # event[i] = topology[permutation[i]] + permutation = channel["diagrams"][0]["permutation"] + return { + i - n_in for i in range(n_in, len(permutation)) + if permutation[i] - n_in in products + } + + +def pid_is_qcd(pid: int): + return abs(pid) in [21, 1, 2, 3, 4, 5, 6, 81] diff --git a/madgraph/iolibs/template_files/mg7/options.py b/madgraph/iolibs/template_files/mg7/options.py new file mode 100644 index 0000000000..9225c37271 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/options.py @@ -0,0 +1,86 @@ +"""The MG5 options of the process directory and the LHAPDF paths they select.""" + +import os + +from madgraph.various import misc + +_ROOTED_OPTIONS = ('lhapdf', 'lhapdf_py3', 'lhapdf_py2', 'heptools_install_dir') + +def load_mg5_options(me_dir=None) -> dict: + """Read the tool paths from the MadGraph configuration, so the launcher + knows where LHAPDF and the optional programs (Pythia8/Delphes/MadSpin/ + reweight/analysis) live. + + Files are read least- to most-specific, each overriding the previous, which + is the same layering CommonRunCmd.set_configuration uses: the card inside + this process directory has the last word, then this installation, then the + per-user file. Relative values are resolved against the root of the file + they came from. + + ``me_dir`` is the process directory whose Cards/me5_configuration.txt has + the last word. It defaults to the working directory, which is right for + bin/generate_events (it chdirs there first) but not for an in-process + caller, which has not moved yet. + """ + + import madgraph + mg5dir = os.path.dirname(os.path.dirname(os.path.abspath(madgraph.__file__))) + me_dir = os.getcwd() if me_dir is None else me_dir + + options = { + 'pythia-pgs_path': None, 'pythia8_path': None, + 'madanalysis5_path': None, 'exrootanalysis_path': None, 'delphes_path': None, + 'rivet_path': None, 'contur_path': None, 'f2py_compiler': None, + 'lhapdf': None, 'lhapdf_py3': None, 'lhapdf_py2': None, 'timeout': 0, + 'heptools_install_dir': None, 'hepmc3_path': None, + } + config_files = [] + base_config = misc.base_config_file() + if base_config: + config_files.append((base_config, mg5dir)) + user_config = misc.user_config_file() + if user_config: + config_files.append((user_config, mg5dir)) + config_files.append((misc.install_config_file(mg5dir), mg5dir)) + config_files.append((os.path.join(me_dir, 'Cards', 'me5_configuration.txt'), + me_dir)) + for cfg, root in config_files: + if not os.path.exists(cfg): + continue + with open(cfg) as fsock: + for line in fsock: + line = line.split('#', 1)[0] + if '=' not in line: + continue + name, value = (x.strip() for x in line.split('=', 1)) + if name not in options or value in ('', 'None'): + continue + if (name.endswith('_path') or name in _ROOTED_OPTIONS) and \ + not os.path.isabs(value) and \ + (os.sep in value or value.startswith('.')): + value = os.path.normpath(os.path.join(root, value)) + options[name] = value + options['mg5_path'] = mg5dir # enables MadSpin/reweight + return options + + +_LHAPDF = None + +def lhapdf_paths(refresh: bool = False) -> misc.LhapdfPaths: + """The LHAPDF installation this run should use, resolved once. + + Uses the same helper as the 'launch' command, so driving a process + directory straight through bin/generate_events behaves identically to + launching it from MadGraph. + """ + + global _LHAPDF + if _LHAPDF is None or refresh: + _LHAPDF = misc.resolve_lhapdf(load_mg5_options()) + # Export what we found: the real LHAPDF library, systematics and the + # MadSpin/reweight subprocesses read these from the environment. + if _LHAPDF.data_paths: + os.environ["LHAPDF_DATA_PATH"] = os.pathsep.join(_LHAPDF.data_paths) + if _LHAPDF.config: + os.environ["MADGRAPH_LHAPDF_CONFIG"] = _LHAPDF.config + return _LHAPDF diff --git a/madgraph/iolibs/template_files/mg7/postprocess.py b/madgraph/iolibs/template_files/mg7/postprocess.py new file mode 100644 index 0000000000..38bb38cf76 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/postprocess.py @@ -0,0 +1,380 @@ +"""Post-processing tools run on the events, and their logging.""" + +import logging +import os +import sys +import time + +from madgraph.iolibs.template_files.mg7 import npy_to_lhe +from madgraph.iolibs.template_files.mg7.options import lhapdf_paths, load_mg5_options + +def _find_event_file(run_path): + """Locate the LHE event file produced by generate_events (if any).""" + for name in ("events.lhe", "events.lhe.gz"): + path = os.path.join(run_path, name) + if os.path.exists(path): + return path + return None + + +def _find_or_convert_event_file(run_path, log): + """The run's LHE event file, converting its npy events when that is all + the run wrote (see npy_to_lhe).""" + lhe_path = _find_event_file(run_path) + if lhe_path is None and npy_to_lhe.can_convert(run_path): + log.info("No LHE event file in %s: converting the npy events.", run_path) + try: + lhe_path = npy_to_lhe.convert(run_path) + except Exception as error: + _report_failure(log, "npy->LHE conversion", error, run_path) + return lhe_path + + +def _report_failure(log, what, error, directory=None): + """Log a post-processing failure and write the full traceback to a file + (whose path is printed) so the problem can be investigated.""" + import traceback + if not directory or not os.path.isdir(directory): + directory = os.getcwd() + path = os.path.join(directory, "%s_crash.log" % what.replace(" ", "_")) + try: + with open(path, "w") as fsock: + fsock.write(traceback.format_exc()) + except Exception: + path = None + log.warning("%s failed: %s", what, error) + if path: + log.warning("full traceback written to: %s", path) + + +def _off(value) -> bool: + """True when a switch value means the tool was not selected.""" + return value in (None, "OFF", "Not Avail.", "Not Avail. (numpy missing)") + + +def _has_effective_handler(lg) -> bool: + """True when ``lg`` (or an ancestor it propagates to) already has a handler, + i.e. some other party is already displaying these records.""" + while lg: + if lg.handlers: + return True + if not lg.propagate: + return False + lg = lg.parent + return False + + +_TOOL_LOGGING_READY = False + + +def _setup_logging(): + """Make the reused madevent tool drivers' progress visible on screen. + + The drivers already narrate what they do through logger.info / + update_status ("Running Pythia8 [arXiv:...]", "Splitting .lhe event + file...", the live Idle/Running/Completed job counters, "Running + MadSpin", ...), but the standalone mg7 launcher never attaches a handler + to the 'madgraph'/'madevent' loggers, so all of it is swallowed. Attach a + colored INFO console handler to them (idempotent).""" + global _TOOL_LOGGING_READY + if _TOOL_LOGGING_READY: + return + if _has_effective_handler(logging.getLogger("madgraph")): + # Someone already configured logging for us -- in practice MG5, which + # now runs the launch in its own process. Its configuration is + # authoritative (it owns stdout_level and the tutorial logger); adding + # a second handler here would print every line twice. + _TOOL_LOGGING_READY = True + return + try: + import madgraph.interface.coloring_logging # registers ColorFormatter + formatter = logging.ColorFormatter("%(message)s") + except Exception: + formatter = logging.Formatter("%(message)s") + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter(formatter) + handler._mg7_tool_handler = True + for name in ("madgraph", "madevent", "cmdprint", "madgraph7"): + lg = logging.getLogger(name) + if not any(getattr(h, "_mg7_tool_handler", False) for h in lg.handlers): + lg.addHandler(handler) + lg.setLevel(logging.INFO) + lg.propagate = False + _TOOL_LOGGING_READY = True + + +def selected_tools(switch) -> list: + """The post-processing drivers run_selected_tools runs for this switch.""" + # MadAnalysis5 hadron level analyses the shower/detector output, so it only + # makes sense when a shower ran (mirrors madevent's card gating: + # analysis == 'MadAnalysis5' and shower != 'OFF'). + ma5 = switch.get("analysis") == "MadAnalysis5" + showered = not _off(switch.get("shower")) + return [t for t, on in ( + ("reweighting", not _off(switch.get("reweight"))), + ("MadSpin", not _off(switch.get("madspin"))), + ("MadAnalysis5 (parton level)", ma5), + ("Pythia8 shower", switch.get("shower") == "Pythia8"), + ("Delphes", switch.get("detector") == "Delphes"), + ("MadAnalysis5 (hadron level)", ma5 and showered), + ("Rivet", switch.get("analysis") == "Rivet"), + ) if on] + + +def run_selected_tools(switch, process) -> None: + """Run the optional post-processing programs selected in the merged + question on the generated events. + + All the tools are driven through :class:`MG7RunCmd`, a thin adapter around + the standard madevent run interface (``MadEventCmd``). This means the mg7 + output reuses the *exact* madevent tool drivers -- in particular Pythia8 is + showered in parallel over the cluster/multicore backend, identically to a + madevent run -- instead of re-implementing each tool here. The command + sequence mirrors ``MadEventCmd.do_launch`` (reweight, MadSpin, MA5 parton, + shower, Delphes, MA5 hadron, Rivet); each tool checks its own card, so a + tool is invoked only when its switch is on. + """ + log = logging.getLogger("madevent") + + tools = selected_tools(switch) + if not tools: + # A switch that is not off does not necessarily select a driver here: + # "Not Avail." is not off, analysis = ExRoot has no mg7 driver, and a + # shower switch set to something other than Pythia8 selects nothing + # either. Without this a plain generate/output/launch announced a + # post-processing step with nothing after the colon, and paid for + # building the run interface (or an npy->LHE conversion) to do nothing. + return + + lhe_path = _find_or_convert_event_file(process.run_path, log) + if lhe_path is None: + log.warning("No LHE event file in %s; cannot run %s.", + process.run_path, ", ".join(tools)) + return + + run_name = os.path.basename(os.path.dirname(os.path.abspath(lhe_path))) + run_dir = os.path.dirname(os.path.abspath(lhe_path)) + + ma5 = switch.get("analysis") == "MadAnalysis5" + showered = not _off(switch.get("shower")) + + log.info("") + log.info("Post-processing the generated events with: %s", ", ".join(tools)) + + try: + from madgraph.iolibs.template_files.mg7.run_interface import MG7RunCmd + cmd = MG7RunCmd(os.getcwd(), load_mg5_options(), run_name, lhe_path) + cmd.set_run_name(run_name, None, "parton") + except Exception as error: + _report_failure(log, "post-processing setup", error, run_dir) + return + + def run(tool, command): + bar = "=" * 60 + log.info("") + log.info(bar) + log.info(" post-processing step: %s", tool) + log.info(bar) + start = time.time() + try: + cmd.exec_cmd(command, postcmd=False, printcmd=False) + except Exception as error: + _report_failure(log, tool, error, run_dir) + else: + log.info(" -> %s done (%.1fs)", tool, time.time() - start) + + # Same order as MadEventCmd.do_launch. The commands take no run name and so + # act on cmd.run_name: reweight runs on the generated events, then + # decay_events (MadSpin) repoints cmd.run_name to the "_decayed_i" run, + # so the shower and the analyses that follow act on the decayed events -- + # exactly as a madevent run does. Each driver still guards on its own card. + if not _off(switch.get("reweight")): + run("reweight", "reweight -from_cards") + if not _off(switch.get("madspin")): + run("MadSpin", "decay_events -from_cards") + if ma5: + run("MadAnalysis5 (parton)", "madanalysis5_parton --no_default") + if switch.get("shower") == "Pythia8": + run("Pythia8 shower", "shower --no_default") + if switch.get("detector") == "Delphes": + run("Delphes", "delphes --no_default") + # hadron-level MA5 needs the shower/detector output: only run it if a shower + # was requested (otherwise there is no hadron-level event file to analyse). + if ma5 and showered: + run("MadAnalysis5 (hadron)", "madanalysis5_hadron --no_default") + if switch.get("analysis") == "Rivet": + run("Rivet", "rivet --no_default") + + # Finalize the run exactly as MadEventCmd.do_launch does after the + # shower/detector/analysis tools: store_result() processes the deferred + # "to_store" actions -- in particular it gzips the Pythia8 HepMC to + # _pythia8_events.hepmc.gz (the pythia8_card default is HEPMCoutput:file + # = hepmc.gz). The deferred Rivet job reads exactly that .gz path, so this + # must run before the Rivet/Contur postprocessor below. + try: + cmd.store_result() + except Exception as error: + _report_failure(log, "store_result", error, run_dir) + + if switch.get("analysis") == "Rivet": + # do_rivet only *prepared* the run (wrote Events//run_rivet.sh) and + # deferred execution to the postprocessor -- the rivet_card default has + # run_rivet_later = True, so with --no_default it logged "Skipping Rivet + # for now, passing it to postprocessor" and appended the run to + # cmd.postprocessing_dirs. MadEventCmd.do_launch runs that postprocessor + # at the end of a run; do the same here so the .yoda (and, when + # run_contur = True, the Contur limits) are actually produced. + bar = "=" * 60 + log.info("") + log.info(bar) + log.info(" post-processing step: Rivet/Contur postprocessing") + log.info(bar) + start = time.time() + try: + cmd.postprocessing() + except Exception as error: + _report_failure(log, "Rivet-Contur postprocessing", error, run_dir) + else: + log.info(" -> Rivet/Contur postprocessing done (%.1fs)", + time.time() - start) + + +def _add_time_of_flight(lhe_path, threshold, param_card_path, log): + """Add invariant-lifetime (vtim) information to the LHE events, drawing each + unstable particle's decay length from its width (in mm). Mirrors + common_run_interface.CommonRunCmd.do_add_time_of_flight but reads the widths + from the local param_card instead of the LHE banner.""" + import random + try: + import madgraph.various.lhe_parser as lhe_parser + except ImportError: + import internal.lhe_parser as lhe_parser + from models import check_param_card as param_card_mod + from madgraph.various import misc as _misc + import madgraph.iolibs.files as files + + need_zip = lhe_path.endswith('.gz') + if need_zip: + _misc.gunzip(lhe_path) + lhe_path = lhe_path[:-3] + + param_card = param_card_mod.ParamCard(param_card_path) + cst = 6.58211915e-25 # hbar in GeV s + c = 299792458000 # speed of light in mm/s + log.info('Adding time of flight information on %s', lhe_path) + lhe = lhe_parser.EventFile(lhe_path) + out = open('%s_2vertex.lhe' % lhe_path, 'w') + out.write(lhe.banner) + for event in lhe: + for particle in event: + # default=0 -> particles without a decay entry are treated as stable + width = param_card['decay'].get((abs(particle.pid),), 0.).value + if width: + vtim = c * random.expovariate(width / cst) + if vtim > threshold: + particle.vtim = vtim + out.write(str(event)) + out.write('\n') + out.close() + lhe.close() + files.mv('%s_2vertex.lhe' % lhe_path, lhe_path) + if need_zip: + _misc.gzip(lhe_path) + + +def _lhapdf_config_path(): + """Path to lhapdf-config so systematics can import the python lhapdf + module (required to compute PDF/scale variations).""" + return lhapdf_paths().config + + +def _run_systematics(lhe_path, cfg, log): + """Run systematics.py (scale/PDF variations) on the LHE file, reusing the + exact worker madevent uses (systematics.call_systematics).""" + try: + import madgraph.various.systematics as systematics + except ImportError: + import internal.systematics as systematics + + def fmt(vals): + return ','.join(str(v) for v in vals) + + opts = [] + if cfg.get('systematics_mur'): + opts.append('--mur=%s' % fmt(cfg['systematics_mur'])) + if cfg.get('systematics_muf'): + opts.append('--muf=%s' % fmt(cfg['systematics_muf'])) + if cfg.get('systematics_pdf'): + opts.append('--pdf=%s' % fmt(cfg['systematics_pdf'])) + extra = cfg.get('systematics_str_options', '') + if extra: + opts.extend(extra.split()) + # The mg7 LHE has no block, so systematics cannot read the per-event + # LO reweighting info. For a process with a single alpha_s power (stored in + # SubProcesses/proc_characteristics at output time) it can be reconstructed + # from the events: pass that power as --lo_nqcd. -1 means the QCD power is + # not uniform, so the reconstruction is not applicable. + if not any(o.startswith('--lo_nqcd') for o in opts): + try: + from madgraph.various import banner as _banner_mod + except ImportError: + import internal.banner as _banner_mod + pc_path = os.path.join('SubProcesses', 'proc_characteristics') + if os.path.exists(pc_path): + try: + nqcd = int(_banner_mod.ProcCharacteristic(pc_path)['single_qcd_order']) + except Exception: + nqcd = -1 + if nqcd >= 0: + opts.append('--lo_nqcd=%d' % nqcd) + # tell systematics where to find lhapdf (so it can link the python module) + if not any(o.startswith('--lhapdf_config') for o in opts): + lhapdf_config = _lhapdf_config_path() + if lhapdf_config: + opts.append('--lhapdf_config=%s' % lhapdf_config) + + log.info('Running systematics on %s %s', lhe_path, ' '.join(opts)) + systematics.call_systematics([lhe_path, lhe_path] + opts, + log=lambda x: log.info(str(x))) + + +def run_lhe_postprocessing(process) -> None: + """Run the LHE-level post-processings configured in the run_card + [postprocessing] section on the generated event file (displaced-vertex + time-of-flight and systematics). Only applies when an LHE file exists.""" + lhe_path = _find_event_file(process.run_path) + if lhe_path is None: + return + # process.run_card is a RunCardMG7; the section view exposes .get(key, def) + try: + cfg = process.run_card["postprocessing"] + except Exception: + return + log = logging.getLogger('madevent') + + if cfg.get('systematics') and getattr(process, 'is_decay', False): + # Scale and PDF variations are a beam quantity. A decay has neither, so + # systematics can only fail here ("not supported for pdlabel=none") -- + # and it is not free: MadSpin reruns the launcher for every pool refill. + log.info("decay mode: skipping systematics (no beams to vary)") + elif cfg.get('systematics'): + native = getattr(process, 'systematics', None) + if native is not None and native.weight_count > 0: + log.warning("[postprocessing] systematics is ignored: the scale/PDF " + "weights were already computed by madspace ([systematics] " + "section); set systematics.enable = false to use the legacy " + "systematics.py path instead.") + else: + try: + _run_systematics(lhe_path, cfg, log) + except Exception as error: + _report_failure(log, "systematics computation", error, + os.path.dirname(lhe_path)) + + tof = cfg.get('time_of_flight', -1.0) + try: + if tof is not None and float(tof) >= 0: + _add_time_of_flight(lhe_path, float(tof), process.param_card_path, log) + except Exception as error: + _report_failure(log, "add_time_of_flight", error, + os.path.dirname(lhe_path)) diff --git a/madgraph/iolibs/template_files/mg7/run_interface.py b/madgraph/iolibs/template_files/mg7/run_interface.py index 4ba4a80ade..bd71bbd364 100644 --- a/madgraph/iolibs/template_files/mg7/run_interface.py +++ b/madgraph/iolibs/template_files/mg7/run_interface.py @@ -198,7 +198,7 @@ def get_model(self): tools this class drives (MadSpin, reweight, the auto widths and the 'update dependent' pass of their card questions) work here too. """ - from madgraph.iolibs.template_files.mg7.launch import load_process_model + from madgraph.iolibs.template_files.mg7.model_widths import load_process_model return load_process_model(self.me_dir) def do_set(self, line, log=True): diff --git a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py new file mode 100644 index 0000000000..8352894842 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py @@ -0,0 +1,1131 @@ +"""Subprocesses with madmatrix matrix elements and MG diagram based phase space.""" + +from __future__ import annotations + +import os +import re +from dataclasses import dataclass, field +from typing import Literal, NamedTuple, TYPE_CHECKING + +from madgraph.iolibs.template_files.mg7.common import logger, me_parameters, ms +from madgraph.iolibs.template_files.mg7.custom_function_loader import load_surrogate +from madgraph.iolibs.template_files.mg7.observables import build_observable, clean_pids, decay_products, observable_pids, pid_is_qcd + +if TYPE_CHECKING: + from madgraph.iolibs.template_files.mg7.madgraph_process import MadgraphProcess + +@dataclass +class Channel: + phasespace_mapping: ms.PhaseSpaceMapping + adaptive_mapping: ms.Flow | ms.VegasMapping + discrete_sym: ms.DiscreteSampler | ms.DiscreteFlow | None + discrete_flavor: ms.DiscreteSampler | ms.DiscreteFlow | None + channel_weight_indices: list[int] | None + name: str + active_flavors: list[int] + event_generator: ms.ChannelEventGenerator | None = None + + +@dataclass +class PhaseSpace: + mode: Literal["multichannel", "flat", "both"] + channels: list[Channel] + symfact: list[int | None] + first_chan_weight_remap: list[list[int]] = field(default_factory=list) + first_remapped_chan_count: int = 0 + second_chan_weight_remap: list[int] = field(default_factory=list) + second_remapped_chan_count: int = 0 + prop_chan_weights: ms.PropagatorChannelWeights | None = None + subchan_weights: ms.SubchannelWeights | None = None + cwnet: ms.ChannelWeightNetwork | None = None + + +class MultiChannelData(NamedTuple): + amp2_remaps: list[list[int]] + symfact: list[int | None] + topologies: list[list[ms.Topology]] + permutations: list[list[list[int]]] + channel_indices: list[list[int]] + channel_weight_indices: list[list[list[int]]] + diagram_indices: list[list[int]] + diagram_color_indices: list[list[list[int]]] + diagram_propagator_pdgs: list[list[list[int]]] + active_flavors: list[list[list[int]]] + qcd_s_channel_count: list[int] + + +def build_topologies( + incoming_masses: list[float], + outgoing_masses: list[float], + channel: dict, + process: MadgraphProcess +) -> list[ms.Topology]: + propagators = [] + for i, (pid, signed_pid) in enumerate(zip( + clean_pids(channel["propagators"]), channel["propagators"] + )): + mass = process.get_mass(pid) + width = process.get_width(pid) + bw_cutoff = process.run_card["phasespace"]["bw_cutoff"] + if i in channel["on_shell_propagators"]: + e_min = mass - bw_cutoff * width + e_max = mass + bw_cutoff * width + else: + e_min = 0 + e_max = 0 + # $-excluded propagator: the matrix element vanishes within bw_cutoff + # widths of the pole, so do not pile the channel's points up there + # (a flat density in that window instead of the Breit-Wigner peak). + # Channels written before the exporter recorded them have none. + flat_window = 0. + if i in channel.get("dollar_propagators", ()) and width > 0: + flat_window = float(bw_cutoff) + propagators.append(ms.Propagator( + mass=mass, + width=width, + integration_order=0, + e_min=e_min, + e_max=e_max, + pdg_id=signed_pid, + flat_window=flat_window, + )) + vertices = channel["vertices"] + diag = ms.Diagram( + incoming_masses, outgoing_masses, propagators, vertices + ) + return ms.Topology.topologies(diag) + + +def build_multi_channel_data( + meta: dict, process: MadgraphProcess, unmerged_meta: dict | None = None +) -> MultiChannelData: + incoming_masses = [ + process.get_mass(pid) for pid in clean_pids(meta["incoming"]) + ] + outgoing_masses = [ + process.get_mass(pid) for pid in clean_pids(meta["outgoing"]) + ] + + if unmerged_meta is None: + diagram_count = meta["diagram_count"] + amp2_remaps = [[-1] * diagram_count] + else: + amp2_remaps = [ + [-1] * unmerged_meta[subproc]["diagram_count"] + for subproc in meta["subprocesses"] + ] + symfact = [] + topologies = [] + permutations = [] + channel_indices = [] + channel_weight_indices = [] + diagram_indices = [] + diagram_color_indices = [] + diagram_propagator_pdgs = [] + active_flavors = [] + channel_index = 0 + qcd_s_channel_count = [] + + for channel in meta["channels"]: + if unmerged_meta is None: + topo_channel = channel + else: + topo_subproc = channel["subprocess"] + topo_channel_index = channel["channel"] + topo_channel = unmerged_meta[topo_subproc]["channels"][topo_channel_index] + chan_topologies = build_topologies( + incoming_masses, outgoing_masses, topo_channel, process + ) + topo_count = len(chan_topologies) + if topo_count == 0: + continue + + topo = chan_topologies[0] + s_chan_count = [0] * len(topo.decays) + non_qcd = [False] * len(topo.decays) + pdg_ids = [decay.pdg_id for decay in topo.decays] + for i, pid in zip(topo.outgoing_indices, meta["outgoing"]): + pdg_ids[i] = pid + for decay in reversed(topo.decays): + if len(decay.child_indices) == 0: + continue + if decay.index == 0 and topo.t_propagator_count > 0: + s_chan_count[0] = sum(s_chan_count[i] for i in decay.child_indices) + break + + is_qcd = pid_is_qcd(pdg_ids[decay.index]) and all( + pid_is_qcd(pdg_ids[i]) for i in decay.child_indices + ) + is_non_qcd = decay.on_shell or not is_qcd or any( + non_qcd[i] for i in decay.child_indices + ) + non_qcd[decay.index] = is_non_qcd + s_chan_count[decay.index] = ( + 0 if is_non_qcd else sum(s_chan_count[i] for i in decay.child_indices) + 1 + ) + qcd_s_channel_count.append(s_chan_count[0]) + + diagrams = channel["diagrams"] + chan_permutations = [d["permutation"] for d in diagrams] + if unmerged_meta is None: + amp2_remaps[0][diagrams[0]["diagram"]] = channel_index + else: + for amp2_remap, diag in zip(amp2_remaps, diagrams[0]["diagram"]): + if diag != -1: + amp2_remap[diag] = channel_index + + channel_index_first = channel_index + symfact_index_first = len(symfact) + channel_index += 1 + symfact.extend([None] * topo_count) + for d in diagrams[1:]: + if unmerged_meta is None: + amp2_remaps[0][d["diagram"]] = channel_index + else: + for amp2_remap, diag in zip(amp2_remaps, d["diagram"]): + if diag != -1: + amp2_remap[diag] = channel_index + channel_index += 1 + symfact.extend(range(symfact_index_first, symfact_index_first + topo_count)) + + topologies.append(chan_topologies) + permutations.append(chan_permutations) + channel_indices.append(list(range(channel_index_first, channel_index))) + channel_weight_indices.append([ + [ + symfact_index_first + topo_index + i * topo_count + for i in range(len(chan_permutations)) + ] + for topo_index in range(topo_count) + ]) + diagram_indices.append([d["diagram"] for d in diagrams]) + if unmerged_meta is None: + diagram_color_indices.append([d["active_colors"] for d in diagrams]) + diagram_propagator_pdgs.append( + [d["propagator_pdgs"] for d in diagrams] + ) + active_flavors.append([d["active_flavors"] for d in diagrams]) + + return MultiChannelData( + amp2_remaps, + symfact, + topologies, + permutations, + channel_indices, + channel_weight_indices, + diagram_indices, + diagram_color_indices, + diagram_propagator_pdgs, + active_flavors, + qcd_s_channel_count, + ) + + +class MadgraphSubprocess: + def __init__( + self, + process: MadgraphProcess, + meta: dict, + subproc_id: int, + unmerged_meta: dict | None = None + ): + self.process = process + self.meta = meta + self.subproc_id = subproc_id + self.multi_channel_data = None + + self.unmerged_meta = None + if unmerged_meta is None: + api_path_formats = [self.meta["me_path"]] + member_metas = [self.meta] + else: + api_path_formats = [] + member_metas = [] + for subproc in self.meta["subprocesses"]: + submeta = unmerged_meta[subproc] + api_path_formats.append(submeta["me_path"]) + member_metas.append(submeta) + if len(api_path_formats) == 1: + self.meta = submeta + else: + self.unmerged_meta = unmerged_meta + + # The libraries were all built up front by MadgraphProcess.compile_matrix_elements + all_api_paths = [ + [api_path_format.format(device=backend) for backend in self.process.backends] + for api_path_format in api_path_formats + ] + + self.incoming_masses = [ + self.process.get_mass(pid) for pid in clean_pids(self.meta["incoming"]) + ] + self.outgoing_masses = [ + self.process.get_mass(pid) for pid in clean_pids(self.meta["outgoing"]) + ] + self.particle_count = len(self.incoming_masses) + len(self.outgoing_masses) + # MadEvent's cut_decays = F: the decay products of on-shell propagators + # are left out of every cut selection (a pair cut then needs both + # members to be cut), while the histograms still see them + if self.process.run_card["phasespace"]["cut_decays"]: + uncut = set() + else: + uncut = decay_products(self.meta, self.unmerged_meta) + n_in = len(self.incoming_masses) + + def cut_pids(observable_kwargs): + pids = observable_pids(self.meta, observable_kwargs, self.unmerged_meta) + # pdg id 0 is in no group, so the particle is never selected + return [0 if i - n_in in uncut else pid for i, pid in enumerate(pids)] + + # custom cuts and scales see the PDG ids of the first flavor, like the + # built-in cuts, which select by flavor-independent groups + representative_pids = [self.flavor_pids()[0]] + custom_cut = self.process.custom_function( + "cut", self.particle_count, representative_pids + ) + if len(self.process.cut_data) > 0: + self.cuts = ms.Cuts([ + ms.CutItem( + observable=build_observable( + cut_pids(cut_item.observable_kwargs), + n_in, cut_item.observable_kwargs + ), + min=cut_item.min, + max=cut_item.max, + mode=cut_item.mode, + ) + for cut_item in self.process.cut_data + ], custom_cut) + elif custom_cut is not None: + self.cuts = ms.Cuts(self.particle_count, custom_cut) + else: + self.cuts = None + # the integration histograms are functions of the momenta; the weight + # distribution is a property of the final event sample, so it is only + # filled by MadgraphProcess.build_event_histograms at combine time + # They cost an observable evaluation per integration point, and the + # plots/HwU are drawn from the post-processing ones instead, so they + # are only filled on request ([run] weighted_histograms). + momentum_hists = [item for item in self.process.hist_data + if not item.from_weight] \ + if self.process.run_card["run"]["weighted_histograms"] else [] + self.histograms = ( + ms.ObservableHistograms([ + ms.HistItem( + observable=build_observable( + observable_pids( + self.meta, hist_item.observable_kwargs, self.unmerged_meta + ), + n_in, hist_item.observable_kwargs + ), + min=hist_item.min, + max=hist_item.max, + bin_count=hist_item.bin_count, + ) + for hist_item in momentum_hists + ]) + if momentum_hists + else None + ) + + self.scale = ms.EnergyScale( + particle_count=self.particle_count, + custom_scale=self.process.custom_function( + "scale", self.particle_count, representative_pids + ), + **self.process.scale_kwargs + ) + + if self.process.run_card["run"]["dummy_matrix_element"]: + self.matrix_elements = [None] * len(all_api_paths) + else: + self.matrix_elements = [] + for api_paths in all_api_paths: + for context, api_path in zip(self.process.contexts, api_paths): + mat = context.load_matrix_element( + api_path, self.process.param_card_path, + me_parameters(self.process.run_card) + ) + self.matrix_elements.append(mat) + self.surrogates = [] + if self.process.surrogate is not None: + for path_format, member_meta in zip(api_path_formats, member_metas): + self.surrogates.append(self.load_surrogate(path_format, member_meta)) + + def load_surrogate(self, me_path_format: str, member_meta: dict): + """The [surrogate] of one matrix element of the subprocess.""" + name = re.sub( + r"^libmadmatrix_(.*)_\{device\}\.so$", r"\1", os.path.basename(me_path_format) + ) + # PDG ids for each flavor index of the matrix element + flavors = member_meta["flavors"] + pdg_ids = [flavors[0]["options"][0]] * (max(f["index"] for f in flavors) + 1) + for flavor in flavors: + pdg_ids[flavor["index"]] = flavor["options"][0] + settings = self.process.surrogate + apis = load_surrogate( + ms, + self.process.contexts, + self.process.backends, + self.process.device_types, + settings["function"], + name, + pdg_ids, + settings["param_card"], + ) + if len({api.index() for api in apis}) != 1: + raise RuntimeError("surrogate loaded at different indices") + return apis[0] + + def flavor_pids(self) -> list[list[int]]: + """PDG ids of the external particles for every flavor of the + integrand, in the order of its flavor index.""" + pids = [] + for flav in self.meta["flavors"]: + if self.unmerged_meta is not None: + subproc_index = self.meta["subprocesses"][flav["subprocess"]] + flav = self.unmerged_meta[subproc_index]["flavors"][flav["flavor"]] + pids.append(flav["options"][0]) + return pids + + def build_multi_channel_data(self) -> MultiChannelData: + if self.multi_channel_data is not None: + return self.multi_channel_data + self.multi_channel_data = build_multi_channel_data( + self.meta, self.process, self.unmerged_meta + ) + return self.multi_channel_data + + def build_multichannel_phasespace(self) -> PhaseSpace: + mcdata = self.build_multi_channel_data() + channel_count = sum(len(topos) for topos in mcdata.topologies) + drop_threshold = self.process.run_card["phasespace"]["drop_qcd_s_channel"] + if drop_threshold >= 0 and channel_count > drop_threshold: + mcdata = self.drop_qcd_s_channels(mcdata) + mcdata = self.drop_empty_channels(mcdata) + + channels = [] + t_channel_mode = self.t_channel_mode( + self.process.run_card["phasespace"]["t_channel"] + ) + for channel_id, (chan_topologies, chan_permutations, chan_indices, active_flavors) in enumerate(zip( + mcdata.topologies, mcdata.permutations, mcdata.channel_weight_indices, + mcdata.active_flavors + )): + topo_count = len(chan_topologies) + for topo_index, (topo, indices) in enumerate(zip(chan_topologies, chan_indices)): + mapping = ms.PhaseSpaceMapping( + chan_topologies[0], + self.process.e_cm, + t_channel_mode=t_channel_mode, + cuts=self.cuts, + invariant_power=self.process.run_card["phasespace"]["invariant_power"], + permutations=chan_permutations, + leptonic=self.process.leptonic, + ) + prefix = f"subproc{self.subproc_id}.channel{channel_id}" + if topo_count > 1: + prefix += f".subchan{topo_index}" + discrete_sym, discrete_flavor = self.build_discrete( + len(chan_permutations), len(self.meta["flavors"]), prefix + ) + channels.append(Channel( + phasespace_mapping = mapping, + adaptive_mapping = self.build_vegas(mapping, prefix), + discrete_sym = discrete_sym, + discrete_flavor = discrete_flavor, + channel_weight_indices = indices, + name = f"{channel_id}", + active_flavors = active_flavors, + )) + + remapped_chan_count = sum( + len(indices) for indices in mcdata.channel_indices + ) + if self.process.run_card["phasespace"]["sde_strategy"] == "denominators": + prop_chan_weights = ms.PropagatorChannelWeights( + [topo[0] for topo in mcdata.topologies], mcdata.permutations, + mcdata.channel_indices + ) + chan_weight_remap = [] + else: + prop_chan_weights = None + chan_weight_remap = [ + [ + len(mcdata.symfact) if remap == -1 else remap + for remap in amp2_remap + ] + for amp2_remap in mcdata.amp2_remaps + ] + + if any(len(topos) > 1 for topos in mcdata.topologies): + subchan_weights = ms.SubchannelWeights( + mcdata.topologies, mcdata.permutations, mcdata.channel_indices + ) + else: + subchan_weights = None + + return PhaseSpace( + mode="multichannel", + channels=channels, + first_chan_weight_remap=chan_weight_remap, + first_remapped_chan_count=remapped_chan_count, + symfact=mcdata.symfact, + prop_chan_weights=prop_chan_weights, + subchan_weights=subchan_weights, + ) + + def build_flat_phasespace(self) -> PhaseSpace: + mapping = ms.PhaseSpaceMapping( + self.incoming_masses + self.outgoing_masses, + self.process.e_cm, + mode=self.t_channel_mode(self.process.run_card["phasespace"]["flat_mode"]), + cuts=self.cuts, + leptonic=self.process.leptonic, + ) + prefix = f"subproc{self.subproc_id}.flat" + discrete_sym, discrete_flavor = self.build_discrete( + 1, len(self.meta["flavors"]), prefix + ) + channel = Channel( + phasespace_mapping = mapping, + adaptive_mapping = self.build_vegas(mapping, prefix), + discrete_sym = discrete_sym, + discrete_flavor = discrete_flavor, + channel_weight_indices = [0], + name = "F", + active_flavors = [], + ) + if self.unmerged_meta is None: + remap = [list(range(self.meta["diagram_count"]))] + else: + remap = [ + list(range(self.unmerged_meta[subproc]["diagram_count"])) + for subproc in self.meta["subprocesses"] + ] + return PhaseSpace( + mode="flat", + channels=[channel], + first_chan_weight_remap=remap, + first_remapped_chan_count=1, + symfact=[None], + ) + + def simplify_phasespace( + self, + multi_phasespace: PhaseSpace, + cross_sections: list[float] + ) -> PhaseSpace | None: + assert multi_phasespace.mode == "multichannel" + + threshold = 1 - self.process.run_card["phasespace"]["combine_channel_threshold"] + kept_channels = [] + tot_cs = sum(cross_sections) + cum_cs = 0. + seen_active_flavors = set() + #seen_resonances = set() + for index, (cs, chan) in sorted( + enumerate(zip(cross_sections, multi_phasespace.channels)), + key=lambda pair: pair[1][0], + reverse=True + ): + cum_cs += cs + has_unseen_flavors = False + has_unseen_resonances = False + for flavs in chan.active_flavors: + for flav in flavs: + if flav not in seen_active_flavors: + has_unseen_flavors = True + seen_active_flavors.add(flav) + #for resonance in chan.resonances: + # if resonance not in seen_resonances: + # has_unseen_resonances = True + # seen_resonances.add(flav) + if has_unseen_flavors or has_unseen_resonances or not tot_cs \ + or cum_cs / tot_cs < threshold: + kept_channels.append(index) + if len(kept_channels) >= len(cross_sections) - 1: + return multi_phasespace + + channels = [] + channel_map = {} + symfact = [] + for old_chan_index in kept_channels: + channel = multi_phasespace.channels[old_chan_index] + perm_count = max(1, channel.phasespace_mapping.channel_count()) + channel_index = len(symfact) + symfact.append(None) + symfact.extend([channel_index] * (perm_count - 1)) + channel_map.update({ + old_index: new_index + for new_index, old_index in enumerate( + channel.channel_weight_indices, start=channel_index + ) + }) + channels.append(Channel( + phasespace_mapping = channel.phasespace_mapping, + adaptive_mapping = channel.adaptive_mapping, + discrete_sym = channel.discrete_sym, + discrete_flavor = channel.discrete_flavor, + channel_weight_indices = list(range( + channel_index, channel_index + perm_count + )), + name = channel.name, + active_flavors = channel.active_flavors, + event_generator = channel.event_generator, + )) + + flat_phasespace = self.build_flat_phasespace() + flat_channel = flat_phasespace.channels[0] + channels.append(Channel( + phasespace_mapping = flat_channel.phasespace_mapping, + adaptive_mapping = flat_channel.adaptive_mapping, + discrete_sym = flat_channel.discrete_sym, + discrete_flavor = flat_channel.discrete_flavor, + channel_weight_indices = [len(symfact)], + name = flat_channel.name, + active_flavors = flat_channel.active_flavors, + )) + flat_index = len(symfact) + symfact.append(None) + channel_map[len(multi_phasespace.symfact)] = len(symfact) + if multi_phasespace.subchan_weights is None and len(multi_phasespace.first_chan_weight_remap) > 0: + first_chan_weight_remap = [ + [ + channel_map.get(remap, flat_index) + for remap in cw_remap + ] + for cw_remap in multi_phasespace.first_chan_weight_remap + ] + first_remapped_chan_count = len(symfact) + second_chan_weight_remap = [] + second_remapped_chan_count = 0 + else: + first_chan_weight_remap = multi_phasespace.first_chan_weight_remap + first_remapped_chan_count = multi_phasespace.first_remapped_chan_count + chan_count = multi_phasespace.first_remapped_chan_count if multi_phasespace.subchan_weights is None else multi_phasespace.subchan_weights.channel_count() + second_chan_weight_remap = [ + channel_map.get(i, flat_index) + for i in range(chan_count) + ] + second_remapped_chan_count = len(symfact) + + return PhaseSpace( + mode="both", + channels=channels, + first_chan_weight_remap=first_chan_weight_remap, + first_remapped_chan_count=first_remapped_chan_count, + second_chan_weight_remap=second_chan_weight_remap, + second_remapped_chan_count=second_remapped_chan_count, + symfact=symfact, + prop_chan_weights=multi_phasespace.prop_chan_weights, + subchan_weights=multi_phasespace.subchan_weights, + ) + + def drop_qcd_s_channels(self, mcdata: MultiChannelData) -> MultiChannelData: + """Drop channels with non-resonant QCD s-channel propagators, to reduce the + channel count for processes with many diagrams. Channels are grouped by the + number of QCD s-channels, allowing for more if necessary to map out all flavor + indices. Channel weights belonging to a dropped channel are left unmapped.""" + groups_by_s_count = {} + for index, count in enumerate(mcdata.qcd_s_channel_count): + groups_by_s_count.setdefault(count, []).append(index) + + covered_flavors = set() + kept_groups = set() + for s_count in sorted(groups_by_s_count): + s_count_groups = groups_by_s_count[s_count] + if s_count == 0: + selected = s_count_groups + else: + selected = [ + index + for index in s_count_groups + if any( + flav not in covered_flavors + for flavs in mcdata.active_flavors[index] + for flav in flavs + ) + ] + kept_groups.update(selected) + for index in selected: + for flavs in mcdata.active_flavors[index]: + covered_flavors.update(flavs) + + return self.keep_channel_groups(mcdata, sorted(kept_groups)) + + def drop_empty_channels(self, mcdata: MultiChannelData) -> MultiChannelData: + """Drop the channels whose phase space the cuts exclude entirely: a pair + mass cut whose floor lies at or above the top of an on-shell window of + the channel (ms.PhaseSpaceMapping.empty). Every point such a channel + generates fails the cuts, so it could only ever fail to find one. + + If that leaves a flavour with no channel at all, its cross section is + zero -- typically a cut that excludes the window of a decay chain every + diagram shares -- and the run cannot sample it: that is reported as an + error naming the cause, rather than as the survey finding no point + passing the cuts.""" + if self.cuts is None: + return mcdata + t_channel_mode = self.t_channel_mode( + self.process.run_card["phasespace"]["t_channel"] + ) + empty = [ + ms.PhaseSpaceMapping( + chan_topologies[0], + self.process.e_cm, + t_channel_mode=t_channel_mode, + cuts=self.cuts, + invariant_power=self.process.run_card["phasespace"]["invariant_power"], + permutations=chan_permutations, + leptonic=self.process.leptonic, + ).empty() + for chan_topologies, chan_permutations in zip( + mcdata.topologies, mcdata.permutations + ) + ] + if not any(empty): + return mcdata + kept_groups = [index for index, is_empty in enumerate(empty) if not is_empty] + covered_flavors = { + flav + for index in kept_groups + for flavs in mcdata.active_flavors[index] + for flav in flavs + } + uncovered = sorted({ + flav + for index, is_empty in enumerate(empty) if is_empty + for flavs in mcdata.active_flavors[index] + for flav in flavs + } - covered_flavors) + if uncovered: + raise ValueError( + f"subprocess {self.meta['incoming']} > {self.meta['outgoing']}: " + "the pair mass cuts exclude the whole on-shell window " + f"(bw_cutoff = {self.process.run_card['phasespace']['bw_cutoff']}) " + f"of every channel of flavour(s) {uncovered}, so their cross " + "section is zero; loosen the cut or remove that decay" + ) + logger.info( + "subprocess %d: %d channel(s) dropped, the cuts exclude the " + "on-shell window they sample", self.subproc_id, len(empty) - len(kept_groups) + ) + return self.keep_channel_groups(mcdata, kept_groups) + + def keep_channel_groups( + self, mcdata: MultiChannelData, kept_groups: list[int] + ) -> MultiChannelData: + """Rebuild the multi-channel data with only the channel groups (indices + into mcdata.topologies) in kept_groups. Channel weights belonging to a + dropped channel are left unmapped.""" + if len(kept_groups) == len(mcdata.topologies): + return mcdata + + amp2_remaps = [[-1] * len(remap) for remap in mcdata.amp2_remaps] + symfact = [] + topologies = [] + permutations = [] + channel_indices = [] + channel_weight_indices = [] + diagram_indices = [] + diagram_color_indices = [] + diagram_propagator_pdgs = [] + active_flavors = [] + qcd_s_channel_count = [] + channel_index = 0 + + for group in kept_groups: + chan_topologies = mcdata.topologies[group] + chan_permutations = mcdata.permutations[group] + chan_diagram_indices = mcdata.diagram_indices[group] + topo_count = len(chan_topologies) + + channel_index_first = channel_index + symfact_index_first = len(symfact) + for i, diag in enumerate(chan_diagram_indices): + if self.unmerged_meta is None: + amp2_remaps[0][diag] = channel_index + else: + for amp2_remap, d in zip(amp2_remaps, diag): + if d != -1: + amp2_remap[d] = channel_index + if i == 0: + symfact.extend([None] * topo_count) + else: + symfact.extend(range(symfact_index_first, symfact_index_first + topo_count)) + channel_index += 1 + + topologies.append(chan_topologies) + permutations.append(chan_permutations) + channel_indices.append(list(range(channel_index_first, channel_index))) + channel_weight_indices.append([ + [ + symfact_index_first + topo_index + i * topo_count + for i in range(len(chan_permutations)) + ] + for topo_index in range(topo_count) + ]) + diagram_indices.append(chan_diagram_indices) + if mcdata.diagram_color_indices: + diagram_color_indices.append(mcdata.diagram_color_indices[group]) + if mcdata.diagram_propagator_pdgs: + diagram_propagator_pdgs.append( + mcdata.diagram_propagator_pdgs[group] + ) + active_flavors.append(mcdata.active_flavors[group]) + qcd_s_channel_count.append(mcdata.qcd_s_channel_count[group]) + + return MultiChannelData( + amp2_remaps, + symfact, + topologies, + permutations, + channel_indices, + channel_weight_indices, + diagram_indices, + diagram_color_indices, + diagram_propagator_pdgs, + active_flavors, + qcd_s_channel_count, + ) + + def set_madnis_auto_settings(self, rsd: float): + madnis_args = self.process.run_card["madnis"] + n_out = len(self.meta["outgoing"]) + n_events = self.process.run_card["generation"]["events"] + is_gridpack = self.process.run_card["gridpack"]["save_gridpack"] + train_batches = madnis_args["train_batches"] + hidden_dim = min(max(int((7 * rsd) / 32) * 32 + 64, 64), 256) + flow_layers = 3 if rsd < 32 else 4 + lr = min(max((10000 - train_batches) / 8000 * 7e-4 + 3e-4, 3e-4), 1e-3) + if n_out <= 2: + enable = False + elif n_out == 3: + enable = rsd > 10. or n_events > 1000000 or is_gridpack + else: + enable = True + fixed_cwnet_fraction = max(0.33, 1.0 - 10000. / train_batches) + batch_size_per_channel = min(max(int((7 * rsd) / 32) * 32 + 64, 128), 512) + + self.madnis_settings = { + "enable": enable, + "flow_layers": flow_layers, + "flow_hidden_dim": hidden_dim, + "discrete_hidden_dim": hidden_dim, + "cwnet_hidden_dim": hidden_dim, + "train_batches": train_batches, + "lr": lr, + "fixed_cwnet_fraction": fixed_cwnet_fraction, + "batch_size_per_channel": batch_size_per_channel, + } + for key, value in self.madnis_settings.items(): + run_card_value = madnis_args[key] + if run_card_value != "auto": + self.madnis_settings[key] = run_card_value + + def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: + madnis_args = self.process.run_card["madnis"] + # Shared across all networks below: initialize_globals() derives an + # independent, non-colliding stream per tensor from this one base seed. + seed = self.process.run_seed + channels = [] + for channel_id, channel in enumerate(phasespace.channels): + prefix = f"subproc{self.subproc_id}.channel{channel_id}" + cond_dim = 0 + + flow_dim = channel.phasespace_mapping.random_dim() + flow = ms.Flow( + input_dim=flow_dim, + condition_dim=cond_dim, + prefix=prefix, + bin_count=madnis_args["flow_spline_bins"], + subnet_hidden_dim=self.madnis_settings["flow_hidden_dim"], + subnet_layers=self.madnis_settings["flow_layers"], + subnet_activation=self.activation(madnis_args["flow_activation"]), + invert_spline=madnis_args["flow_invert_spline"], + ) + if channel.adaptive_mapping is None: + flow.initialize_globals(self.process.contexts[0], seed) + else: + flow.initialize_from_vegas( + self.process.contexts[0], channel.adaptive_mapping.grid_name(), + seed + ) + cond_dim += flow_dim + + # discrete_sym runs after the adaptive map, so it can condition on its latent. + discrete_sym = channel.discrete_sym + if discrete_sym is not None: + perm_count = channel.phasespace_mapping.channel_count() + discrete_sym = ms.DiscreteFlow( + option_counts=[perm_count], + prefix=f"{prefix}.discrete_flow_sym", + dims_with_prior=[], + condition_dim=cond_dim, + subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], + subnet_layers=madnis_args["discrete_layers"], + subnet_activation=self.activation(madnis_args["discrete_activation"]), + ) + discrete_sym.initialize_globals(self.process.contexts[0], seed) + cond_dim += perm_count + + discrete_flavor = channel.discrete_flavor + if discrete_flavor is not None: + discrete_flavor = ms.DiscreteFlow( + option_counts=[len(self.meta["flavors"])], + prefix=f"{prefix}.discrete_flow_flavor", + dims_with_prior=[0], + condition_dim=cond_dim, + subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], + subnet_layers=madnis_args["discrete_layers"], + subnet_activation=self.activation(madnis_args["discrete_activation"]), + ) + discrete_flavor.initialize_globals(self.process.contexts[0], seed) + + channels.append(Channel( + phasespace_mapping = channel.phasespace_mapping, + adaptive_mapping = flow, + discrete_sym = discrete_sym, + discrete_flavor = discrete_flavor, + channel_weight_indices = channel.channel_weight_indices, + name = channel.name, + active_flavors = channel.active_flavors, + )) + + return PhaseSpace( + mode="both", + channels=channels, + first_chan_weight_remap=phasespace.first_chan_weight_remap, + first_remapped_chan_count=phasespace.first_remapped_chan_count, + second_chan_weight_remap=phasespace.second_chan_weight_remap, + second_remapped_chan_count=phasespace.second_remapped_chan_count, + symfact=phasespace.symfact, + cwnet=self.build_cwnet(len(phasespace.symfact)), + prop_chan_weights=phasespace.prop_chan_weights, + subchan_weights=phasespace.subchan_weights, + ) + + def build_vegas(self, mapping: ms.PhaseSpaceMapping, prefix: str) -> ms.VegasMapping: + if not self.process.run_card["vegas"]["enable"]: + return None + + vegas = ms.VegasMapping( + mapping.random_dim(), + self.process.run_card["vegas"]["bins"], + prefix, + ) + for context in self.process.contexts: + vegas.initialize_globals(context) + return vegas + + def build_discrete( + self, permutation_count: int, flavor_count: int, prefix: str + ) -> tuple[ms.DiscreteSampler | None, ms.DiscreteSampler | None]: + is_adaptive = self.process.run_card["phasespace"]["adaptive_symmetry_sampling"] + if is_adaptive and permutation_count > 1: + discrete_sym = ms.DiscreteSampler( + [permutation_count], f"{prefix}.discrete_sym" + ) + for context in self.process.contexts: + discrete_sym.initialize_globals(context) + else: + discrete_sym = None + + if flavor_count > 1: + # the flavor probabilities are conditioned on the PDF prior; a + # leptonic process (or a decay) has none, and the Integrand then + # passes no condition: declaring one anyway fails to build + # ("keys and values must have the same size"), e.g. z > q q~ with + # both u- and d-type rows + discrete_flavor = ms.DiscreteSampler( + [flavor_count], f"{prefix}.discrete_flavor", + [] if self.process.leptonic else [0] + ) + for context in self.process.contexts: + discrete_flavor.initialize_globals(context) + else: + discrete_flavor = None + + return discrete_sym, discrete_flavor + + def build_cwnet(self, channel_count: int) -> ms.ChannelWeightNetwork: + #if channel_count == 1: + # return None + madnis_args = self.process.run_card["madnis"] + cwnet = ms.ChannelWeightNetwork( + channel_count=channel_count, + particle_count=self.particle_count, + hidden_dim=self.madnis_settings["cwnet_hidden_dim"], + layers=madnis_args["cwnet_layers"], + activation=self.activation(madnis_args["cwnet_activation"]), + prefix=f"subproc{self.subproc_id}.cwnet", + ) + cwnet.initialize_globals( + self.process.contexts[0], self.process.run_seed + ) + return cwnet + + def t_channel_mode(self, name: str) -> ms.PhaseSpaceMapping.TChannelMode: + modes = { + "propagator": ms.PhaseSpaceMapping.propagator, + "rambo": ms.PhaseSpaceMapping.rambo, + "chili": ms.PhaseSpaceMapping.chili, + } + if name in modes: + return modes[name] + else: + raise ValueError(f"Invalid t-channel mode '{name}'") + + def activation(self, name: str) -> ms.MLP.Activation: + activations = { + "relu": ms.MLP.relu, + "leaky_relu": ms.MLP.leaky_relu, + "elu": ms.MLP.elu, + "gelu": ms.MLP.gelu, + "sigmoid": ms.MLP.sigmoid, + "softplus": ms.MLP.softplus, + } + if name in activations: + return activations[name] + else: + raise ValueError(f"Invalid activation function '{name}'") + + def build_integrands( + self, + phasespace: PhaseSpace, + madnis_training: bool = False, + drop_cuts_and_rescale: bool = False + ) -> list[ms.Integrand]: + flavors = [] + flavor_remap = [] + flavor_factors = [] + flavor_mirror = [] + flavor_diff_xs_indices = [] + flavor_subproc_indices = [] + flavor_per_subproc_remap = [] + + for flav in self.meta["flavors"]: + if self.unmerged_meta is not None: + diff_xs_index = flav["subprocess"] + subproc_index = self.meta["subprocesses"][diff_xs_index] + ps_flavor = flav["flavor"] + flavor_diff_xs_indices.append(diff_xs_index) + flavor_subproc_indices.append(subproc_index) + flavor_per_subproc_remap.append(ps_flavor) + flav = self.unmerged_meta[subproc_index]["flavors"][ps_flavor] + flavors.append(flav["options"][0]) + flavor_remap.append(flav["index"]) + flavor_factors.append(len(flav["options"])) + flavor_mirror.append(flav["mirror"]) + + cross_sections = [] + self.step_matrix_elements = [] + for me_index, matrix_element in enumerate(self.matrix_elements): + if self.surrogates: + # the surrogate in the integrand, the full matrix element in + # the second pass, with the random numbers passed in + mat = ms.MatrixElement( + self.surrogates[me_index].index(), + self.particle_count, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in], + [ms.MatrixElement.matrix_element_out], + 1, + False, + self.process.me_frame, + self.process.incoming_count, + ) + self.step_matrix_elements.append(ms.MatrixElement( + matrix_element, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in, ms.MatrixElement.random_color_in, + ms.MatrixElement.random_helicity_in, + ms.MatrixElement.random_diagram_in], + [ms.MatrixElement.matrix_element_out, + ms.MatrixElement.color_index_out, + ms.MatrixElement.helicity_index_out, + ms.MatrixElement.diagram_index_out], + False, + self.process.me_frame, + self.process.incoming_count, + )) + elif matrix_element: + mat = ms.MatrixElement( + matrix_element, + ms.Integrand.matrix_element_inputs, + ms.Integrand.matrix_element_outputs, + True, + self.process.me_frame, + self.process.incoming_count, + ) + else: + #TODO: not working in merged mode + mat = ms.MatrixElement( + 0xBADCAFE, + self.particle_count, + ms.Integrand.matrix_element_inputs, + ms.Integrand.matrix_element_outputs, + self.meta["diagram_count"], + True, + self.process.me_frame, + self.process.incoming_count, + ) + pdf_grid = None if self.process.leptonic else self.process.pdf_grid + pdf_arg = None if self.process.leptonic else ms.CachedPdf() + cross_sections.append( + ms.DifferentialCrossSection( + matrix_element=mat, + # For a decay this is the decaying particle's mass, and the + # flux becomes 1/(2M): the result is a partial width in + # GeV, not a cross section in pb. + cm_energy=self.process.e_cm, + running_coupling=None, + energy_scale=ms.CachedScale(), + pid_options=[], + pdf1=pdf_arg, + pdf2=pdf_arg, + input_momentum_fraction=True, + decay=self.process.is_decay, + return_matrix_element=bool(self.surrogates), + ) + ) + partial_weights = self.process.systematics_enabled() + madnis_args = self.process.run_card["madnis"] + custom_bias = self.process.custom_function("bias", self.particle_count, flavors) + integrands = [] + for channel in phasespace.channels: + integrands.append(ms.Integrand( + channel.phasespace_mapping, + cross_sections, + channel.adaptive_mapping, + channel.discrete_sym, + channel.discrete_flavor, + flavors, + pdf_grid, + self.process.running_coupling, + self.scale, + phasespace.prop_chan_weights, + phasespace.subchan_weights, + phasespace.cwnet, + phasespace.first_chan_weight_remap, + phasespace.first_remapped_chan_count, + phasespace.second_chan_weight_remap, + phasespace.second_remapped_chan_count, + madnis_training, + drop_cuts_and_rescale, + partial_weights, + channel.channel_weight_indices, + channel.active_flavors, + flavor_remap, + flavor_factors, + flavor_mirror, + flavor_diff_xs_indices, + flavor_subproc_indices, + flavor_per_subproc_remap, + madnis_args["compressed_channel_weight_count"], + custom_bias, + )) + #print(integrands[1].function()) + #for i in integrands: print(i.function()) + return integrands diff --git a/tests/unit_tests/interface/test_mg7_launch.py b/tests/unit_tests/interface/test_mg7_launch.py index 8f7fd39788..12f851e099 100644 --- a/tests/unit_tests/interface/test_mg7_launch.py +++ b/tests/unit_tests/interface/test_mg7_launch.py @@ -846,14 +846,15 @@ def test_setup_logging_defers_to_an_existing_handler(self): lg = logging.getLogger('madgraph') handler = logging.NullHandler() lg.addHandler(handler) - saved = self.launch._TOOL_LOGGING_READY - self.launch._TOOL_LOGGING_READY = False + from madgraph.iolibs.template_files.mg7 import postprocess + saved = postprocess._TOOL_LOGGING_READY + postprocess._TOOL_LOGGING_READY = False try: before = list(lg.handlers) self.launch._setup_logging() self.assertEqual(list(lg.handlers), before) finally: - self.launch._TOOL_LOGGING_READY = saved + postprocess._TOOL_LOGGING_READY = saved lg.removeHandler(handler) From 636deadf4d4a5dcf6c0bd433cbe1f610d9171641 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:08:46 +0200 Subject: [PATCH 19/88] mg7: weighting steps of each subprocess The subprocess returns its weighting steps, so the weighting groups do not depend on how a subprocess builds them. A [surrogate] subprocess returns its second pass with the full matrix element, as before. --- .../template_files/mg7/madgraph_process.py | 27 ++++++------------- .../mg7/subprocess_madmatrix.py | 13 +++++++++ 2 files changed, 21 insertions(+), 19 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/madgraph_process.py b/madgraph/iolibs/template_files/mg7/madgraph_process.py index 7321b86638..f1f70ae409 100644 --- a/madgraph/iolibs/template_files/mg7/madgraph_process.py +++ b/madgraph/iolibs/template_files/mg7/madgraph_process.py @@ -330,25 +330,14 @@ def init_surrogate(self) -> None: ) def weighting_groups(self, subprocess_channels: list[list[int]]) -> list: - """One ms.WeightingGroup per subprocess with a [surrogate]: the second - pass with the full matrix element over its channels, given by their - indices in the event generator.""" - if self.surrogate is None: - return [] - gen_args = self.run_card["generation"] - return [ - ms.WeightingGroup( - channel_indices, - [ms.WeightingStep( - subproc.step_matrix_elements, - mode=self.surrogate["mode"], - max_overweight_truncation=gen_args["max_overweight_truncation"], - freeze_max_weight_after=gen_args["freeze_max_weight_after"], - )], - ) - for subproc, channel_indices in zip(self.subprocesses, subprocess_channels) - if channel_indices - ] + """One ms.WeightingGroup per subprocess with weighting steps, over its + channels, given by their indices in the event generator.""" + groups = [] + for subproc, channel_indices in zip(self.subprocesses, subprocess_channels): + steps = subproc.weighting_steps() + if channel_indices and steps: + groups.append(ms.WeightingGroup(channel_indices, steps)) + return groups @staticmethod def build_custom_function(build_dir: str) -> None: diff --git a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py index 8352894842..98eb957a93 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py +++ b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py @@ -377,6 +377,19 @@ def load_surrogate(self, me_path_format: str, member_meta: dict): raise RuntimeError("surrogate loaded at different indices") return apis[0] + def weighting_steps(self) -> list: + """The steps after the first unweighting: with a [surrogate], the + second pass with the full matrix element.""" + if not self.surrogates: + return [] + gen_args = self.process.run_card["generation"] + return [ms.WeightingStep( + self.step_matrix_elements, + mode=self.process.surrogate["mode"], + max_overweight_truncation=gen_args["max_overweight_truncation"], + freeze_max_weight_after=gen_args["freeze_max_weight_after"], + )] + def flavor_pids(self) -> list[list[int]]: """PDG ids of the external particles for every flavor of the integrand, in the order of its flavor index.""" From 68714e6f104596723a052f723bceff14aa95a4dc Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 22:52:43 +0200 Subject: [PATCH 20/88] interface: '[amplicol]' processes 'generate p p > j j j [amplicol]' keeps the process definition for the mg7 output, which computes its matrix elements with AmpliCol, instead of generating MG diagrams. Such processes can be mixed with ordinary ones, and are refused by the other outputs and for decay chains, coupling orders, s-channel or particle restrictions and polarizations. --- madgraph/interface/madgraph_interface.py | 62 +++++++++++++- madgraph/interface/master_interface.py | 4 + .../interface/test_amplicol_syntax.py | 85 +++++++++++++++++++ 3 files changed, 150 insertions(+), 1 deletion(-) create mode 100644 tests/unit_tests/interface/test_amplicol_syntax.py diff --git a/madgraph/interface/madgraph_interface.py b/madgraph/interface/madgraph_interface.py index 836112fc28..32f008229c 100755 --- a/madgraph/interface/madgraph_interface.py +++ b/madgraph/interface/madgraph_interface.py @@ -178,6 +178,26 @@ def interference_mode_of(option, orders): return '' +# '[amplicol]' takes the matrix elements of a process from AmpliCol, a +# colour-ordered generator, instead of MG diagrams (mg7 output only) +ME_BACKENDS = ('amplicol',) +_BACKEND_BRACKET = re.compile(r"\[\s*(?P\w+)\s*\]") + + +def split_backend_line(line): + """(line without the bracket, backend) for a process with a matrix-element + backend bracket like '[amplicol]', (line, '') for any other line.""" + + match = _BACKEND_BRACKET.search(line) + if not match or match.group('backend').lower() not in ME_BACKENDS: + return line, '' + rest = line[:match.start()] + ' ' + line[match.end():] + if '[' in rest or ']' in rest: + raise madgraph.InvalidCmd( + "[%s] cannot be combined with another '[...]'." % match.group('backend')) + return ' '.join(rest.split()), match.group('backend').lower() + + def split_interference_line(line): """Split a process definition of the form @@ -2345,7 +2365,11 @@ def check_output(self, args, default='mg7'): return - if not self._curr_amps: + if self._curr_amplicol_defs and self._export_format != 'mg7': + raise self.InvalidCmd( + 'The [amplicol] processes are only supported by the mg7 output.') + + if not self._curr_amps and not self._curr_amplicol_defs: text = 'No processes generated. Please generate a process first.' raise self.InvalidCmd(text) @@ -3866,6 +3890,8 @@ def _tutorial_opts(self): _curr_model = None #base_objects.Model() _curr_amps = diagram_generation.AmplitudeList() _curr_proc_defs = base_objects.ProcessDefinitionList() + # processes whose matrix elements come from AmpliCol, see split_backend_line + _curr_amplicol_defs = base_objects.ProcessDefinitionList() _curr_matrix_elements = helas_objects.HelasMultiProcess() _curr_helas_model = None _curr_exporter = None @@ -3954,6 +3980,7 @@ def setup(self): self._done_export=False self._curr_amps = diagram_generation.AmplitudeList() self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() self._curr_matrix_elements = helas_objects.HelasMultiProcess() self._v4_export_formats = ['madevent', 'standalone_fortran','standalone_msP','standalone_msF', @@ -3978,6 +4005,30 @@ def do_quit(self, line): # Add a process to the existing multiprocess definition # Generate a new amplitude + def add_backend_process(self, procdef, backend): + """Keep a process whose matrix elements come from another backend than + MG diagrams. It is exported by the mg7 output, without diagrams.""" + + if procdef.get('decay_chains') or procdef.get('required_s_channels') or \ + procdef.get('forbidden_onsh_s_channels') or \ + procdef.get('forbidden_s_channels') or \ + procdef.get('forbidden_particles') or procdef.get('orders') or \ + procdef.get('squared_orders') or procdef.get('is_decay_chain') or \ + any(leg.get('polarization') for leg in procdef.get('legs')) or \ + procdef.get_ninitial() != 2: + raise self.InvalidCmd( + "[%s] supports 2 -> n processes without decay chains, coupling " + "orders, s-channel or particle restrictions and polarizations." + % backend) + initial = [amp.get_ninitial() for amp in self._curr_amps] + if initial and initial[0] != procdef.get_ninitial(): + raise self.InvalidCmd( + "Can not mix processes with different number of initial states.") + self._curr_amplicol_defs.append(procdef) + self._done_export = False + logger.info("Process %s is computed by %s at output time" % + (procdef.nice_string(prefix=False), backend)) + def do_add(self, line, counter=0): """Generate an amplitude for a given process and add to existing amplitudes @@ -4079,7 +4130,10 @@ def do_add(self, line, counter=0): else: nb_proc = len([l for l in self.history if l.startswith(('generate','add process'))]) nb_proc += counter + line, backend = split_backend_line(line) myprocdef = self.extract_process(line, proc_number=nb_proc) + if backend: + return self.add_backend_process(myprocdef, backend) # 'left [treextree] right': hidden order selecting the # interference, minimal orders of the right-hand process diagram_generation.prepare_interference_process(myprocdef) @@ -6045,6 +6099,7 @@ def clean_process(self): self._curr_amps = diagram_generation.AmplitudeList() # Reset Process definition self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() # Reset Helas matrix elements self._curr_matrix_elements = helas_objects.HelasMultiProcess() self._generate_info = "" @@ -7555,6 +7610,7 @@ def do_import(self, line, force=False, options={}): self._curr_amps = diagram_generation.AmplitudeList() # Reset proc defs self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() self._curr_matrix_elements = helas_objects.HelasMultiProcess() process_checks.store_aloha = [] self.advise_neglected_masses() @@ -10938,6 +10994,7 @@ def set2_group_subprocesses(self, args, log=True): logger.info('Note that you need to regenerate all processes') self._curr_amps = diagram_generation.AmplitudeList() self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() self._curr_matrix_elements = helas_objects.HelasMultiProcess() if self.options[args[0]] == 'gpu': export_v4.ProcessExporterFortranMEGroup.grouped_mode = 'gpu' @@ -11113,6 +11170,7 @@ def set2_gauge(self, args, log=True): self._curr_model = None self._curr_amps = diagram_generation.AmplitudeList() self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() self._curr_matrix_elements = helas_objects.HelasMultiProcess() self._curr_helas_model = None self._curr_exporter = None @@ -11592,6 +11650,7 @@ def set2_OLP(self, args, log=True): # depend on this option self._curr_amps = diagram_generation.AmplitudeList() self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() self._curr_matrix_elements = helas_objects.HelasMultiProcess() self._curr_exporter = None self._second_exporter = None @@ -13544,6 +13603,7 @@ def do_decay_diagram(self, line, skip_2body=False, model=None): # Reset amplitudes self._curr_amps = diagram_generation.AmplitudeList() self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() # Reset Helas matrix elements self._curr_matrix_elements = helas_objects.HelasMultiProcess() # Reset _done_export, since we have new process diff --git a/madgraph/interface/master_interface.py b/madgraph/interface/master_interface.py index 7db93f4e83..dd440419a9 100755 --- a/madgraph/interface/master_interface.py +++ b/madgraph/interface/master_interface.py @@ -173,6 +173,10 @@ def extract_process_type(line): space_before = re.compile(r"(?P\S)(?P[\\[\\]/\,\\$\\>|])(?P\S)") line2 = space_before.sub(r'\g \g \g', line) + # a matrix-element backend bracket like [amplicol] is a tree process + if MGcmd.split_backend_line(re.split(r'\s\-\-', line, 1)[0])[1]: + return ('tree', None, []) + # Use regular expressions to extract the loop mode (if present) and its # option, specified in the line with format [ option = loop_orders ] or # [ loop_orders ] which implicitly select the 'all' option. diff --git a/tests/unit_tests/interface/test_amplicol_syntax.py b/tests/unit_tests/interface/test_amplicol_syntax.py new file mode 100644 index 0000000000..1e20f0d35e --- /dev/null +++ b/tests/unit_tests/interface/test_amplicol_syntax.py @@ -0,0 +1,85 @@ +################################################################################ +# +# Copyright (c) 2026 The MadGraph5_aMC@NLO Development team and Contributors +# +# This file is a part of the MadGraph5_aMC@NLO project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph5_aMC@NLO license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""Syntax of the processes computed by AmpliCol + + generate p p > j j j [amplicol] + +The process is kept as a definition for the mg7 output, which generates its +matrix elements with pyamplicol; no MG diagrams are generated for it. +""" + +import os +import sys + +root_path = os.path.split(os.path.dirname(os.path.realpath(__file__)))[0] +sys.path.append(os.path.join(root_path, os.path.pardir, os.path.pardir)) + +import tests.unit_tests as unittest + +import madgraph.interface.madgraph_interface as mg_interface +import madgraph.interface.master_interface as MGCmd + +from madgraph import InvalidCmd + +split = mg_interface.split_backend_line + + +class BackendLineSplitTest(unittest.TestCase): + """split_backend_line, without any model.""" + + def test_split(self): + self.assertEqual(split('p p > j j j [amplicol]'), ('p p > j j j', 'amplicol')) + self.assertEqual(split('p p > j j j [ AmpliCol ] @2'), + ('p p > j j j @2', 'amplicol')) + self.assertEqual(split('p p > j j j'), ('p p > j j j', '')) + self.assertEqual(split('p p > t t~ [QCD]'), ('p p > t t~ [QCD]', '')) + + def test_other_brackets_are_refused(self): + self.assertRaises(InvalidCmd, split, 'p p > j j [amplicol] [QCD]') + + def test_process_type(self): + process_type = MGCmd.MasterCmd.extract_process_type + self.assertEqual(process_type('p p > j j j [amplicol]'), ('tree', None, [])) + self.assertEqual(process_type('p p > t t~ [QCD]'), ('NLO', 'all', ['QCD'])) + + +class AmplicolProcessTest(unittest.TestCase): + """generate/add process with [amplicol], with the sm model.""" + + def setUp(self): + self.cmd = MGCmd.MasterCmd() + self.cmd.exec_cmd('import model sm', printcmd=False, precmd=True) + + def test_no_diagrams_are_generated(self): + self.cmd.exec_cmd('generate g g > g g g [amplicol]', printcmd=False, + precmd=True) + self.cmd.exec_cmd('add process p p > e+ e-', printcmd=False, precmd=True) + self.assertEqual(len(self.cmd._curr_amplicol_defs), 1) + self.assertEqual( + self.cmd._curr_amplicol_defs[0].nice_string(prefix=False).split('@')[0].strip(), + 'g g > g g g') + # the madmatrix process keeps its own process number + self.assertTrue(self.cmd._curr_amps) + self.assertEqual( + {amp.get('process').get('id') for amp in self.cmd._curr_amps}, {2}) + self.cmd.exec_cmd('generate p p > e+ e-', printcmd=False, precmd=True) + self.assertEqual(len(self.cmd._curr_amplicol_defs), 0) + + def test_unsupported_options_are_refused(self): + for line in ['p p > j j QED=0 [amplicol]', + 'p p > z > e+ e- [amplicol]', + 'p p > j j / b [amplicol]', + 'z > e+ e- [amplicol]']: + self.assertRaises(InvalidCmd, self.cmd.do_generate, line) From 0b2f9c3556ac8dba1b4eab1635fe6d47dd669a5a Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 22:52:43 +0200 Subject: [PATCH 21/88] mg7: output of the [amplicol] processes Each [amplicol] process gets a SubProcesses/A_ directory with the leading- and full-colour pyamplicol artifacts, the sources of their ratio library and a makefile building both libraries into lib/. Every colour ordering becomes its own entry of subprocesses.json, with backend 'amplicol', and is never merged with other subprocesses. --- madgraph/iolibs/export_amplicol.py | 247 +++++++++++++++++++++++++++++ madgraph/iolibs/export_cpp.py | 20 ++- 2 files changed, 266 insertions(+), 1 deletion(-) create mode 100644 madgraph/iolibs/export_amplicol.py diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py new file mode 100644 index 0000000000..585d4c9b7a --- /dev/null +++ b/madgraph/iolibs/export_amplicol.py @@ -0,0 +1,247 @@ +################################################################################ +# +# Copyright (c) 2026 The MadGraph5_aMC@NLO Development team and Contributors +# +# This file is a part of the MadGraph5_aMC@NLO project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph5_aMC@NLO license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""mg7 output of the processes computed by AmpliCol ('[amplicol]'). + +pyamplicol generates the colour-ordered matrix elements of a process definition +in SubProcesses/A_: a leading-colour artifact (lc), a full-colour one +(full), and the sources of the UMAMI library returning their ratio (reweight). +Its makefile builds the libraries into lib/. Every colour ordering (UMAMI +channel) of every provider becomes one entry of subprocesses.json, integrated +on its own with a colour-ordered phase-space mapping. +""" + +import json +import logging +import os +import re +import shutil + +import madgraph + +logger = logging.getLogger('madgraph.export_amplicol') +pjoin = os.path.join + +MODEL_FILE = 'amplicol_model.pyamplicol-model' + +MAKEFILE = """\ +# Builds the AmpliCol matrix elements of this process into ../../lib. +RUSTICOL_CONFIG ?= %(rusticol_config)s +BUILD := $(CURDIR)/build + +all: +\t$(MAKE) -C lc/API/umami shared BUILD_DIR=$(BUILD)/lc RUSTICOL_CONFIG=$(RUSTICOL_CONFIG) +\t$(MAKE) -C reweight shared BUILD_DIR=$(BUILD)/reweight RUSTICOL_CONFIG=$(RUSTICOL_CONFIG) +%(copies)s +clean: +\trm -rf $(BUILD) +""" + + +def import_pyamplicol(): + try: + import pyamplicol + from pyamplicol.artifacts import umami_reweight + except ImportError as error: + raise madgraph.InvalidCmd( + 'The [amplicol] processes need pyamplicol (%s). Install it in the ' + 'Python environment of MadGraph.' % error) + return pyamplicol, umami_reweight + + +def process_name(procdef): + """A directory name for a process definition, e.g. 'pp_jjj'.""" + string = procdef.nice_string(prefix=False).split('@')[0] + initial, final = [side.split() for side in string.split('>')[:2]] + name = ''.join(initial) + '_' + ''.join(final) + return re.sub(r'[^A-Za-z0-9_]', '', + name.replace('~', 'x').replace('+', 'p').replace('-', 'm')) + + +def process_expression(procdef, model): + """The process in pyamplicol syntax, and its multiparticle definitions.""" + particles = model.get('particle_dict') + multiparticles = {} + labels = [] + for leg in procdef.get('legs'): + names = [particles[pid].get_name() for pid in leg.get('ids')] + if len(names) == 1: + labels.append(names[0]) + continue + label = next((key for key, value in multiparticles.items() + if value == names), None) + if label is None: + label = 'mp%d' % len(multiparticles) + multiparticles[label] = names + labels.append(label) + initial = [label for leg, label in zip(procdef.get('legs'), labels) + if not leg.get('state')] + final = [label for leg, label in zip(procdef.get('legs'), labels) + if leg.get('state')] + return '%s > %s' % (' '.join(initial), ' '.join(final)), multiparticles + + +def model_source(pyamplicol, model, cache_path): + """The UFO model of the output with its restriction, compiled for + pyamplicol once per output.""" + if os.path.exists(cache_path): + return pyamplicol.ModelSource.from_path(cache_path) + path = model.get('modelpath') + reference = os.path.basename(model.get('modelpath+restriction')) + restriction = 'default' + if '-' in reference: + suffix = reference.split('-', 1)[1] + if suffix == 'full': + restriction = 'none' + else: + restriction = pjoin(path, 'restrict_%s.dat' % suffix) + from pyamplicol.config import EvaluatorConfig, JITConfig + source = pyamplicol.ModelSource.from_path(path, restriction=restriction) + logger.info('Compiling the model %s for AmpliCol' % reference) + source.compile(prepared_output=cache_path, + evaluator=EvaluatorConfig(jit=JITConfig(optimization_level=2))) + return pyamplicol.ModelSource.from_path(cache_path) + + +def generate_artifact(pyamplicol, artifact, accuracy, expression, multiparticles, + model, workers): + from pyamplicol.config import resolve_config + config = resolve_config({ + 'schema_version': 1, + 'action': 'generate', + 'color': {'accuracy': accuracy}, + 'process': {'multiparticles': multiparticles}, + 'generation': {'workers': workers, 'emit_api_bundle': True, + 'relation_discovery': {'mode': 'off'}}, + 'evaluator': {'backend': 'jit', 'execution_mode': 'recurrence', + 'optimization': {'cores': workers}, + 'jit': {'optimization_level': 2}}, + }) + pyamplicol.Generator(config).generate(expression, artifact, model=model, + mode='replace') + + +def color_types(model, pdgs): + types = {} + for pdg in pdgs: + particle = model.get_particle(pdg) + types[pdg] = particle.get_color() + types[-pdg] = model.get_particle(particle.get_anti_pdg_code()).get_color() + return types + + +def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path): + """One subprocesses.json entry per channel of a UMAMI provider.""" + pdg_ids = metadata['pdg_ids'] + provider = metadata['provider'] + incoming = provider['incoming_count'] + entries = [] + for channel_index, channel in enumerate(metadata['channels']): + flavors = [] + for flavor_index, process in enumerate(channel['processes']): + if process['multichannels'] != [channel_index]: + raise madgraph.MadGraph5Error( + 'AmpliCol contributions shared by several channels are not ' + 'supported yet') + for term in process['matrix_elements']: + flavors.append({ + 'index': flavor_index, + 'options': [pdg_ids[term['pdg_ids']]], + 'factor': term['factor'], + 'mirror': False, + }) + first = flavors[0]['options'][0] + colored = sum(1 for pdg in first + if abs(model.get_particle(pdg).get_color()) > 1) + entries.append({ + 'backend': 'amplicol', + 'incoming': first[:incoming], + 'outgoing': first[incoming:], + 'channels': [], + 'path': rel_dir, + 'me_path': me_path, + 'ratio_path': ratio_path, + 'artifact': pjoin(rel_dir, 'lc'), + 'ratio_artifact': pjoin(rel_dir, 'reweight'), + 'umami_channel': channel_index, + 'phasespace_order': channel['phasespace_order'], + 'masses': provider['masses'], + 'flavors': flavors, + # tree level with the minimal coupling orders: every coloured + # parton beyond the first two comes with one power of alpha_s + 'qcd_power': max(colored - 2, 0), + 'color_flows': metadata['color_flows'], + 'pdg_color_types': color_types( + model, {abs(pdg) for pdgs in pdg_ids for pdg in pdgs}), + 'helicities': metadata['helicities'], + 'diagram_count': 0, + 'interference': False, + 'parameters': parameter_references(metadata), + }) + return entries + + +def parameter_references(metadata): + """The parameter card entry of every mutable runtime parameter.""" + references = {} + for process in metadata['runtime_processes']: + for parameter in process['parameters']: + if parameter.get('mutable') and parameter.get('lhablock'): + references[parameter['name']] = [parameter['lhablock'], + parameter['lhacode']] + return references + + +def export_amplicol_processes(dir_path, proc_defs, model, nb_core=None): + """Write SubProcesses/A_ for every '[amplicol]' process definition + and return the subprocesses.json entries.""" + pyamplicol, umami_reweight = import_pyamplicol() + workers = nb_core or 1 + subproc_dir = pjoin(dir_path, 'SubProcesses') + source = model_source(pyamplicol, model, pjoin(subproc_dir, MODEL_FILE)) + rusticol_config = shutil.which('rusticol-config') or 'rusticol-config' + entries = [] + for number, procdef in enumerate(proc_defs, start=1): + name = 'A%d_%s' % (number, process_name(procdef)) + path = pjoin(subproc_dir, name) + rel_dir = pjoin('SubProcesses', name) + expression, multiparticles = process_expression(procdef, model) + logger.info('Generating the AmpliCol matrix elements of %s in %s' % + (procdef.nice_string(prefix=False), path)) + os.makedirs(path, exist_ok=True) + for accuracy in ('lc', 'full'): + generate_artifact(pyamplicol, pjoin(path, accuracy), accuracy, + expression, multiparticles, source, workers) + provider_ids = umami_reweight.write_reweight_provider( + pjoin(path, 'lc'), pjoin(path, 'full'), pjoin(path, 'reweight')) + documents = umami_reweight.load_umami_metadata(pjoin(path, 'lc')) + copies = [] + for provider_id, metadata in zip(provider_ids, documents): + library = metadata['provider']['library'] + me_path = pjoin('lib', 'libamplicol_%s_%s.so' % (name, provider_id)) + ratio_path = pjoin('lib', 'libamplicol_%s_%s_ratio.so' % + (name, provider_id)) + copies.append('\tcp $(BUILD)/lc/%s ../../%s' % (library, me_path)) + copies.append('\tcp $(BUILD)/reweight/libreweight_%s.so ../../%s' % + (provider_id, ratio_path)) + entries.extend(subprocess_entries(metadata, model, rel_dir, me_path, + ratio_path)) + with open(pjoin(path, 'makefile'), 'w') as stream: + stream.write(MAKEFILE % {'rusticol_config': rusticol_config, + 'copies': '\n'.join(copies)}) + with open(pjoin(path, 'process.json'), 'w') as stream: + json.dump({'process': procdef.nice_string(prefix=False), + 'expression': expression, + 'multiparticles': multiparticles}, stream, indent=2) + return entries diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index e70cbfe98c..e791c79fa4 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3458,6 +3458,7 @@ def get_merged_info(self): return merged_subproc_info def finalize(self, matrix_elements=None, history='', *args, **kwargs): + self.export_amplicol() file_name = os.path.normpath(os.path.join( self.dir_path, "SubProcesses", "subprocesses.json" )) @@ -3504,8 +3505,25 @@ def finalize(self, matrix_elements=None, history='', *args, **kwargs): def pass_information_from_cmd(self, cmd): """Capture the process definitions from the command interface; needed to - generate the MadAnalysis5 default cards at output time.""" + generate the MadAnalysis5 default cards at output time, and the + [amplicol] processes.""" self.proc_defs = getattr(cmd, '_curr_proc_defs', None) + self.amplicol_defs = getattr(cmd, '_curr_amplicol_defs', None) or [] + self.amplicol_model = cmd._curr_model + self.amplicol_nb_core = cmd.options.get('nb_core') + + def export_amplicol(self): + """The [amplicol] processes, as subprocesses after the madmatrix ones.""" + if not getattr(self, 'amplicol_defs', None): + return + import madgraph.iolibs.export_amplicol as export_amplicol + for entry in export_amplicol.export_amplicol_processes( + self.dir_path, self.amplicol_defs, self.amplicol_model, + self.amplicol_nb_core): + # never merged with another subprocess + self.merged_subprocesses[('amplicol', len(self.process_info))].append( + (len(self.process_info), [])) + self.process_info.append(entry) def create_ma5_default_cards(self, matrix_elements, history): """Call MadAnalysis5 to write process-tailored default analysis cards From 40ce2e932feed303e1ee432774429172c5b9a433 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:01:25 +0200 Subject: [PATCH 22/88] madspace: weighting step matrix elements with a fixed channel A fixed channel is not an external input of a matrix element. The check that the random inputs of a step are passed in now only looks at them. --- madspace/src/phasespace/weighting_step.cpp | 11 +++++++- madspace/tests/test_integrand_me_outputs.py | 12 +++++---- madspace/tests/test_weighting_steps.py | 29 ++++++++++++++++++++- 3 files changed, 45 insertions(+), 7 deletions(-) diff --git a/madspace/src/phasespace/weighting_step.cpp b/madspace/src/phasespace/weighting_step.cpp index c0e75f7075..8982b1bf24 100644 --- a/madspace/src/phasespace/weighting_step.cpp +++ b/madspace/src/phasespace/weighting_step.cpp @@ -107,7 +107,16 @@ WeightingStepFunction::WeightingStepFunction( ); } } - if (matrix_element.external_inputs().size() != matrix_element.inputs().size()) { + bool has_random_inputs = std::any_of( + matrix_element.inputs().begin(), + matrix_element.inputs().end(), + [](auto input) { + return input == MatrixElement::random_color_in || + input == MatrixElement::random_helicity_in || + input == MatrixElement::random_diagram_in; + } + ); + if (has_random_inputs && matrix_element.sample_random_inputs()) { throw std::invalid_argument( "the random inputs of a weighting step matrix element are passed in " "and must not be sampled internally" diff --git a/madspace/tests/test_integrand_me_outputs.py b/madspace/tests/test_integrand_me_outputs.py index 82e400a5af..a5a9fc5c53 100644 --- a/madspace/tests/test_integrand_me_outputs.py +++ b/madspace/tests/test_integrand_me_outputs.py @@ -21,11 +21,11 @@ import madspace as ms -def python_matrix_element(context, function, outputs=("matrix_element",)): +def python_matrix_element(context, function, outputs=("matrix_element",), channel=-1): + """With a channel, the function also gets the fixed channel index.""" + inputs = ["momenta"] if channel < 0 else ["momenta", "channel_index"] api = context.add_matrix_element( - ms.PythonFunction( - function, ["momenta"], list(outputs), particle_count=len(PIDS) - ) + ms.PythonFunction(function, inputs, list(outputs), particle_count=len(PIDS)) ) output_enums = { "matrix_element": ms.MatrixElement.matrix_element_out, @@ -37,8 +37,10 @@ def python_matrix_element(context, function, outputs=("matrix_element",)): return ms.MatrixElement( api.index(), len(PIDS), - [ms.MatrixElement.momenta_in], + [ms.MatrixElement.momenta_in] + + ([] if channel < 0 else [ms.MatrixElement.channel_in]), [output_enums[output] for output in outputs], + channel=channel, ) diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py index fef4cbcf8b..1967cc74f8 100644 --- a/madspace/tests/test_weighting_steps.py +++ b/madspace/tests/test_weighting_steps.py @@ -76,6 +76,7 @@ def generate( step_outputs=("matrix_element",), override_outputs=(), integrand_outputs=("matrix_element",), + step_channel=-1, ): """Survey and generate, returns the status, the combined events and the channels. step_me and step_outputs can be lists, one entry per step.""" @@ -86,7 +87,7 @@ def generate( step_me = [step_me] * len(step_modes) step_outputs = [step_outputs] * len(step_modes) step_mes = [ - python_matrix_element(context, me, outputs) + python_matrix_element(context, me, outputs, step_channel) for me, outputs in zip(step_me, step_outputs) ] config = make_config() @@ -256,6 +257,32 @@ def test_ratio_step(tmp_path, step_modes): assert status_me.mean == status_ratio.mean +def test_ratio_with_fixed_channel(tmp_path): + """A step matrix element evaluated at a fixed channel gets it as an input, + which is not an input of the step.""" + + def ratio(momenta, channel): + return true_over_surrogate(momenta) * (channel == 3) + + results = [ + generate( + str(tmp_path / name), + surrogate_me, + ["reweight"], + step_me=step_me, + step_outputs=("matrix_element_ratio",), + step_channel=channel, + ) + for name, step_me, channel in [ + ("plain", true_over_surrogate, -1), + ("channel", ratio, 3), + ] + ] + (status_plain, events_plain, _), (status_channel, events_channel, _) = results + assert np.array_equal(events_plain, events_channel) + assert status_plain.mean == status_channel.mean + + def test_ratio_then_matrix_element(tmp_path, reference): """After a ratio step, the matrix element it implies is the reference of the next step, so the ratios still multiply to |M_true|^2 / |M_surrogate|^2.""" From 0e137ac8fafca8e9473d84cc63fd759f5a95c351 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:09:12 +0200 Subject: [PATCH 23/88] mg7: outputs with only [amplicol] processes An output without MG diagrams skips their export. The [amplicol] processes stand in for them where the output needs MG processes: the run card defaults, proc_characteristics and the model reference. --- madgraph/interface/madgraph_interface.py | 4 ++-- madgraph/iolibs/export_amplicol.py | 25 ++++++++++++++++++++++++ madgraph/iolibs/export_cpp.py | 18 +++++++++++++++-- 3 files changed, 43 insertions(+), 4 deletions(-) diff --git a/madgraph/interface/madgraph_interface.py b/madgraph/interface/madgraph_interface.py index 32f008229c..ce76124ce1 100755 --- a/madgraph/interface/madgraph_interface.py +++ b/madgraph/interface/madgraph_interface.py @@ -12544,7 +12544,7 @@ def do_output(self, line): # have been defined with multiparticle labels, because then # branching ratios necessitates to keep subprocesses independent. # That applies only if there is more than one subprocess of course. - if self._curr_amps[0].get_ninitial() == 1 and \ + if self._curr_amps and self._curr_amps[0].get_ninitial() == 1 and \ len(self._curr_amps)>1: processes = [amp.get('process') for amp in self._curr_amps if 'process' in list(amp.keys())] @@ -12940,7 +12940,7 @@ def generate_matrix_elements(self, group_processes=True): matrix_elements = self._curr_matrix_elements.get_matrix_elements() - if any(l.get('flavor') is None for l in matrix_elements[0].get('processes')[0].get('legs')): + if matrix_elements and any(l.get('flavor') is None for l in matrix_elements[0].get('processes')[0].get('legs')): [m.get_external_flavors() for m in matrix_elements] #precompute all flavors # Just the matrix.f files if self._export_format == 'matrix': diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index 585d4c9b7a..2f8a35e362 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -132,6 +132,30 @@ def generate_artifact(pyamplicol, artifact, accuracy, expression, multiparticles mode='replace') +def representative_processes(proc_defs, model): + """Processes standing for the [amplicol] definitions where the output needs + MG processes (run card defaults, proc_characteristics): every initial state, + with the first particle of each final-state leg.""" + import itertools + import madgraph.core.base_objects as base_objects + processes = [] + for procdef in proc_defs: + legs = procdef.get('legs') + initial = [leg.get('ids') for leg in legs if not leg.get('state')] + final = [leg.get('ids')[0] for leg in legs if leg.get('state')] + for ids in itertools.product(*initial): + pdgs = list(ids) + final + processes.append(base_objects.Process({ + 'legs': base_objects.LegList([ + base_objects.Leg({'id': pdg, 'number': number + 1, + 'state': number >= len(ids)}) + for number, pdg in enumerate(pdgs)]), + 'model': model, + 'id': procdef.get('id'), + })) + return processes + + def color_types(model, pdgs): types = {} for pdg in pdgs: @@ -177,6 +201,7 @@ def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path): 'umami_channel': channel_index, 'phasespace_order': channel['phasespace_order'], 'masses': provider['masses'], + 'supports_alpha_s': provider['supports_alpha_s'], 'flavors': flavors, # tree level with the minimal coupling orders: every coloured # parton beyond the first two comes with one power of alpha_s diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index e791c79fa4..7e34afcb35 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3512,6 +3512,14 @@ def pass_information_from_cmd(self, cmd): self.amplicol_model = cmd._curr_model self.amplicol_nb_core = cmd.options.get('nb_core') + def amplicol_processes(self): + """MG processes standing for the [amplicol] processes.""" + if not getattr(self, 'amplicol_defs', None): + return [] + import madgraph.iolibs.export_amplicol as export_amplicol + return export_amplicol.representative_processes( + self.amplicol_defs, self.amplicol_model) + def export_amplicol(self): """The [amplicol] processes, as subprocesses after the madmatrix ones.""" if not getattr(self, 'amplicol_defs', None): @@ -3613,6 +3621,8 @@ def create_run_card(self, matrix_elements, history): for me in matrix_elements['matrix_elements']] except (KeyError, TypeError): processes = None + if self.amplicol_processes(): + processes = (processes or []) + [self.amplicol_processes()] if processes: run_card.create_default_for_process(self.proc_characteristic, @@ -3649,8 +3659,12 @@ def create_proc_characteristics(self, matrix_elements): else: me_list = [] - procs = [] - qcd_orders = set() + procs = list(self.amplicol_processes()) + qcd_orders = set(entry['qcd_power'] for entry in self.process_info + if entry.get('backend') == 'amplicol') + if procs: + pc['nexternal'] = max(pc['nexternal'], len(procs[0].get('legs'))) + pc['ninitial'] = 2 for me in me_list: if not me.get('processes'): continue From 62cc98992efb30b2d75dacfb154a1fb4b3c2bb84 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:09:12 +0200 Subject: [PATCH 24/88] mg7: [amplicol] run card section The colour accuracy of the [amplicol] events, lc or full, whether the leading-colour events are re- or unweighted to it, and the settings of the leading-colour unweighting. --- madgraph/iolibs/template_files/mg7/run_card.toml | 9 +++++++++ madgraph/various/banner.py | 11 +++++++++++ 2 files changed, 20 insertions(+) diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index ea618c4862..0837f4258e 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -142,6 +142,15 @@ first_pass_max_overweight_truncation = %(surrogate.first_pass_max_overweight_tru first_pass_freeze_max_weight_after = %(surrogate.first_pass_freeze_max_weight_after)s param_card = %(surrogate.param_card)s # for a .so surrogate, "" uses Cards/param_card.dat +[amplicol] +# the [amplicol] processes are integrated at leading colour; their events are +# re- or unweighted to color_accuracy in a second pass (lc: no second pass) +color_accuracy = %(amplicol.color_accuracy)s +mode = %(amplicol.mode)s +# unweighting at leading colour (unweight mode); [generation] sets the second pass +first_pass_max_overweight_truncation = %(amplicol.first_pass_max_overweight_truncation)s +first_pass_freeze_max_weight_after = %(amplicol.first_pass_freeze_max_weight_after)s + [multiparticles] $multiparticles diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index c904f8691d..bcc6c6fb68 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6917,6 +6917,17 @@ def default_setup(self): comment="freeze_max_weight_after of the unweighting with the surrogate") self.add_toml_param('surrogate', 'param_card', "", comment="param card for a .so surrogate; \"\" uses Cards/param_card.dat") + # ---------------------------- [amplicol] ---------------------- + # the leading-colour [amplicol] subprocesses are re- or unweighted to + # the colour accuracy in a second pass + self.add_toml_param('amplicol', 'color_accuracy', "full", allowed=["full", "lc"], + comment="colour accuracy of the [amplicol] events, lc skips the second pass") + self.add_toml_param('amplicol', 'mode', "reweight", allowed=["unweight", "reweight"], + comment="unweight or reweight the events to the colour accuracy") + self.add_toml_param('amplicol', 'first_pass_max_overweight_truncation', 0.01, + comment="max_overweight_truncation of the leading-colour unweighting") + self.add_toml_param('amplicol', 'first_pass_freeze_max_weight_after', 10000, + comment="freeze_max_weight_after of the leading-colour unweighting") # ----------------------------- [madnis] ----------------------- self.add_toml_param('madnis', 'enable', False, allowed=["auto", True, False], auto=True, From b74dee3de7eb8fd6860dc4d00de999698e947d3b Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:09:13 +0200 Subject: [PATCH 25/88] mg7: subprocess hooks for other matrix-element backends The matrix elements of a subprocess are loaded by load_matrix_elements, and those of its integrand built by integrand_matrix_elements, which a subprocess of another backend overrides. A flavor can carry its own factor. --- .../mg7/subprocess_madmatrix.py | 75 +++++++++++-------- 1 file changed, 45 insertions(+), 30 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py index 98eb957a93..25ad94d7ea 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py +++ b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py @@ -336,6 +336,10 @@ def cut_pids(observable_kwargs): **self.process.scale_kwargs ) + self.load_matrix_elements(all_api_paths, api_path_formats, member_metas) + + def load_matrix_elements(self, all_api_paths, api_path_formats, member_metas): + """The matrix element and [surrogate] of every member subprocess.""" if self.process.run_card["run"]["dummy_matrix_element"]: self.matrix_elements = [None] * len(all_api_paths) else: @@ -377,6 +381,9 @@ def load_surrogate(self, me_path_format: str, member_meta: dict): raise RuntimeError("surrogate loaded at different indices") return apis[0] + def has_weighting_steps(self) -> bool: + return bool(self.surrogates) + def weighting_steps(self) -> list: """The steps after the first unweighting: with a [surrogate], the second pass with the full matrix element.""" @@ -1004,35 +1011,10 @@ def activation(self, name: str) -> ms.MLP.Activation: else: raise ValueError(f"Invalid activation function '{name}'") - def build_integrands( - self, - phasespace: PhaseSpace, - madnis_training: bool = False, - drop_cuts_and_rescale: bool = False - ) -> list[ms.Integrand]: - flavors = [] - flavor_remap = [] - flavor_factors = [] - flavor_mirror = [] - flavor_diff_xs_indices = [] - flavor_subproc_indices = [] - flavor_per_subproc_remap = [] - - for flav in self.meta["flavors"]: - if self.unmerged_meta is not None: - diff_xs_index = flav["subprocess"] - subproc_index = self.meta["subprocesses"][diff_xs_index] - ps_flavor = flav["flavor"] - flavor_diff_xs_indices.append(diff_xs_index) - flavor_subproc_indices.append(subproc_index) - flavor_per_subproc_remap.append(ps_flavor) - flav = self.unmerged_meta[subproc_index]["flavors"][ps_flavor] - flavors.append(flav["options"][0]) - flavor_remap.append(flav["index"]) - flavor_factors.append(len(flav["options"])) - flavor_mirror.append(flav["mirror"]) - - cross_sections = [] + def integrand_matrix_elements(self) -> list: + """The matrix element of every member subprocess in the integrand. Sets + the step_matrix_elements of the second pass of a [surrogate].""" + matrix_elements = [] self.step_matrix_elements = [] for me_index, matrix_element in enumerate(self.matrix_elements): if self.surrogates: @@ -1084,6 +1066,39 @@ def build_integrands( self.process.me_frame, self.process.incoming_count, ) + matrix_elements.append(mat) + return matrix_elements + + def build_integrands( + self, + phasespace: PhaseSpace, + madnis_training: bool = False, + drop_cuts_and_rescale: bool = False + ) -> list[ms.Integrand]: + flavors = [] + flavor_remap = [] + flavor_factors = [] + flavor_mirror = [] + flavor_diff_xs_indices = [] + flavor_subproc_indices = [] + flavor_per_subproc_remap = [] + + for flav in self.meta["flavors"]: + if self.unmerged_meta is not None: + diff_xs_index = flav["subprocess"] + subproc_index = self.meta["subprocesses"][diff_xs_index] + ps_flavor = flav["flavor"] + flavor_diff_xs_indices.append(diff_xs_index) + flavor_subproc_indices.append(subproc_index) + flavor_per_subproc_remap.append(ps_flavor) + flav = self.unmerged_meta[subproc_index]["flavors"][ps_flavor] + flavors.append(flav["options"][0]) + flavor_remap.append(flav["index"]) + flavor_factors.append(flav.get("factor", len(flav["options"]))) + flavor_mirror.append(flav["mirror"]) + + cross_sections = [] + for mat in self.integrand_matrix_elements(): pdf_grid = None if self.process.leptonic else self.process.pdf_grid pdf_arg = None if self.process.leptonic else ms.CachedPdf() cross_sections.append( @@ -1100,7 +1115,7 @@ def build_integrands( pdf2=pdf_arg, input_momentum_fraction=True, decay=self.process.is_decay, - return_matrix_element=bool(self.surrogates), + return_matrix_element=self.has_weighting_steps(), ) ) partial_weights = self.process.systematics_enabled() From 6f2c76f9f88a33ccfafde1acb46a8f561e238c6b Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:09:13 +0200 Subject: [PATCH 26/88] mg7: AmpliCol subprocesses Every colour ordering of an [amplicol] process is integrated at leading colour with the colour-ordered phase-space mapping of its UMAMI channel, and its events are re- or unweighted to full colour with the ratio library. The parameters of the param card are passed to the AmpliCol libraries, which the launcher builds when they are missing. GPUs, gridpacks and dummy matrix elements are refused. --- .../template_files/mg7/madgraph_process.py | 64 ++++-- .../template_files/mg7/subprocess_amplicol.py | 184 ++++++++++++++++++ 2 files changed, 235 insertions(+), 13 deletions(-) create mode 100644 madgraph/iolibs/template_files/mg7/subprocess_amplicol.py diff --git a/madgraph/iolibs/template_files/mg7/madgraph_process.py b/madgraph/iolibs/template_files/mg7/madgraph_process.py index f1f70ae409..a7de60ee76 100644 --- a/madgraph/iolibs/template_files/mg7/madgraph_process.py +++ b/madgraph/iolibs/template_files/mg7/madgraph_process.py @@ -21,6 +21,7 @@ from madgraph.iolibs.template_files.mg7.custom_function_loader import load_custom_function from madgraph.iolibs.template_files.mg7.observables import CutItem, HistItem, build_observable, clean_pids, observable_pids from madgraph.iolibs.template_files.mg7.options import lhapdf_paths +from madgraph.iolibs.template_files.mg7.subprocess_amplicol import AmplicolSubprocess, amplicol_subproc_args, is_amplicol from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import MadgraphSubprocess, PhaseSpace, build_multi_channel_data class MadgraphProcess: @@ -1021,13 +1022,23 @@ def init_generator_config(self) -> None: cfg.vegas_damping = vegas_args["damping"] cfg.max_overweight_truncation = gen_args["max_overweight_truncation"] cfg.freeze_max_weight_after = gen_args["freeze_max_weight_after"] + first_pass = None if self.surrogate is not None and self.surrogate["mode"] == "unweight": + first_pass = self.surrogate + amplicol = self.run_card["amplicol"] + if ( + first_pass is None and amplicol["mode"] == "unweight" + and amplicol["color_accuracy"] != "lc" + and any(is_amplicol(meta) for meta in self.subprocess_data) + ): + first_pass = amplicol + if first_pass is not None: # [generation] applies to the second pass, which gives the events cfg.max_overweight_truncation = ( - self.surrogate["first_pass_max_overweight_truncation"] + first_pass["first_pass_max_overweight_truncation"] ) cfg.freeze_max_weight_after = ( - self.surrogate["first_pass_freeze_max_weight_after"] + first_pass["first_pass_freeze_max_weight_after"] ) cfg.start_batch_size = vegas_args["start_batch_size"] cfg.max_batch_size = vegas_args["max_batch_size"] @@ -1065,11 +1076,19 @@ def init_subprocesses(self) -> None: self.subprocesses = [] if self.merged_subprocess_data is None: for subproc_id, meta in enumerate(self.subprocess_data): - self.subprocesses.append(MadgraphSubprocess(self, meta, subproc_id)) + subproc_class = ( + AmplicolSubprocess if is_amplicol(meta) else MadgraphSubprocess + ) + self.subprocesses.append(subproc_class(self, meta, subproc_id)) else: for subproc_id, meta in enumerate(self.merged_subprocess_data): + # an AmpliCol subprocess is never merged with another one + first = self.subprocess_data[meta["subprocesses"][0]] + subproc_class = ( + AmplicolSubprocess if is_amplicol(first) else MadgraphSubprocess + ) self.subprocesses.append( - MadgraphSubprocess(self, meta, subproc_id, self.subprocess_data) + subproc_class(self, meta, subproc_id, self.subprocess_data) ) def compile_matrix_elements(self) -> list[str]: @@ -1091,10 +1110,14 @@ def compile_matrix_elements(self) -> list[str]: device_names = [device_names] cpu_mode = self.run_card["run"]["cpu_mode"] backends = [backend_of(name, cpu_mode) for name in device_names] - if not self.subprocess_data: + self.compile_amplicol() + madmatrix_data = [ + meta for meta in self.subprocess_data if not is_amplicol(meta) + ] + if not madmatrix_data: return backends - first_proc_path = self.subprocess_data[0]["path"] + first_proc_path = madmatrix_data[0]["path"] subproc_path = os.path.dirname(first_proc_path) fptype = fptype_of(self.run_card["run"]["precision"]) @@ -1117,7 +1140,7 @@ def compile_matrix_elements(self) -> list[str]: log_path = os.path.join(self.run_path, "compile_subprocesses.log") for backend in resolved: missing = [ - meta for meta in self.subprocess_data + meta for meta in madmatrix_data if not os.path.isfile(meta["me_path"].format(device=backend)) ] if not missing: @@ -1138,6 +1161,19 @@ def compile_matrix_elements(self) -> list[str]: ) return resolved + def compile_amplicol(self) -> None: + """Build the libraries of the [amplicol] processes that are missing.""" + paths = sorted({ + meta["path"] for meta in self.subprocess_data + if is_amplicol(meta) and not ( + os.path.isfile(meta["me_path"]) and os.path.isfile(meta["ratio_path"]) + ) + }) + log_path = os.path.join(self.run_path, "compile_amplicol.log") + for path in paths: + logger.info(f"Compiling the AmpliCol matrix elements in {path}, see {log_path}") + self.make_subprocesses(path, [], 1, log_path) + @staticmethod def make_subprocesses( subproc_path: str, args: list[str], nb_core: int, log_path: str @@ -1759,13 +1795,15 @@ def write_lhe_header(self, systematics) -> None: del writer # closes the file (writes the closing tag) def build_lhe_completer(self): - all_mcdata = ( - [subproc.build_multi_channel_data() for subproc in self.subprocesses] - if self.merged_subprocess_data is None else - [build_multi_channel_data(meta, self) for meta in self.subprocess_data] - ) + all_mcdata = [ + None if is_amplicol(meta) + else self.subprocesses[index].build_multi_channel_data() + if self.merged_subprocess_data is None + else build_multi_channel_data(meta, self) + for index, meta in enumerate(self.subprocess_data) + ] subproc_args = [ - ms.SubprocArgs( + amplicol_subproc_args(meta, self) if mcdata is None else ms.SubprocArgs( topologies = [topo[0] for topo in mcdata.topologies], permutations = mcdata.permutations, diagram_indices = mcdata.diagram_indices, diff --git a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py new file mode 100644 index 0000000000..0c66a8ce85 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py @@ -0,0 +1,184 @@ +"""Subprocesses with AmpliCol matrix elements ('[amplicol]' processes). + +Each subprocess is one colour ordering (UMAMI channel) of an AmpliCol provider, +integrated on its own at leading colour with a colour-ordered phase-space +mapping. A second pass re- or unweights its events to full colour with the +ratio of the full-colour to the leading-colour matrix element summed over all +colour orderings, at the helicity of the event. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +from madgraph.iolibs.template_files.mg7.common import lhe_helicities, ms +from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import ( + Channel, + MadgraphSubprocess, + PhaseSpace, +) + +if TYPE_CHECKING: + from madgraph.iolibs.template_files.mg7.madgraph_process import MadgraphProcess + + +def is_amplicol(meta: dict) -> bool: + return meta.get("backend") == "amplicol" + + +def amplicol_parameters(meta: dict, process: MadgraphProcess) -> dict[str, float]: + """The values of the mutable model parameters in the param card.""" + parameters = {} + for name, (block, code) in meta["parameters"].items(): + try: + parameters[name] = float( + process.param_card.get_value(block.lower(), tuple(code)) + ) + except (KeyError, ValueError): + continue + return parameters + + +def chain_topology(incoming_masses: list[float], outgoing_masses: list[float]): + """A t-channel chain without propagators, for the LHE completion.""" + n_out = len(outgoing_masses) + vertices = [["i0", "p0", "o0"]] + vertices += [[f"p{i - 1}", f"p{i}", f"o{i}"] for i in range(1, n_out - 1)] + vertices.append(["i1", f"p{n_out - 2}", f"o{n_out - 1}"]) + propagators = [ms.Propagator() for _ in range(n_out - 1)] + return ms.Topology( + ms.Diagram(incoming_masses, outgoing_masses, propagators, vertices) + ) + + +def amplicol_subproc_args(meta: dict, process: MadgraphProcess) -> ms.SubprocArgs: + """LHE completion of an AmpliCol subprocess: its events have no diagram.""" + return ms.SubprocArgs( + topologies=[chain_topology( + [process.get_mass(pid) for pid in meta["incoming"]], + [process.get_mass(pid) for pid in meta["outgoing"]], + )], + permutations=[[]], + diagram_indices=[[]], + diagram_color_indices=[[]], + color_flows=meta["color_flows"], + pdg_color_types={ + int(key): value for key, value in meta["pdg_color_types"].items() + }, + helicities=lhe_helicities(meta, process.run_card), + pdg_ids=[flavor["options"] for flavor in meta["flavors"]], + ) + + +class AmplicolSubprocess(MadgraphSubprocess): + def load_matrix_elements(self, all_api_paths, api_path_formats, member_metas): + run_args = self.process.run_card["run"] + if any(device != "cpu" for device in self.process.device_types): + raise RuntimeError("[amplicol] processes are only supported on CPUs") + if run_args["dummy_matrix_element"]: + raise RuntimeError("[amplicol] processes need their matrix elements") + if self.process.run_card["gridpack"]["save_gridpack"]: + raise RuntimeError("[amplicol] processes do not support gridpacks yet") + if not self.meta.get("supports_alpha_s", True): + raise RuntimeError( + "the AmpliCol matrix elements of this model cannot take the running " + "alpha_s, as the model has no SMINPUTS 3 parameter" + ) + parameters = amplicol_parameters(self.meta, self.process) + self.channel_index = self.meta["umami_channel"] + self.matrix_elements = [ + self.load(self.meta["me_path"], self.meta["artifact"], parameters) + ] + self.surrogates = [] + self.ratio_matrix_element = None + if self.process.run_card["amplicol"]["color_accuracy"] != "lc": + self.ratio_matrix_element = self.load( + self.meta["ratio_path"], self.meta["ratio_artifact"], parameters + ) + + def load(self, path: str, artifact: str, parameters: dict): + # the libraries find their pyamplicol artifacts from an absolute path + for context in self.process.contexts: + api = context.load_matrix_element(path, os.path.abspath(artifact), parameters) + return api + + def build_multi_channel_data(self): + return None + + def build_multichannel_phasespace(self) -> PhaseSpace: + mapping = ms.PhaseSpaceMapping( + self.incoming_masses + self.outgoing_masses, + self.process.e_cm, + leptonic=self.process.leptonic, + invariant_power=self.process.run_card["phasespace"]["invariant_power"], + mode=ms.PhaseSpaceMapping.color_ordered, + cuts=self.cuts, + color_order=self.meta["phasespace_order"], + ) + prefix = f"subproc{self.subproc_id}.amplicol" + discrete_sym, discrete_flavor = self.build_discrete( + 1, len(self.meta["flavors"]), prefix + ) + channel = Channel( + phasespace_mapping=mapping, + adaptive_mapping=self.build_vegas(mapping, prefix), + discrete_sym=discrete_sym, + discrete_flavor=discrete_flavor, + channel_weight_indices=[0], + name="C", + active_flavors=[], + ) + return PhaseSpace(mode="multichannel", channels=[channel], symfact=[None]) + + def build_flat_phasespace(self) -> PhaseSpace: + return self.build_multichannel_phasespace() + + def simplify_phasespace(self, multi_phasespace, cross_sections) -> PhaseSpace: + return multi_phasespace + + def integrand_matrix_elements(self) -> list: + inputs = [ + ms.MatrixElement.momenta_in, + ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in, + ms.MatrixElement.channel_in, + ms.MatrixElement.random_helicity_in, + ] + outputs = [ + ms.MatrixElement.matrix_element_out, + ms.MatrixElement.color_index_out, + ms.MatrixElement.helicity_index_out, + ] + return [ms.MatrixElement( + self.matrix_elements[0].index(), self.particle_count, inputs, outputs, 1, + True, self.process.me_frame, self.process.incoming_count, + self.channel_index, + )] + + def has_weighting_steps(self) -> bool: + return self.ratio_matrix_element is not None + + def weighting_steps(self) -> list: + """The second pass to full colour.""" + if self.ratio_matrix_element is None: + return [] + ratio = ms.MatrixElement( + self.ratio_matrix_element.index(), + self.particle_count, + [ms.MatrixElement.momenta_in, ms.MatrixElement.flavor_in, + ms.MatrixElement.channel_in, ms.MatrixElement.helicity_in], + [ms.MatrixElement.matrix_element_ratio_out], + 1, + False, + self.process.me_frame, + self.process.incoming_count, + self.channel_index, + ) + gen_args = self.process.run_card["generation"] + return [ms.WeightingStep( + [ratio], + mode=self.process.run_card["amplicol"]["mode"], + max_overweight_truncation=gen_args["max_overweight_truncation"], + freeze_max_weight_after=gen_args["freeze_max_weight_after"], + )] From a8bfeee108fade4f06d4df445f1fcaf381960607 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:11:00 +0200 Subject: [PATCH 27/88] mg7: [amplicol] processes with flavour-grouped legs MG replaces the light quarks and leptons of a process definition by flavour groups, which pyamplicol does not know: the export passes the original flavours of a leg, or the members of its groups. The directory is named with the multiparticle labels of the definition. --- madgraph/iolibs/export_amplicol.py | 45 ++++++++++++++----- madgraph/iolibs/export_cpp.py | 3 +- .../interface/test_amplicol_syntax.py | 42 +++++++++++++++++ 3 files changed, 79 insertions(+), 11 deletions(-) diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index 2f8a35e362..ec8d9e5f0a 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -60,11 +60,35 @@ def import_pyamplicol(): return pyamplicol, umami_reweight -def process_name(procdef): - """A directory name for a process definition, e.g. 'pp_jjj'.""" - string = procdef.nice_string(prefix=False).split('@')[0] - initial, final = [side.split() for side in string.split('>')[:2]] - name = ''.join(initial) + '_' + ''.join(final) +def leg_pdgs(leg, model): + """The pdg ids of a leg: the flavours MG merged into a group id, or the ids + with the group ids expanded.""" + if leg.get('flavor'): + return list(leg.get('flavor')) + merged = model.get('merged_particles') + pdgs = [] + for pid in leg.get('ids'): + members = merged.get(abs(pid), [abs(pid)]) + pdgs.extend(member if pid > 0 else -member for member in members) + return pdgs + + +def process_name(procdef, model, multiparticles): + """A directory name for a process definition, e.g. 'pp_jjj', with the + multiparticle labels of the interface.""" + particles = model.get('particle_dict') + labels = [] + for leg in procdef.get('legs'): + pdgs = leg_pdgs(leg, model) + matches = [label for label, ids in multiparticles.items() + if sorted(ids) == sorted(pdgs)] + preferred = 'j' if leg.get('state') else 'p' + label = preferred if preferred in matches else next(iter(matches), None) + if label is None: + label = ''.join(particles[pid].get_name() for pid in pdgs) + labels.append((leg.get('state'), label)) + name = ''.join(l for state, l in labels if not state) + '_' + \ + ''.join(l for state, l in labels if state) return re.sub(r'[^A-Za-z0-9_]', '', name.replace('~', 'x').replace('+', 'p').replace('-', 'm')) @@ -75,7 +99,7 @@ def process_expression(procdef, model): multiparticles = {} labels = [] for leg in procdef.get('legs'): - names = [particles[pid].get_name() for pid in leg.get('ids')] + names = [particles[pid].get_name() for pid in leg_pdgs(leg, model)] if len(names) == 1: labels.append(names[0]) continue @@ -141,8 +165,8 @@ def representative_processes(proc_defs, model): processes = [] for procdef in proc_defs: legs = procdef.get('legs') - initial = [leg.get('ids') for leg in legs if not leg.get('state')] - final = [leg.get('ids')[0] for leg in legs if leg.get('state')] + initial = [leg_pdgs(leg, model) for leg in legs if not leg.get('state')] + final = [leg_pdgs(leg, model)[0] for leg in legs if leg.get('state')] for ids in itertools.product(*initial): pdgs = list(ids) + final processes.append(base_objects.Process({ @@ -228,7 +252,8 @@ def parameter_references(metadata): return references -def export_amplicol_processes(dir_path, proc_defs, model, nb_core=None): +def export_amplicol_processes(dir_path, proc_defs, model, nb_core=None, + multiparticles=None): """Write SubProcesses/A_ for every '[amplicol]' process definition and return the subprocesses.json entries.""" pyamplicol, umami_reweight = import_pyamplicol() @@ -238,7 +263,7 @@ def export_amplicol_processes(dir_path, proc_defs, model, nb_core=None): rusticol_config = shutil.which('rusticol-config') or 'rusticol-config' entries = [] for number, procdef in enumerate(proc_defs, start=1): - name = 'A%d_%s' % (number, process_name(procdef)) + name = 'A%d_%s' % (number, process_name(procdef, model, multiparticles or {})) path = pjoin(subproc_dir, name) rel_dir = pjoin('SubProcesses', name) expression, multiparticles = process_expression(procdef, model) diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index 7e34afcb35..dcd3d0be43 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3511,6 +3511,7 @@ def pass_information_from_cmd(self, cmd): self.amplicol_defs = getattr(cmd, '_curr_amplicol_defs', None) or [] self.amplicol_model = cmd._curr_model self.amplicol_nb_core = cmd.options.get('nb_core') + self.amplicol_multiparticles = dict(getattr(cmd, '_multiparticles', {})) def amplicol_processes(self): """MG processes standing for the [amplicol] processes.""" @@ -3527,7 +3528,7 @@ def export_amplicol(self): import madgraph.iolibs.export_amplicol as export_amplicol for entry in export_amplicol.export_amplicol_processes( self.dir_path, self.amplicol_defs, self.amplicol_model, - self.amplicol_nb_core): + self.amplicol_nb_core, self.amplicol_multiparticles): # never merged with another subprocess self.merged_subprocesses[('amplicol', len(self.process_info))].append( (len(self.process_info), [])) diff --git a/tests/unit_tests/interface/test_amplicol_syntax.py b/tests/unit_tests/interface/test_amplicol_syntax.py index 1e20f0d35e..493801d065 100644 --- a/tests/unit_tests/interface/test_amplicol_syntax.py +++ b/tests/unit_tests/interface/test_amplicol_syntax.py @@ -83,3 +83,45 @@ def test_unsupported_options_are_refused(self): 'p p > j j / b [amplicol]', 'z > e+ e- [amplicol]']: self.assertRaises(InvalidCmd, self.cmd.do_generate, line) + + +class AmplicolExportTest(unittest.TestCase): + """The process definitions as pyamplicol sees them, without pyamplicol.""" + + def setUp(self): + self.cmd = MGCmd.MasterCmd() + self.cmd.exec_cmd('import model sm', printcmd=False, precmd=True) + self.cmd.exec_cmd('define q = u d', printcmd=False, precmd=True) + self.cmd.exec_cmd('define qx = u~ d~', printcmd=False, precmd=True) + + def export(self, line): + import madgraph.iolibs.export_amplicol as export_amplicol + self.cmd.exec_cmd('generate %s [amplicol]' % line, printcmd=False, + precmd=True) + procdef = self.cmd._curr_amplicol_defs[0] + model = self.cmd._curr_model + return (export_amplicol.process_name(procdef, model, + self.cmd._multiparticles), + export_amplicol.process_expression(procdef, model), + export_amplicol.representative_processes([procdef], model)) + + def test_flavour_groups_are_expanded(self): + """MG replaces the light quarks by a flavour group, which pyamplicol + does not know: the original flavours are passed on.""" + name, expression, processes = self.export('u u~ > e+ e- g') + self.assertEqual(name, 'uux_epemg') + self.assertEqual(expression, ('u u~ > e+ e- g', {})) + self.assertEqual(len(processes), 1) + name, expression, processes = self.export('q qx > e+ e- g') + self.assertEqual(name, 'qqx_epemg') + self.assertEqual(expression, ('mp0 mp1 > e+ e- g', + {'mp0': ['u', 'd'], 'mp1': ['u~', 'd~']})) + self.assertEqual(len(processes), 4) + + def test_multiparticles(self): + name, (expression, multiparticles), processes = self.export('p p > j j') + self.assertEqual(name, 'pp_jj') + self.assertEqual(expression, 'mp0 mp0 > mp0 mp0') + self.assertEqual(sorted(multiparticles['mp0']), + sorted(['g', 'u', 'd', 's', 'c', 'u~', 'd~', 's~', 'c~'])) + self.assertEqual(len(processes), 81) From 4e6d7164ab778598ba89d7b4053889eb93883ea4 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:22:04 +0200 Subject: [PATCH 28/88] mg7: coupling orders of the [amplicol] processes pyamplicol gets the largest orders of the couplings that MG selects by default, the minimal WEIGHTED order of the process definition, instead of its minimal orders per process, which differ for the processes of a multiparticle definition. The power of alpha_s follows from them. pyamplicol runs with one worker, as Symbolica's restricted mode allows a single thread. --- madgraph/iolibs/export_amplicol.py | 48 ++++++++++++------- madgraph/iolibs/export_cpp.py | 3 +- .../interface/test_amplicol_syntax.py | 11 +++++ 3 files changed, 44 insertions(+), 18 deletions(-) diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index ec8d9e5f0a..fe3315fb6c 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -116,6 +116,20 @@ def process_expression(procdef, model): return '%s > %s' % (' '.join(initial), ' '.join(final)), multiparticles +def coupling_orders(procdef, model): + """The largest orders of the couplings above the lowest hierarchy, as MG + selects them by default: the minimal WEIGHTED order of the definition, + with n - 2 vertices at tree level.""" + import madgraph.core.diagram_generation as diagram_generation + weighted = diagram_generation.MultiProcess.find_optimal_process_orders( + procdef)['WEIGHTED'] + hierarchy = model.get('order_hierarchy') + lowest = min(hierarchy.values()) + vertices = len(procdef.get('legs')) - 2 + return {order: max((weighted - vertices * lowest) // (weight - lowest), 0) + for order, weight in hierarchy.items() if weight > lowest} + + def model_source(pyamplicol, model, cache_path): """The UFO model of the output with its restriction, compiled for pyamplicol once per output.""" @@ -139,17 +153,20 @@ def model_source(pyamplicol, model, cache_path): def generate_artifact(pyamplicol, artifact, accuracy, expression, multiparticles, - model, workers): + orders, model): from pyamplicol.config import resolve_config + # one worker: Symbolica's restricted mode allows a single thread config = resolve_config({ 'schema_version': 1, 'action': 'generate', 'color': {'accuracy': accuracy}, - 'process': {'multiparticles': multiparticles}, - 'generation': {'workers': workers, 'emit_api_bundle': True, + 'process': {'multiparticles': multiparticles, + 'coupling_order_policy': 'explicit', + 'max_coupling_orders': orders}, + 'generation': {'workers': 1, 'emit_api_bundle': True, 'relation_discovery': {'mode': 'off'}}, 'evaluator': {'backend': 'jit', 'execution_mode': 'recurrence', - 'optimization': {'cores': workers}, + 'optimization': {'cores': 1}, 'jit': {'optimization_level': 2}}, }) pyamplicol.Generator(config).generate(expression, artifact, model=model, @@ -189,7 +206,7 @@ def color_types(model, pdgs): return types -def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path): +def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power): """One subprocesses.json entry per channel of a UMAMI provider.""" pdg_ids = metadata['pdg_ids'] provider = metadata['provider'] @@ -210,8 +227,6 @@ def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path): 'mirror': False, }) first = flavors[0]['options'][0] - colored = sum(1 for pdg in first - if abs(model.get_particle(pdg).get_color()) > 1) entries.append({ 'backend': 'amplicol', 'incoming': first[:incoming], @@ -227,9 +242,7 @@ def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path): 'masses': provider['masses'], 'supports_alpha_s': provider['supports_alpha_s'], 'flavors': flavors, - # tree level with the minimal coupling orders: every coloured - # parton beyond the first two comes with one power of alpha_s - 'qcd_power': max(colored - 2, 0), + 'qcd_power': qcd_power, 'color_flows': metadata['color_flows'], 'pdg_color_types': color_types( model, {abs(pdg) for pdgs in pdg_ids for pdg in pdgs}), @@ -252,12 +265,10 @@ def parameter_references(metadata): return references -def export_amplicol_processes(dir_path, proc_defs, model, nb_core=None, - multiparticles=None): +def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None): """Write SubProcesses/A_ for every '[amplicol]' process definition and return the subprocesses.json entries.""" pyamplicol, umami_reweight = import_pyamplicol() - workers = nb_core or 1 subproc_dir = pjoin(dir_path, 'SubProcesses') source = model_source(pyamplicol, model, pjoin(subproc_dir, MODEL_FILE)) rusticol_config = shutil.which('rusticol-config') or 'rusticol-config' @@ -267,12 +278,16 @@ def export_amplicol_processes(dir_path, proc_defs, model, nb_core=None, path = pjoin(subproc_dir, name) rel_dir = pjoin('SubProcesses', name) expression, multiparticles = process_expression(procdef, model) + orders = coupling_orders(procdef, model) + # every vertex carries one coupling: the rest are powers of g_s + qcd_power = (len(procdef.get('legs')) - 2 - sum(orders.values()) + if len(orders) == 1 else -1) logger.info('Generating the AmpliCol matrix elements of %s in %s' % (procdef.nice_string(prefix=False), path)) os.makedirs(path, exist_ok=True) for accuracy in ('lc', 'full'): generate_artifact(pyamplicol, pjoin(path, accuracy), accuracy, - expression, multiparticles, source, workers) + expression, multiparticles, orders, source) provider_ids = umami_reweight.write_reweight_provider( pjoin(path, 'lc'), pjoin(path, 'full'), pjoin(path, 'reweight')) documents = umami_reweight.load_umami_metadata(pjoin(path, 'lc')) @@ -286,12 +301,13 @@ def export_amplicol_processes(dir_path, proc_defs, model, nb_core=None, copies.append('\tcp $(BUILD)/reweight/libreweight_%s.so ../../%s' % (provider_id, ratio_path)) entries.extend(subprocess_entries(metadata, model, rel_dir, me_path, - ratio_path)) + ratio_path, qcd_power)) with open(pjoin(path, 'makefile'), 'w') as stream: stream.write(MAKEFILE % {'rusticol_config': rusticol_config, 'copies': '\n'.join(copies)}) with open(pjoin(path, 'process.json'), 'w') as stream: json.dump({'process': procdef.nice_string(prefix=False), 'expression': expression, - 'multiparticles': multiparticles}, stream, indent=2) + 'multiparticles': multiparticles, + 'max_coupling_orders': orders}, stream, indent=2) return entries diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index dcd3d0be43..de2d2f116f 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3510,7 +3510,6 @@ def pass_information_from_cmd(self, cmd): self.proc_defs = getattr(cmd, '_curr_proc_defs', None) self.amplicol_defs = getattr(cmd, '_curr_amplicol_defs', None) or [] self.amplicol_model = cmd._curr_model - self.amplicol_nb_core = cmd.options.get('nb_core') self.amplicol_multiparticles = dict(getattr(cmd, '_multiparticles', {})) def amplicol_processes(self): @@ -3528,7 +3527,7 @@ def export_amplicol(self): import madgraph.iolibs.export_amplicol as export_amplicol for entry in export_amplicol.export_amplicol_processes( self.dir_path, self.amplicol_defs, self.amplicol_model, - self.amplicol_nb_core, self.amplicol_multiparticles): + self.amplicol_multiparticles): # never merged with another subprocess self.merged_subprocesses[('amplicol', len(self.process_info))].append( (len(self.process_info), [])) diff --git a/tests/unit_tests/interface/test_amplicol_syntax.py b/tests/unit_tests/interface/test_amplicol_syntax.py index 493801d065..d52a290ad6 100644 --- a/tests/unit_tests/interface/test_amplicol_syntax.py +++ b/tests/unit_tests/interface/test_amplicol_syntax.py @@ -125,3 +125,14 @@ def test_multiparticles(self): self.assertEqual(sorted(multiparticles['mp0']), sorted(['g', 'u', 'd', 's', 'c', 'u~', 'd~', 's~', 'c~'])) self.assertEqual(len(processes), 81) + + def test_coupling_orders(self): + """The orders MG selects by default, as the largest non-QCD orders.""" + import madgraph.iolibs.export_amplicol as export_amplicol + for line, orders in [('p p > j j j', {'QED': 0}), + ('p p > e+ e- j', {'QED': 2}), + ('p p > w+ j j', {'QED': 1})]: + self.cmd.exec_cmd('generate %s [amplicol]' % line, printcmd=False, + precmd=True) + self.assertEqual(export_amplicol.coupling_orders( + self.cmd._curr_amplicol_defs[0], self.cmd._curr_model), orders) From 327e2b5dc7191bc0750f483278b73b600b1cd237 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:22:04 +0200 Subject: [PATCH 29/88] mg7: refuse cuts that are not flavour blind for [amplicol] The flavours of an AmpliCol channel share its momenta and its cuts, built for the first flavour, so every cut has to select the same particles for all of them. --- .../template_files/mg7/subprocess_amplicol.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py index 0c66a8ce85..eedabb5118 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py +++ b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING from madgraph.iolibs.template_files.mg7.common import lhe_helicities, ms +from madgraph.iolibs.template_files.mg7.observables import clean_pids from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import ( Channel, MadgraphSubprocess, @@ -72,6 +73,27 @@ def amplicol_subproc_args(meta: dict, process: MadgraphProcess) -> ms.SubprocArg class AmplicolSubprocess(MadgraphSubprocess): + def __init__(self, process, meta, subproc_id, unmerged_meta=None): + super().__init__(process, meta, subproc_id, unmerged_meta) + self.check_flavor_blind_cuts() + + def check_flavor_blind_cuts(self): + """AmpliCol groups the flavours and identical-particle orderings that + share a matrix element into one channel, so its cuts, built for the + first flavour, must select the same particles for all of them.""" + options = [ + clean_pids(pdgs) for flavor in self.meta["flavors"] + for pdgs in flavor["options"] + ] + for item in self.process.cut_data: + for group in item.observable_kwargs["select_pids"]: + if len({tuple(pid in group for pid in pdgs) for pdgs in options}) > 1: + raise RuntimeError( + f"the cut {item.observable_kwargs['name']} does not select " + f"the same particles in all flavours of the [amplicol] " + f"subprocess {self.meta['incoming']} > {self.meta['outgoing']}" + ) + def load_matrix_elements(self, all_api_paths, api_path_formats, member_metas): run_args = self.process.run_card["run"] if any(device != "cpu" for device in self.process.device_types): From b38793b27c1c9dea3b2f73939b3b4101bfabc1a6 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:22:05 +0200 Subject: [PATCH 30/88] tests: mg7 [amplicol] cross sections against madmatrix gg > ggg in both colour modes and at leading colour, p p > j j j, and an output mixing madmatrix and [amplicol] processes. --- tests/acceptance_tests/test_mg7_amplicol.py | 158 ++++++++++++++++++++ 1 file changed, 158 insertions(+) create mode 100644 tests/acceptance_tests/test_mg7_amplicol.py diff --git a/tests/acceptance_tests/test_mg7_amplicol.py b/tests/acceptance_tests/test_mg7_amplicol.py new file mode 100644 index 0000000000..539f31ffbd --- /dev/null +++ b/tests/acceptance_tests/test_mg7_amplicol.py @@ -0,0 +1,158 @@ +################################################################################ +# +# Copyright (c) 2026 The MadGraph5_aMC@NLO Development team and Contributors +# +# This file is a part of the MadGraph5_aMC@NLO project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph5_aMC@NLO license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""mg7 runs of '[amplicol]' processes against the same processes with madmatrix. + +The AmpliCol subprocesses are integrated at leading colour and re- or +unweighted to full colour, so their cross sections must agree with the +madmatrix ones within the Monte Carlo errors, alone and mixed with madmatrix +processes in one output. Needs madspace, pyamplicol and LHAPDF; the tests skip +otherwise. Knobs: MG7_AMPLICOL_EVENTS (default 10000) and MG7_AMPLICOL_NSIGMA +(default 4). + + ./tests/test_manager.py test_mg7_amplicol -pA -t0 +""" + +import glob +import gzip +import json +import os +import re +import shutil +import subprocess +import sys +import tempfile +import unittest + +import madgraph.interface.master_interface as MGCmd +import madgraph.various.misc as misc + +pjoin = os.path.join + +_EVENTS = int(os.environ.get('MG7_AMPLICOL_EVENTS', 10000)) +_NSIGMA = float(os.environ.get('MG7_AMPLICOL_NSIGMA', 4)) +_PDF = 'NNPDF23_lo_as_0130_qed' + + +def _edit_run_card(path, **amplicol): + card = open(path).read() + card = card.replace('fixed_ren_scale = false', 'fixed_ren_scale = true') + card = card.replace('fixed_fact_scale = false', 'fixed_fact_scale = true') + card = re.sub(r'(?m)^pdf = ".*"$', 'pdf = "%s"' % _PDF, card) + card = re.sub(r'(?m)^events = \d+', 'events = %d' % _EVENTS, card) + card = re.sub(r'(?m)^seed = -1', 'seed = 7', card) + card = re.sub(r'(?m)^make_plots = true', 'make_plots = false', card) + card = re.sub(r'(\[systematics\]\nenable = )true', r'\1false', card) + card = re.sub(r'(\[madnis\]\n(?:#.*\n)*enable = )"auto"', r'\1false', card) + head, section = card.split('\n[amplicol]\n') + for key, value in amplicol.items(): + section, count = re.subn(r'(?m)^%s = .*$' % key, + '%s = "%s"' % (key, value), section, count=1) + assert count == 1, key + card = head + '\n[amplicol]\n' + section + open(path, 'w').write(card) + + +def _color_balanced(lhe_path): + """Whether every colour line of every event starts and ends once.""" + with gzip.open(lhe_path, 'rt') as stream: + text = stream.read() + for event in re.findall(r'\n(.*?)', text, re.S): + lines = [line.split() for line in event.strip().splitlines()[1:] + if line.strip() and not line.startswith('<')] + tags = {} + for particle in lines: + status, color, anticolor = int(particle[1]), int(particle[4]), int(particle[5]) + if status == -1: + color, anticolor = anticolor, color + for tag, sign in ((color, 1), (anticolor, -1)): + if tag: + tags[tag] = tags.get(tag, 0) + sign + if any(tags.values()): + return False + return True + + +class MG7AmplicolTest(unittest.TestCase): + + def setUp(self): + try: + import madspace # noqa: F401 + import pyamplicol # noqa: F401 + except ImportError as error: + self.skipTest('madspace or pyamplicol unavailable: %s' % error) + if not misc.resolve_lhapdf().find_set(_PDF): + self.skipTest('%s LHAPDF data not found' % _PDF) + self.path = tempfile.mkdtemp(prefix='mg7_amplicol_') + + def tearDown(self): + shutil.rmtree(self.path, ignore_errors=True) + + def run_mg7(self, name, processes, **amplicol): + """Output and run the processes, returns the cross section, its error + and the run directory.""" + run_dir = pjoin(self.path, name) + mg = MGCmd.MasterCmd() + mg.no_notification() + mg.exec_cmd('set automatic_html_opening False --no_save') + for index, process in enumerate(processes): + mg.exec_cmd('%s %s' % ('generate' if index == 0 else 'add process', + process)) + mg.exec_cmd('output mg7 %s' % run_dir) + _edit_run_card(pjoin(run_dir, 'Cards', 'run_card.toml'), **amplicol) + log = pjoin(run_dir, 'generate_events.log') + with open(log, 'w') as stream: + status = subprocess.call( + [sys.executable, pjoin(run_dir, 'bin', 'generate_events'), '-f'], + cwd=run_dir, stdout=stream, stderr=subprocess.STDOUT) + self.assertEqual(status, 0, open(log).read()[-3000:]) + info = sorted(glob.glob(pjoin(run_dir, 'Events', '*', 'info.json')))[-1] + with open(info) as stream: + result = json.load(stream)['process'] + return float(result['mean']), float(result['error']), run_dir + + def assert_compatible(self, first, second, factor=1.): + (mean1, error1), (mean2, error2) = first, second + error = (error1**2 + (factor * error2)**2)**0.5 + self.assertLess(abs(mean1 - factor * mean2), _NSIGMA * error, + '%g +- %g vs %g +- %g' % (mean1, error1, factor * mean2, + factor * error2)) + + def test_gg_ggg(self): + """Full colour in both modes, and leading colour, which differs by + (N^2 - 1) / N^2 for five gluons.""" + madmatrix = self.run_mg7('mm', ['g g > g g g'])[:2] + reweight = self.run_mg7('rw', ['g g > g g g [amplicol]']) + self.assert_compatible(reweight[:2], madmatrix) + events = glob.glob(pjoin(reweight[2], 'Events', '*', 'events.lhe.gz')) + if events: + self.assertTrue(_color_balanced(events[-1])) + unweight = self.run_mg7('uw', ['g g > g g g [amplicol]'], mode='unweight') + self.assert_compatible(unweight[:2], madmatrix) + leading = self.run_mg7('lc', ['g g > g g g [amplicol]'], + color_accuracy='lc') + self.assert_compatible(madmatrix, leading[:2], 8. / 9.) + + def test_pp_jjj(self): + madmatrix = self.run_mg7('mm', ['p p > j j j'])[:2] + amplicol = self.run_mg7('amp', ['p p > j j j [amplicol]'])[:2] + self.assert_compatible(amplicol, madmatrix) + + def test_mixed(self): + """[amplicol] and madmatrix processes in one output.""" + jj = self.run_mg7('jj', ['p p > j j'])[:2] + ggg = self.run_mg7('ggg', ['g g > g g g'])[:2] + mixed = self.run_mg7('mixed', ['p p > j j', 'g g > g g g [amplicol]'])[:2] + self.assert_compatible( + mixed, (jj[0] + ggg[0], (jj[1]**2 + ggg[1]**2)**0.5)) From 3ace55773a20e3430792f604a6e19b9292162b04 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:28:13 +0200 Subject: [PATCH 31/88] mg7: [amplicol] fixes for quark processes and models without alpha_s The LHE completion takes the masses of antiparticles too. A model without an alpha_s parameter can be used with a fixed renormalization scale, at the alpha_s of the model. The p p > j j j test is skipped until pyamplicol drops the processes without amplitudes from its multiparticle expansion and its UFO quark processes agree with madmatrix. --- .../template_files/mg7/subprocess_amplicol.py | 13 ++++++++----- tests/acceptance_tests/test_mg7_amplicol.py | 3 +++ 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py index eedabb5118..b75ca6b990 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py +++ b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py @@ -57,8 +57,8 @@ def amplicol_subproc_args(meta: dict, process: MadgraphProcess) -> ms.SubprocArg """LHE completion of an AmpliCol subprocess: its events have no diagram.""" return ms.SubprocArgs( topologies=[chain_topology( - [process.get_mass(pid) for pid in meta["incoming"]], - [process.get_mass(pid) for pid in meta["outgoing"]], + [process.get_mass(pid) for pid in clean_pids(meta["incoming"])], + [process.get_mass(pid) for pid in clean_pids(meta["outgoing"])], )], permutations=[[]], diagram_indices=[[]], @@ -102,10 +102,12 @@ def load_matrix_elements(self, all_api_paths, api_path_formats, member_metas): raise RuntimeError("[amplicol] processes need their matrix elements") if self.process.run_card["gridpack"]["save_gridpack"]: raise RuntimeError("[amplicol] processes do not support gridpacks yet") - if not self.meta.get("supports_alpha_s", True): + self.supports_alpha_s = self.meta.get("supports_alpha_s", True) + if not self.supports_alpha_s and not self.process.run_card["beam"]["fixed_ren_scale"]: raise RuntimeError( "the AmpliCol matrix elements of this model cannot take the running " - "alpha_s, as the model has no SMINPUTS 3 parameter" + "alpha_s, as the model has no SMINPUTS 3 parameter: use a fixed " + "renormalization scale" ) parameters = amplicol_parameters(self.meta, self.process) self.channel_index = self.meta["umami_channel"] @@ -160,9 +162,10 @@ def simplify_phasespace(self, multi_phasespace, cross_sections) -> PhaseSpace: return multi_phasespace def integrand_matrix_elements(self) -> list: + # without an alpha_s input, the model value of the fixed scale is used inputs = [ ms.MatrixElement.momenta_in, - ms.MatrixElement.alpha_s_in, + *([ms.MatrixElement.alpha_s_in] if self.supports_alpha_s else []), ms.MatrixElement.flavor_in, ms.MatrixElement.channel_in, ms.MatrixElement.random_helicity_in, diff --git a/tests/acceptance_tests/test_mg7_amplicol.py b/tests/acceptance_tests/test_mg7_amplicol.py index 539f31ffbd..d8fee330c8 100644 --- a/tests/acceptance_tests/test_mg7_amplicol.py +++ b/tests/acceptance_tests/test_mg7_amplicol.py @@ -144,6 +144,9 @@ def test_gg_ggg(self): color_accuracy='lc') self.assert_compatible(madmatrix, leading[:2], 8. / 9.) + @unittest.skip('pyamplicol expands p p > j j j into processes without ' + 'amplitudes, and its quark processes from the sm UFO model ' + 'differ from madmatrix (the built-in pyamplicol model agrees)') def test_pp_jjj(self): madmatrix = self.run_mg7('mm', ['p p > j j j'])[:2] amplicol = self.run_mg7('amp', ['p p > j j j [amplicol]'])[:2] From f0dd992defb9682e6e43d54d3769d1de2c3ee410 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:47:17 +0200 Subject: [PATCH 32/88] mg7: no madmatrix files in outputs with only [amplicol] processes Without MG diagrams the output gets no src/, backend/, madmatrix makefiles or model code, which nothing would build: only the cards, the launcher and the [amplicol] subprocesses. --- madgraph/iolibs/export_cpp.py | 2 ++ madmatrix/output.py | 8 ++++++++ 2 files changed, 10 insertions(+) diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index de2d2f116f..e5d766e9c0 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3509,6 +3509,8 @@ def pass_information_from_cmd(self, cmd): [amplicol] processes.""" self.proc_defs = getattr(cmd, '_curr_proc_defs', None) self.amplicol_defs = getattr(cmd, '_curr_amplicol_defs', None) or [] + # without MG diagrams there is no madmatrix code to write or build + self.has_madmatrix = bool(cmd._curr_amps) or not self.amplicol_defs self.amplicol_model = cmd._curr_model self.amplicol_multiparticles = dict(getattr(cmd, '_multiparticles', {})) diff --git a/madmatrix/output.py b/madmatrix/output.py index 6e9cc06a88..7b299e5576 100644 --- a/madmatrix/output.py +++ b/madmatrix/output.py @@ -216,6 +216,12 @@ def get_makefile_replace_dict(self, model): # AV - overload the default version: create CMake directory, do not create lib directory def copy_template(self, model): + if not getattr(self, 'has_madmatrix', True): + # only the cards and the launcher: no src/, backend/ or makefiles + self.dirs_to_create = ['bin', 'lib', 'Cards', 'SubProcesses'] + self.from_template = {key: self.from_template[key] for key in ('.', 'Cards')} + self.template_src_make = self.template_Sub_make = None + return super().copy_template(model) super().copy_template(model) # Copy Arithmetics headers for the double-word expansion (FPTYPE=e) arithmetics_src = pjoin(self.madmatrix_templates, 'Arithmetics') @@ -293,6 +299,8 @@ def generate_subprocess_directory(self, matrix_element, cpp_helas_call_writer, p def convert_model(self, model, wanted_lorentz=[], wanted_couplings=[], **opts): # **opts: an option meant for one exporter (npwave, for the dual HELAS # libraries of P-wave bound states) is passed by keyword to all of them + if not getattr(self, 'has_madmatrix', True): + return if hasattr(model , 'cudacpp_wanted_ordered_couplings'): wanted_couplings = model.cudacpp_wanted_ordered_couplings del model.cudacpp_wanted_ordered_couplings From 4de6aedd46852650ee4f41517dd399d564a7a480 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 15 Oct 2026 14:59:07 +0200 Subject: [PATCH 33/88] mg7: load the [amplicol] libraries once for all their channels Each colour-ordered channel is its own subprocess, and each loaded the same two libraries again, with an instance per thread that reads the whole pyamplicol artifact. --- .../iolibs/template_files/mg7/madgraph_process.py | 2 ++ .../iolibs/template_files/mg7/subprocess_amplicol.py | 12 ++++++++---- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/madgraph_process.py b/madgraph/iolibs/template_files/mg7/madgraph_process.py index a7de60ee76..0f910aba49 100644 --- a/madgraph/iolibs/template_files/mg7/madgraph_process.py +++ b/madgraph/iolibs/template_files/mg7/madgraph_process.py @@ -251,6 +251,8 @@ def init_event_dir(self) -> None: def init_context(self) -> None: device_names = self.run_card["run"]["device"] self.contexts = [] + # AmpliCol libraries, shared by all channels of their process + self.amplicol_apis = {} self.device_types = [] self.devices = [] self.pool_sizes = [] diff --git a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py index b75ca6b990..807edefd08 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py +++ b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py @@ -122,10 +122,14 @@ def load_matrix_elements(self, all_api_paths, api_path_formats, member_metas): ) def load(self, path: str, artifact: str, parameters: dict): - # the libraries find their pyamplicol artifacts from an absolute path - for context in self.process.contexts: - api = context.load_matrix_element(path, os.path.abspath(artifact), parameters) - return api + # loaded once for all channels; each load reads the whole pyamplicol artifact + apis = self.process.amplicol_apis + if path not in apis: + for context in self.process.contexts: + apis[path] = context.load_matrix_element( + path, os.path.abspath(artifact), parameters + ) + return apis[path] def build_multi_channel_data(self): return None From d0567581811ccfcb9f7082a9d793fa1c1eb47356 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 15 Oct 2026 15:41:41 +0200 Subject: [PATCH 34/88] mg7: mirrored beams of the [amplicol] processes pyamplicol keeps one process per unordered initial state, so p p > j j j had g u > ... but not u g > ..., and its cross section was 21% low. Such flavours are integrated with the mirrored beams, as madmatrix does, when the process definition allows the other order. --- madgraph/iolibs/export_amplicol.py | 27 +++++++++++++++---- .../interface/test_amplicol_syntax.py | 15 +++++++++++ 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index fe3315fb6c..f8b919587a 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -206,11 +206,25 @@ def color_types(model, pdgs): return types -def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power): - """One subprocesses.json entry per channel of a UMAMI provider.""" +def mirrored(pdgs, incoming, initial_legs, exported): + """Whether the process with swapped beams is requested but not exported: + pyamplicol keeps one process per unordered initial state.""" + if incoming != 2 or pdgs[0] == pdgs[1]: + return False + swapped = ((pdgs[1], pdgs[0]), tuple(sorted(pdgs[2:]))) + return (pdgs[1] in initial_legs[0] and pdgs[0] in initial_legs[1] + and swapped not in exported) + + +def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power, + initial_legs): + """One subprocesses.json entry per channel of a UMAMI provider; initial_legs + are the pdg ids allowed for each initial leg.""" pdg_ids = metadata['pdg_ids'] provider = metadata['provider'] incoming = provider['incoming_count'] + exported = {(tuple(pdgs[:incoming]), tuple(sorted(pdgs[incoming:]))) + for pdgs in pdg_ids} entries = [] for channel_index, channel in enumerate(metadata['channels']): flavors = [] @@ -220,11 +234,12 @@ def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power) 'AmpliCol contributions shared by several channels are not ' 'supported yet') for term in process['matrix_elements']: + pdgs = pdg_ids[term['pdg_ids']] flavors.append({ 'index': flavor_index, - 'options': [pdg_ids[term['pdg_ids']]], + 'options': [pdgs], 'factor': term['factor'], - 'mirror': False, + 'mirror': mirrored(pdgs, incoming, initial_legs, exported), }) first = flavors[0]['options'][0] entries.append({ @@ -300,8 +315,10 @@ def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None): copies.append('\tcp $(BUILD)/lc/%s ../../%s' % (library, me_path)) copies.append('\tcp $(BUILD)/reweight/libreweight_%s.so ../../%s' % (provider_id, ratio_path)) + initial_legs = [set(leg_pdgs(leg, model)) + for leg in procdef.get('legs') if not leg.get('state')] entries.extend(subprocess_entries(metadata, model, rel_dir, me_path, - ratio_path, qcd_power)) + ratio_path, qcd_power, initial_legs)) with open(pjoin(path, 'makefile'), 'w') as stream: stream.write(MAKEFILE % {'rusticol_config': rusticol_config, 'copies': '\n'.join(copies)}) diff --git a/tests/unit_tests/interface/test_amplicol_syntax.py b/tests/unit_tests/interface/test_amplicol_syntax.py index d52a290ad6..854b7bb3ff 100644 --- a/tests/unit_tests/interface/test_amplicol_syntax.py +++ b/tests/unit_tests/interface/test_amplicol_syntax.py @@ -136,3 +136,18 @@ def test_coupling_orders(self): precmd=True) self.assertEqual(export_amplicol.coupling_orders( self.cmd._curr_amplicol_defs[0], self.cmd._curr_model), orders) + + def test_mirrored_beams(self): + """pyamplicol keeps one process per unordered initial state; the other + beam order is integrated as the mirror where the definition allows it.""" + import madgraph.iolibs.export_amplicol as export_amplicol + mirrored = export_amplicol.mirrored + partons = {21, 1, 2, -1, -2} + exported = {((21, 2), (2, 21, 21)), ((2, 2), (2, 2, 21)), + ((2, -2), (-2, 2, 21)), ((-2, 2), (-2, 2, 21))} + self.assertTrue(mirrored([21, 2, 21, 21, 2], 2, [partons, partons], exported)) + self.assertFalse(mirrored([2, 2, 2, 2, 21], 2, [partons, partons], exported)) + # both beam orders are exported + self.assertFalse(mirrored([2, -2, 2, -2, 21], 2, [partons, partons], exported)) + # g u > u g g only + self.assertFalse(mirrored([21, 2, 21, 21, 2], 2, [{21}, {2}], exported)) From 8f045497cecd727c5eb60c15f5136c1f73fbe9cd Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 15 Oct 2026 15:41:43 +0200 Subject: [PATCH 35/88] tests: [amplicol] p p > j j j and g g > t t~ g against madmatrix --- tests/acceptance_tests/test_mg7_amplicol.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/acceptance_tests/test_mg7_amplicol.py b/tests/acceptance_tests/test_mg7_amplicol.py index d8fee330c8..880de0ab8c 100644 --- a/tests/acceptance_tests/test_mg7_amplicol.py +++ b/tests/acceptance_tests/test_mg7_amplicol.py @@ -144,14 +144,18 @@ def test_gg_ggg(self): color_accuracy='lc') self.assert_compatible(madmatrix, leading[:2], 8. / 9.) - @unittest.skip('pyamplicol expands p p > j j j into processes without ' - 'amplitudes, and its quark processes from the sm UFO model ' - 'differ from madmatrix (the built-in pyamplicol model agrees)') def test_pp_jjj(self): + """Quark lines and the expansion of the multiparticle labels.""" madmatrix = self.run_mg7('mm', ['p p > j j j'])[:2] amplicol = self.run_mg7('amp', ['p p > j j j [amplicol]'])[:2] self.assert_compatible(amplicol, madmatrix) + def test_gg_ttxg(self): + """Massive quarks in the colour-ordered phase space.""" + madmatrix = self.run_mg7('mm', ['g g > t t~ g'])[:2] + amplicol = self.run_mg7('amp', ['g g > t t~ g [amplicol]'])[:2] + self.assert_compatible(amplicol, madmatrix) + def test_mixed(self): """[amplicol] and madmatrix processes in one output.""" jj = self.run_mg7('jj', ['p p > j j'])[:2] From 2d8f6039365891c130efdc41ccb7675685901f62 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 15 Oct 2026 20:14:20 +0200 Subject: [PATCH 36/88] mg7: generate the [amplicol] artifacts in a worker process Symbolica, imported by pyamplicol, holds a per-user lock until its process exits, so a MadGraph session that had output [amplicol] processes blocked all others. pyamplicol now only runs in amplicol_worker.py, which waits while another process holds the lock, and its workers and cores are left to pyamplicol, licensed through its Symbolica library key. --- madgraph/iolibs/amplicol_worker.py | 81 +++++++++++ madgraph/iolibs/export_amplicol.py | 146 ++++++++++++-------- tests/acceptance_tests/test_mg7_amplicol.py | 11 +- 3 files changed, 176 insertions(+), 62 deletions(-) create mode 100644 madgraph/iolibs/amplicol_worker.py diff --git a/madgraph/iolibs/amplicol_worker.py b/madgraph/iolibs/amplicol_worker.py new file mode 100644 index 0000000000..74001176c6 --- /dev/null +++ b/madgraph/iolibs/amplicol_worker.py @@ -0,0 +1,81 @@ +################################################################################ +# +# Copyright (c) 2026 The MadGraph5_aMC@NLO Development team and Contributors +# +# This file is a part of the MadGraph5_aMC@NLO project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph5_aMC@NLO license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""Generation of the '[amplicol]' artifacts with pyamplicol, run in its own +process by export_amplicol. + +Symbolica, which pyamplicol imports, holds a per-user lock until the process +exits, so it is never imported by MadGraph itself: sessions do not block each +other while they are open. Usage: python amplicol_worker.py JOB.json +""" + +import json +import logging +import sys + + +def model_source(pyamplicol, model): + """The model compiled for pyamplicol, cached in the output.""" + import os + if not os.path.exists(model['cache']): + from pyamplicol.config import EvaluatorConfig, JITConfig + source = pyamplicol.ModelSource.from_path(model['path'], + restriction=model['restriction']) + source.compile(prepared_output=model['cache'], + evaluator=EvaluatorConfig(jit=JITConfig(optimization_level=2))) + return pyamplicol.ModelSource.from_path(model['cache']) + + +def generate_artifact(pyamplicol, artifact, accuracy, process, model): + from pyamplicol.config import resolve_config + # workers and cores are left to pyamplicol, which clamps them without a license + config = resolve_config({ + 'schema_version': 1, + 'action': 'generate', + 'color': {'accuracy': accuracy}, + 'process': {'multiparticles': process['multiparticles'], + 'coupling_order_policy': 'explicit', + 'max_coupling_orders': process['orders']}, + 'generation': {'emit_api_bundle': True, + 'relation_discovery': {'mode': 'off'}}, + 'evaluator': {'backend': 'jit', 'execution_mode': 'recurrence', + 'jit': {'optimization_level': 2}}, + }) + pyamplicol.Generator(config).generate(process['expression'], artifact, + model=model, mode='replace') + + +def main(job_path): + import os + logging.basicConfig(level=logging.WARNING, format='%(name)s: %(message)s') + with open(job_path) as stream: + job = json.load(stream) + import pyamplicol + from pyamplicol.artifacts import umami_reweight + model = model_source(pyamplicol, job['model']) + providers = [] + for process in job['processes']: + path = process['path'] + for accuracy in ('lc', 'full'): + generate_artifact(pyamplicol, os.path.join(path, accuracy), accuracy, + process, model) + providers.append(umami_reweight.write_reweight_provider( + os.path.join(path, 'lc'), os.path.join(path, 'full'), + os.path.join(path, 'reweight'))) + with open(job['result'], 'w') as stream: + json.dump(providers, stream) + + +if __name__ == '__main__': + main(sys.argv[1]) diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index f8b919587a..c278c40453 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -22,11 +22,16 @@ on its own with a colour-ordered phase-space mapping. """ +import importlib.util import json import logging import os import re import shutil +import subprocess +import sys +import tempfile +import time import madgraph @@ -49,15 +54,61 @@ """ -def import_pyamplicol(): - try: - import pyamplicol - from pyamplicol.artifacts import umami_reweight - except ImportError as error: +WORKER = pjoin(os.path.dirname(os.path.abspath(__file__)), 'amplicol_worker.py') +# printed by Symbolica when another process holds its per-user thread +SYMBOLICA_BUSY = 'Cannot start another restricted Symbolica thread' + + +def check_pyamplicol(): + """pyamplicol is only imported by the worker, see amplicol_worker.""" + if importlib.util.find_spec('pyamplicol') is None: raise madgraph.InvalidCmd( - 'The [amplicol] processes need pyamplicol (%s). Install it in the ' - 'Python environment of MadGraph.' % error) - return pyamplicol, umami_reweight + 'The [amplicol] processes need pyamplicol. Install it in the Python ' + 'environment of MadGraph.') + + +def run_worker(job): + """Run amplicol_worker on a job, waiting while Symbolica is busy elsewhere.""" + env = dict(os.environ, SYMBOLICA_HIDE_BANNER='1', + PYTHONPATH=os.pathsep.join(p for p in sys.path if p)) + waiting = False + while True: + worker = subprocess.Popen([sys.executable, WORKER, job], env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True) + busy = False + for line in worker.stdout: + if SYMBOLICA_BUSY in line: + busy = True + elif not line.startswith(('\u250c', '\u2502', '\u2514', '\u251c')): + sys.stdout.write(line) + if worker.wait() == 0: + return + if not busy: + raise madgraph.InvalidCmd('pyamplicol failed to generate the [amplicol] ' + 'matrix elements, see the messages above') + if 'symbolica' in sys.modules: + raise madgraph.InvalidCmd( + 'Symbolica is imported in this MadGraph process, so the pyamplicol ' + 'worker cannot start: do not import pyamplicol or symbolica here') + if not waiting: + logger.info('Symbolica is used by another process: waiting for it') + waiting = True + time.sleep(10) + + +def load_metadata(artifact): + """The UMAMI metadata documents of an artifact, one per provider.""" + root = pjoin(artifact, 'API', 'umami') + with open(pjoin(root, 'metadata.json')) as stream: + data = json.load(stream) + if data.get('kind') != 'pyamplicol-umami-provider-index': + return [data] + documents = [] + for provider in data['providers']: + with open(pjoin(root, provider['metadata'])) as stream: + documents.append(json.load(stream)) + return documents def leg_pdgs(leg, model): @@ -130,11 +181,9 @@ def coupling_orders(procdef, model): for order, weight in hierarchy.items() if weight > lowest} -def model_source(pyamplicol, model, cache_path): +def model_job(model, cache_path): """The UFO model of the output with its restriction, compiled for pyamplicol once per output.""" - if os.path.exists(cache_path): - return pyamplicol.ModelSource.from_path(cache_path) path = model.get('modelpath') reference = os.path.basename(model.get('modelpath+restriction')) restriction = 'default' @@ -144,33 +193,9 @@ def model_source(pyamplicol, model, cache_path): restriction = 'none' else: restriction = pjoin(path, 'restrict_%s.dat' % suffix) - from pyamplicol.config import EvaluatorConfig, JITConfig - source = pyamplicol.ModelSource.from_path(path, restriction=restriction) - logger.info('Compiling the model %s for AmpliCol' % reference) - source.compile(prepared_output=cache_path, - evaluator=EvaluatorConfig(jit=JITConfig(optimization_level=2))) - return pyamplicol.ModelSource.from_path(cache_path) - - -def generate_artifact(pyamplicol, artifact, accuracy, expression, multiparticles, - orders, model): - from pyamplicol.config import resolve_config - # one worker: Symbolica's restricted mode allows a single thread - config = resolve_config({ - 'schema_version': 1, - 'action': 'generate', - 'color': {'accuracy': accuracy}, - 'process': {'multiparticles': multiparticles, - 'coupling_order_policy': 'explicit', - 'max_coupling_orders': orders}, - 'generation': {'workers': 1, 'emit_api_bundle': True, - 'relation_discovery': {'mode': 'off'}}, - 'evaluator': {'backend': 'jit', 'execution_mode': 'recurrence', - 'optimization': {'cores': 1}, - 'jit': {'optimization_level': 2}}, - }) - pyamplicol.Generator(config).generate(expression, artifact, model=model, - mode='replace') + if not os.path.exists(cache_path): + logger.info('Compiling the model %s for AmpliCol' % reference) + return {'path': path, 'restriction': restriction, 'cache': cache_path} def representative_processes(proc_defs, model): @@ -283,31 +308,40 @@ def parameter_references(metadata): def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None): """Write SubProcesses/A_ for every '[amplicol]' process definition and return the subprocesses.json entries.""" - pyamplicol, umami_reweight = import_pyamplicol() + check_pyamplicol() subproc_dir = pjoin(dir_path, 'SubProcesses') - source = model_source(pyamplicol, model, pjoin(subproc_dir, MODEL_FILE)) rusticol_config = shutil.which('rusticol-config') or 'rusticol-config' - entries = [] + processes = [] for number, procdef in enumerate(proc_defs, start=1): name = 'A%d_%s' % (number, process_name(procdef, model, multiparticles or {})) path = pjoin(subproc_dir, name) + expression, process_multiparticles = process_expression(procdef, model) + os.makedirs(path, exist_ok=True) + logger.info('Generating the AmpliCol matrix elements of %s in %s' % + (procdef.nice_string(prefix=False), path)) + processes.append({'name': name, 'path': path, 'expression': expression, + 'multiparticles': process_multiparticles, + 'orders': coupling_orders(procdef, model)}) + with tempfile.TemporaryDirectory(prefix='amplicol_job_') as scratch: + job = pjoin(scratch, 'job.json') + with open(job, 'w') as stream: + json.dump({'model': model_job(model, pjoin(subproc_dir, MODEL_FILE)), + 'processes': processes, + 'result': pjoin(scratch, 'providers.json')}, stream) + run_worker(job) + with open(pjoin(scratch, 'providers.json')) as stream: + providers = json.load(stream) + entries = [] + for procdef, process, provider_ids in zip(proc_defs, processes, providers): + name, path, orders = process['name'], process['path'], process['orders'] rel_dir = pjoin('SubProcesses', name) - expression, multiparticles = process_expression(procdef, model) - orders = coupling_orders(procdef, model) # every vertex carries one coupling: the rest are powers of g_s qcd_power = (len(procdef.get('legs')) - 2 - sum(orders.values()) if len(orders) == 1 else -1) - logger.info('Generating the AmpliCol matrix elements of %s in %s' % - (procdef.nice_string(prefix=False), path)) - os.makedirs(path, exist_ok=True) - for accuracy in ('lc', 'full'): - generate_artifact(pyamplicol, pjoin(path, accuracy), accuracy, - expression, multiparticles, orders, source) - provider_ids = umami_reweight.write_reweight_provider( - pjoin(path, 'lc'), pjoin(path, 'full'), pjoin(path, 'reweight')) - documents = umami_reweight.load_umami_metadata(pjoin(path, 'lc')) + initial_legs = [set(leg_pdgs(leg, model)) + for leg in procdef.get('legs') if not leg.get('state')] copies = [] - for provider_id, metadata in zip(provider_ids, documents): + for provider_id, metadata in zip(provider_ids, load_metadata(pjoin(path, 'lc'))): library = metadata['provider']['library'] me_path = pjoin('lib', 'libamplicol_%s_%s.so' % (name, provider_id)) ratio_path = pjoin('lib', 'libamplicol_%s_%s_ratio.so' % @@ -315,8 +349,6 @@ def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None): copies.append('\tcp $(BUILD)/lc/%s ../../%s' % (library, me_path)) copies.append('\tcp $(BUILD)/reweight/libreweight_%s.so ../../%s' % (provider_id, ratio_path)) - initial_legs = [set(leg_pdgs(leg, model)) - for leg in procdef.get('legs') if not leg.get('state')] entries.extend(subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power, initial_legs)) with open(pjoin(path, 'makefile'), 'w') as stream: @@ -324,7 +356,7 @@ def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None): 'copies': '\n'.join(copies)}) with open(pjoin(path, 'process.json'), 'w') as stream: json.dump({'process': procdef.nice_string(prefix=False), - 'expression': expression, - 'multiparticles': multiparticles, + 'expression': process['expression'], + 'multiparticles': process['multiparticles'], 'max_coupling_orders': orders}, stream, indent=2) return entries diff --git a/tests/acceptance_tests/test_mg7_amplicol.py b/tests/acceptance_tests/test_mg7_amplicol.py index 880de0ab8c..802aa2ca72 100644 --- a/tests/acceptance_tests/test_mg7_amplicol.py +++ b/tests/acceptance_tests/test_mg7_amplicol.py @@ -26,6 +26,7 @@ import glob import gzip +import importlib.util import json import os import re @@ -87,11 +88,11 @@ def _color_balanced(lhe_path): class MG7AmplicolTest(unittest.TestCase): def setUp(self): - try: - import madspace # noqa: F401 - import pyamplicol # noqa: F401 - except ImportError as error: - self.skipTest('madspace or pyamplicol unavailable: %s' % error) + # not imported: Symbolica would hold its lock against the pyamplicol worker + missing = [name for name in ('madspace', 'pyamplicol') + if importlib.util.find_spec(name) is None] + if missing: + self.skipTest('%s unavailable' % ' and '.join(missing)) if not misc.resolve_lhapdf().find_set(_PDF): self.skipTest('%s LHAPDF data not found' % _PDF) self.path = tempfile.mkdtemp(prefix='mg7_amplicol_') From 3ba676e832a0847e72ac963a97fe977f357bac90 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 26 Oct 2026 12:55:27 +0100 Subject: [PATCH 37/88] mg7: [amplicol] with the pyamplicol of pyamplicol/install The output installs pyamplicol there if needed (bootstrap.ensure_pyamplicol) and puts it first on the worker's path. The makefiles call a SubProcesses/rusticol-config that runs its rusticol_config, and the worker warns when Symbolica does not accept the package key. check_pyamplicol_source_hash reads the hash file instead of importing pyamplicol, which took Symbolica's lock in MadGraph's process and blocked the worker. The tests check that the runs finished. --- madgraph/iolibs/amplicol_worker.py | 4 ++ madgraph/iolibs/export_amplicol.py | 51 +++++++++++++------ madgraph/iolibs/export_cpp.py | 8 ++- .../iolibs/template_files/mg7/bootstrap.py | 16 ++++-- tests/acceptance_tests/test_mg7_amplicol.py | 17 ++++--- 5 files changed, 68 insertions(+), 28 deletions(-) diff --git a/madgraph/iolibs/amplicol_worker.py b/madgraph/iolibs/amplicol_worker.py index 74001176c6..b546a6284a 100644 --- a/madgraph/iolibs/amplicol_worker.py +++ b/madgraph/iolibs/amplicol_worker.py @@ -63,6 +63,10 @@ def main(job_path): job = json.load(stream) import pyamplicol from pyamplicol.artifacts import umami_reweight + from pyamplicol.licensing import detect_symbolica_license + if not detect_symbolica_license(suggest=False).licensed: + logging.warning('Symbolica does not accept the license key of pyamplicol ' + '(%s): generating in its restricted mode', pyamplicol.__file__) model = model_source(pyamplicol, job['model']) providers = [] for process in job['processes']: diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index c278c40453..84141aa7f7 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -22,12 +22,11 @@ on its own with a colour-ordered phase-space mapping. """ -import importlib.util import json import logging import os import re -import shutil +import shlex import subprocess import sys import tempfile @@ -42,7 +41,7 @@ MAKEFILE = """\ # Builds the AmpliCol matrix elements of this process into ../../lib. -RUSTICOL_CONFIG ?= %(rusticol_config)s +RUSTICOL_CONFIG ?= $(CURDIR)/../rusticol-config BUILD := $(CURDIR)/build all: @@ -59,12 +58,33 @@ SYMBOLICA_BUSY = 'Cannot start another restricted Symbolica thread' -def check_pyamplicol(): - """pyamplicol is only imported by the worker, see amplicol_worker.""" - if importlib.util.find_spec('pyamplicol') is None: - raise madgraph.InvalidCmd( - 'The [amplicol] processes need pyamplicol. Install it in the Python ' - 'environment of MadGraph.') +# rusticol-config of the pyamplicol in pyamplicol/install, for the makefiles +RUSTICOL_CONFIG = """\ +#!/bin/sh +PYTHONPATH=%(install)s${PYTHONPATH:+:$PYTHONPATH} exec %(python)s -m rusticol_config "$@" +""" + + +def check_pyamplicol(interactive=None, jobs=None): + """Install pyamplicol in pyamplicol/install if needed and put it first on + sys.path, which the worker inherits. Never imported here, see amplicol_worker.""" + from madgraph.iolibs.template_files.mg7 import bootstrap + try: + bootstrap.ensure_pyamplicol(interactive=interactive, jobs=jobs) + except RuntimeError as error: + raise madgraph.InvalidCmd('The [amplicol] processes need pyamplicol: %s' % error) + bootstrap.check_pyamplicol_source_hash() + + +def write_rusticol_config(subproc_dir): + """SubProcesses/rusticol-config, the RUSTICOL_CONFIG of the makefiles.""" + from madgraph.iolibs.template_files.mg7 import bootstrap + path = pjoin(subproc_dir, 'rusticol-config') + with open(path, 'w') as stream: + stream.write(RUSTICOL_CONFIG % { + 'install': shlex.quote(str(bootstrap.PYAMPLICOL_INSTALL_DIR)), + 'python': shlex.quote(sys.executable)}) + os.chmod(path, 0o755) def run_worker(job): @@ -305,12 +325,14 @@ def parameter_references(metadata): return references -def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None): +def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None, + interactive=None, jobs=None): """Write SubProcesses/A_ for every '[amplicol]' process definition - and return the subprocesses.json entries.""" - check_pyamplicol() + and return the subprocesses.json entries. ``interactive`` and ``jobs`` are + those of the pyamplicol installation, see bootstrap.ensure_pyamplicol.""" + check_pyamplicol(interactive, jobs) subproc_dir = pjoin(dir_path, 'SubProcesses') - rusticol_config = shutil.which('rusticol-config') or 'rusticol-config' + write_rusticol_config(subproc_dir) processes = [] for number, procdef in enumerate(proc_defs, start=1): name = 'A%d_%s' % (number, process_name(procdef, model, multiparticles or {})) @@ -352,8 +374,7 @@ def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None): entries.extend(subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power, initial_legs)) with open(pjoin(path, 'makefile'), 'w') as stream: - stream.write(MAKEFILE % {'rusticol_config': rusticol_config, - 'copies': '\n'.join(copies)}) + stream.write(MAKEFILE % {'copies': '\n'.join(copies)}) with open(pjoin(path, 'process.json'), 'w') as stream: json.dump({'process': procdef.nice_string(prefix=False), 'expression': process['expression'], diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index e5d766e9c0..55c122f543 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3513,6 +3513,11 @@ def pass_information_from_cmd(self, cmd): self.has_madmatrix = bool(cmd._curr_amps) or not self.amplicol_defs self.amplicol_model = cmd._curr_model self.amplicol_multiparticles = dict(getattr(cmd, '_multiparticles', {})) + # how pyamplicol is installed if it is missing + self.amplicol_install = { + 'interactive': bool(getattr(cmd, 'use_rawinput', False)), + 'jobs': getattr(cmd, 'options', {}).get('nb_core'), + } def amplicol_processes(self): """MG processes standing for the [amplicol] processes.""" @@ -3529,7 +3534,8 @@ def export_amplicol(self): import madgraph.iolibs.export_amplicol as export_amplicol for entry in export_amplicol.export_amplicol_processes( self.dir_path, self.amplicol_defs, self.amplicol_model, - self.amplicol_multiparticles): + self.amplicol_multiparticles, + **getattr(self, 'amplicol_install', {})): # never merged with another subprocess self.merged_subprocesses[('amplicol', len(self.process_info))].append( (len(self.process_info), [])) diff --git a/madgraph/iolibs/template_files/mg7/bootstrap.py b/madgraph/iolibs/template_files/mg7/bootstrap.py index 416415d5bf..74fae15d21 100644 --- a/madgraph/iolibs/template_files/mg7/bootstrap.py +++ b/madgraph/iolibs/template_files/mg7/bootstrap.py @@ -19,9 +19,10 @@ safe to do lazily. """ -import importlib +import importlib.util import json import os +import re import subprocess import sys from pathlib import Path @@ -118,10 +119,15 @@ def check_pyamplicol_source_hash() -> bool: Returns whether the hashes match. Must run after :func:`ensure_pyamplicol`. """ - try: - installed = importlib.import_module("pyamplicol._source_hash").SOURCE_HASH - except ImportError: - installed = None + # read, not imported: importing pyamplicol would take Symbolica's per-user + # lock in this process and block the generation worker + installed = None + spec = importlib.util.find_spec("pyamplicol") + if spec is not None and spec.origin: + hash_file = Path(spec.origin).parent / "_source_hash.py" + if hash_file.is_file(): + match = re.search(r'SOURCE_HASH\s*=\s*"([0-9a-f]*)"', hash_file.read_text()) + installed = match.group(1) if match else None source = subprocess.run( [sys.executable, str(PYAMPLICOL_DIR / "source_hash.py")], capture_output=True, text=True, check=True, diff --git a/tests/acceptance_tests/test_mg7_amplicol.py b/tests/acceptance_tests/test_mg7_amplicol.py index 88d5e349ad..e5f5559cd1 100644 --- a/tests/acceptance_tests/test_mg7_amplicol.py +++ b/tests/acceptance_tests/test_mg7_amplicol.py @@ -17,16 +17,15 @@ The AmpliCol subprocesses are integrated at leading colour and re- or unweighted to full colour, so their cross sections must agree with the madmatrix ones within the Monte Carlo errors, alone and mixed with madmatrix -processes in one output. Needs madspace, pyamplicol and LHAPDF; the tests skip -otherwise. Knobs: MG7_AMPLICOL_EVENTS (default 10000) and MG7_AMPLICOL_NSIGMA -(default 4). +processes in one output. Needs madspace, pyamplicol (in pyamplicol/install) and +LHAPDF; the tests skip otherwise. Knobs: MG7_AMPLICOL_EVENTS (default 10000) +and MG7_AMPLICOL_NSIGMA (default 4). ./tests/test_manager.py test_mg7_amplicol -pA -t0 """ import glob import gzip -import importlib.util import json import os import re @@ -88,11 +87,13 @@ def _color_balanced(lhe_path): class MG7AmplicolTest(unittest.TestCase): def setUp(self): + from madgraph.iolibs.template_files.mg7 import bootstrap # not imported: Symbolica would hold its lock against the pyamplicol worker - missing = [name for name in ('madspace', 'pyamplicol') - if importlib.util.find_spec(name) is None] + missing = [name for name, installed in ( + ('madspace', bootstrap.madspace_is_installed()), + ('pyamplicol', bootstrap.pyamplicol_is_installed())) if not installed] if missing: - self.skipTest('%s unavailable' % ' and '.join(missing)) + self.skipTest('%s not installed' % ' and '.join(missing)) if not misc.resolve_lhapdf().find_set(_PDF): self.skipTest('%s LHAPDF data not found' % _PDF) self.path = tempfile.mkdtemp(prefix='mg7_amplicol_') @@ -121,6 +122,8 @@ def run_mg7(self, name, processes, **amplicol): info = sorted(glob.glob(pjoin(run_dir, 'Events', '*', 'info.json')))[-1] with open(info) as stream: result = json.load(stream)['process'] + # a failed generation can still exit with 0, after the survey + self.assertTrue(result['done'], open(log).read()[-3000:]) return float(result['mean']), float(result['error']), run_dir def assert_compatible(self, first, second, factor=1.): From 988523af37aef087209c580815fbf4f2928200be Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 26 Oct 2026 12:55:33 +0100 Subject: [PATCH 38/88] pyamplicol: umami.h is madspace's A symlink to madspace/include/madspace/umami.h, as for madmatrix, which the build backend copies by content. main changed the scale outputs of the header, so the libraries built with pyamplicol's copy rejected the output keys of madspace. --- .../assets/api_templates/umami/umami.h | 325 +----------------- .../tests/unit/test_repository_policy.py | 2 + 2 files changed, 3 insertions(+), 324 deletions(-) mode change 100644 => 120000 pyamplicol/src/pyamplicol/assets/api_templates/umami/umami.h diff --git a/pyamplicol/src/pyamplicol/assets/api_templates/umami/umami.h b/pyamplicol/src/pyamplicol/assets/api_templates/umami/umami.h deleted file mode 100644 index 4eb2dba380..0000000000 --- a/pyamplicol/src/pyamplicol/assets/api_templates/umami/umami.h +++ /dev/null @@ -1,324 +0,0 @@ -/* - * _ - * (_) - * _ _ _ __ ___ __ _ _ __ ___ _ - * | | | | '_ ` _ \ / _` | '_ ` _ \| | - * | |_| | | | | | | (_| | | | | | | | - * \__,_|_| |_| |_|\__,_|_| |_| |_|_| - * - * Unified MAtrix eleMent Interface - * - * - */ - -#ifndef UMAMI_HEADER -#define UMAMI_HEADER 1 - -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif - -/** - * Major version number of the UMAMI interface. If the major version is the same - * between caller and implementation, binary compatibility is ensured. - */ -#define UMAMI_MAJOR_VERSION 1 -/** - * Minor version number of the UMAMI interface. Between minor versions, new keys for - * errors, devices, metadata, inputs and outputs can be added. - */ -#define UMAMI_MINOR_VERSION 2 - -typedef enum { - /** operation was executed successfully */ - UMAMI_SUCCESS, - /** an unspecified error */ - UMAMI_ERROR, - /** operation not implemented */ - UMAMI_ERROR_NOT_IMPLEMENTED, - /** the provided input key is not supported by this matrix element implementation */ - UMAMI_ERROR_UNSUPPORTED_INPUT, - /** the provided output key is not supported by this matrix element implementation - */ - UMAMI_ERROR_UNSUPPORTED_OUTPUT, - /** the meta_key was not yet initialized */ - UMAMI_ERROR_UNINITIALIZED_META, - /** the provided metadata key is not supported by this matrix element implementation - */ - UMAMI_ERROR_UNSUPPORTED_META, - /** a mandatory matrix element input was not provided */ - UMAMI_ERROR_MISSING_INPUT, -} UmamiStatus; - -typedef enum { - UMAMI_DEVICE_CPU, - UMAMI_DEVICE_CUDA, - UMAMI_DEVICE_HIP, -} UmamiDevice; - -typedef enum { - /** `UmamiDevice` specifying the type of device */ - UMAMI_META_DEVICE, - /** `int` specifying the number of external particles */ - UMAMI_META_PARTICLE_COUNT, - /** `int` specifying the number of Feynman diagrams */ - UMAMI_META_DIAGRAM_COUNT, - /** `int` specifying the number of helicities */ - UMAMI_META_HELICITY_COUNT, - /** `int` specifying the number of colors */ - UMAMI_META_COLOR_COUNT, - /** `double` array of length particle count: the external-leg masses (GeV) */ - UMAMI_META_MASSES, -} UmamiMetaKey; - -typedef enum { - /** momenta of the external legs, type: `double`, shape: `(particle count, 4)` */ - UMAMI_IN_MOMENTA, - /** value for the strong coupling, type: `double`, shape: `()` */ - UMAMI_IN_ALPHA_S, - /** flavor index, type: `int`, shape: `()` */ - UMAMI_IN_FLAVOR_INDEX, - /** random number for color selection, type: `double`, shape: `()` */ - UMAMI_IN_RANDOM_COLOR, - /** random number for helicity selection, type: `double`, shape: `()` */ - UMAMI_IN_RANDOM_HELICITY, - /** random number for diagram selection, type: `double`, shape: `()` */ - UMAMI_IN_RANDOM_DIAGRAM, - /** externally selected helicity index, type: `int`, shape: `()` */ - UMAMI_IN_HELICITY_INDEX, - /** externally selected diagram index, type: `unsigned int`, shape: `()` */ - UMAMI_IN_DIAGRAM_INDEX, - /** externally selected channel index, type: `unsigned int`, shape: `()` */ - UMAMI_IN_CHANNEL_INDEX, - /** number of external legs, type: `int`, shape: `()` */ - UMAMI_IN_PARTICLE_COUNT, - /** PDG ids of the external legs, type: `int`, shape: `(particle count)` */ - UMAMI_IN_PDG_IDS, - /** momentum fraction of an incoming parton, type: `double`, shape: `()` */ - UMAMI_IN_MOMENTUM_FRACTION, - /** energy scale, e.g. the factorization scale, type: `double`, shape: `()` */ - UMAMI_IN_SCALE, - /** PDG id of a single parton, type: `int`, shape: `()` */ - UMAMI_IN_PDG_ID, - /** index of the beam (0 or 1), type: `int`, shape: `()` */ - UMAMI_IN_BEAM_INDEX, -} UmamiInputKey; - -/** Number of values in `UmamiInputKey` */ -#define UMAMI_INPUT_KEY_COUNT (UMAMI_IN_BEAM_INDEX + 1) - -typedef enum { - /** value of the matrix element, type: `double`, shape: `()` */ - UMAMI_OUT_MATRIX_ELEMENT, - /** squared amplitudes of the diagrams, type: `double`, shape: `(diagram count)` */ - UMAMI_OUT_DIAGRAM_AMP2, - /** selected color index, type: `int`, shape: `()` */ - UMAMI_OUT_COLOR_INDEX, - /** selected helicity index, type: `int`, shape: `()` */ - UMAMI_OUT_HELICITY_INDEX, - /** selected diagram index, type: `int`, shape: `()` */ - UMAMI_OUT_DIAGRAM_INDEX, - /** CUDA or HIP stream for asynchronous execution. Listed as an output as it is a - * mutable pointer */ - UMAMI_OUT_GPU_STREAM, - /** value of a custom cut, 0 rejects the event, type: `double`, shape: `()` */ - UMAMI_OUT_CUT, - /** value of a custom energy scale, type: `double`, shape: `()` */ - UMAMI_OUT_SCALE, - /** value x f(x, Q) of a custom parton density, type: `double`, shape: `()` */ - UMAMI_OUT_PDF, - /** value of a custom event bias, type: `double`, shape: `()` */ - UMAMI_OUT_BIAS, - /** ratio of the matrix element to a reference matrix element defined by the - * caller, type: `double`, shape: `()` */ - UMAMI_OUT_MATRIX_ELEMENT_RATIO, -} UmamiOutputKey; - -/** Number of values in `UmamiOutputKey` */ -#define UMAMI_OUTPUT_KEY_COUNT (UMAMI_OUT_MATRIX_ELEMENT_RATIO + 1) - -/** Implementation-defined pointer to a matrix element instance */ -typedef void* UmamiHandle; - -/** - * Retrieve metadata about the implemented matrix element - * - * @param meta_key - * key specifying the type of metadata to be retrieved - * @param result - * pointer to store the result. It's type depends on the metadata key - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_get_meta(UmamiMetaKey meta_key, void* result); - -/** - * Reports which input keys can be passed to this matrix element implementation, to - * be written by whoever implements the UMAMI interface. Implementing this function - * is optional; a caller that fails to resolve it should treat it as unavailable. - * - * @param supported - * pointer set on return to an implementation-owned array of booleans, where - * element `i` indicates whether the input key with numeric value `i` - * (see `UmamiInputKey`) is accepted by this implementation. - * @param count - * pointer set on return to the number of elements in `*supported`, reflecting - * the number of `UmamiInputKey` values known to this implementation. A caller - * built against a newer header should treat any key index beyond `*count` as - * not supported. - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_supported_inputs(bool const** supported, int* count); - -/** - * Reports which input keys must be passed to this matrix element implementation for - * `umami_matrix_element` to succeed, to be written by whoever implements the UMAMI - * interface. This is a subset of the keys reported by `umami_supported_inputs`. - * Implementing this function is optional; a caller that fails to resolve it should - * treat it as unavailable. - * - * @param required - * pointer set on return to an implementation-owned array of booleans, where - * element `i` indicates whether the input key with numeric value `i` - * (see `UmamiInputKey`) is mandatory for this implementation. - * @param count - * pointer set on return to the number of elements in `*required`. A caller - * built against a newer header should treat any key index beyond `*count` as - * not required. - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_required_inputs(bool const** required, int* count); - -/** - * Reports which output keys can be requested from this matrix element - * implementation, to be written by whoever implements the UMAMI interface. - * Implementing this function is optional; a caller that fails to resolve it should - * treat it as unavailable. - * - * @param supported - * pointer set on return to an implementation-owned array of booleans, where - * element `i` indicates whether the output key with numeric value `i` - * (see `UmamiOutputKey`) can be produced by this implementation. - * @param count - * pointer set on return to the number of elements in `*supported`. A caller - * built against a newer header should treat any key index beyond `*count` as - * not supported. - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_supported_outputs(bool const** supported, int* count); - -/** - * Creates an instance of the matrix element. Each instance is independent, so thread - * safety can be achieved by creating a separate one for every thread. - * - * @param handle - * pointer to an instance of the subprocess. Has to be cleaned up by - * the caller with `free_subprocess`. - * @param param_card_path - * path to the parameter file - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_initialize(UmamiHandle* handle, char const* param_card_path); - -/** - * Sets the value of a model parameter - * - * @param handle - * handle of a matrix element instance - * @param name - * name of the parameter - * @param parameter_real - * real part of the parameter value - * @param parameter_imag - * imaginary part of the parameter value. Ignored for real valued parameters. - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_set_parameter( - UmamiHandle handle, char const* name, double parameter_real, double parameter_imag -); - -/** - * Retrieves the value of a model parameter - * - * @param handle - * handle of a matrix element instance - * @param name - * name of the parameter - * @param parameter_real - * pointer to double to return real part of the parameter value - * @param parameter_imag - * pointer to double to return imaginary part of the parameter value. Ignored - * for real-valued parameters (i.e. you may pass a null pointer) - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_get_parameter( - UmamiHandle handle, char const* name, double* parameter_real, double* parameter_imag -); - -/** - * Evaluates the matrix element as a function of the given inputs, filling the - * requested outputs. Unless otherwise specified, all inputs and outputs have a - * column-major memory layout and have a batch dimension that is contiguous in memory. - * - * @param handle - * handle of a matrix element instance - * @param count - * number of events to evaluate the matrix element for - * @param stride - * stride of the batch dimension of the input and output arrays to simplify - * parallel execution on CPUs, see memory layout - * @param offset - * offset of the event index - * @param input_count - * number of inputs to the matrix element - * @param input_keys - * pointer to an array of input keys, length `input_count` - * @param inputs - * pointer to an array of void pointers to the inputs. The type of the inputs - * depends on the input key - * @param output_count - * number of outputs to the matrix element - * @param output_keys - * pointer to an array of output keys, length `output_count` - * @param outputs - * pointer to an array of void pointers to the outputs. The type of the outputs - * depends on the output key. The caller is responsible for allocating memory for - * the outputs. - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_matrix_element( - UmamiHandle handle, - size_t count, - size_t stride, - size_t offset, - size_t input_count, - UmamiInputKey const* input_keys, - void const* const* inputs, - size_t output_count, - UmamiOutputKey const* output_keys, - void* const* outputs -); - -/** - * Frees matrix element instance - * - * @param handle - * handle of a matrix element instance - */ -UmamiStatus umami_free(UmamiHandle handle); - -#ifdef __cplusplus -} -#endif - -#endif // UMAMI_HEADER diff --git a/pyamplicol/src/pyamplicol/assets/api_templates/umami/umami.h b/pyamplicol/src/pyamplicol/assets/api_templates/umami/umami.h new file mode 120000 index 0000000000..dfa21b9ab2 --- /dev/null +++ b/pyamplicol/src/pyamplicol/assets/api_templates/umami/umami.h @@ -0,0 +1 @@ +../../../../../../madspace/include/madspace/umami.h \ No newline at end of file diff --git a/pyamplicol/tests/unit/test_repository_policy.py b/pyamplicol/tests/unit/test_repository_policy.py index 38dc6ed7d6..8273b51dd9 100644 --- a/pyamplicol/tests/unit/test_repository_policy.py +++ b/pyamplicol/tests/unit/test_repository_policy.py @@ -112,6 +112,8 @@ def test_release_source_roots_contain_no_symlinks() -> None: for source_root in FIRST_PARTY_ROOTS for path in source_root.rglob("*") if _eligible(path) and path.is_symlink() + # shared with madspace, copied by content (_SHARED_SYMLINKS) + and path != ROOT / "src/pyamplicol/assets/api_templates/umami/umami.h" ] assert not links, "release source contains symlinks:\n" + "\n".join(links) From 1c0f5c896ce7540bcb9fa496ac1b74b5157ff3c5 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 26 Oct 2026 13:41:16 +0100 Subject: [PATCH 39/88] pyamplicol: model-wide generation data once per artifact The template catalog of a prepared kernel pack is kept on the pack instead of being rebuilt at every access, the executor packs share the validated model-wide part of the direct-template catalog and kernel pack, and the structural source proofs of one artifact share the records of their common payloads. For p p > j j j [amplicol] the output takes 85 s instead of 346 s, with the same artifacts up to timings and the identities derived from them. --- .../pyamplicol/generation/artifact_writer.py | 7 +- .../generation/recurrence_schedule_sharing.py | 126 +++++++++++++----- .../generation/structural_source_proof.py | 93 ++++++++++--- pyamplicol/src/pyamplicol/models/prepared.py | 14 +- 4 files changed, 180 insertions(+), 60 deletions(-) diff --git a/pyamplicol/src/pyamplicol/generation/artifact_writer.py b/pyamplicol/src/pyamplicol/generation/artifact_writer.py index be38636aff..c81c620edc 100644 --- a/pyamplicol/src/pyamplicol/generation/artifact_writer.py +++ b/pyamplicol/src/pyamplicol/generation/artifact_writer.py @@ -118,7 +118,10 @@ from .structural_source_proof import ( ROLE as STRUCTURAL_SOURCE_PROOF_ROLE, ) -from .structural_source_proof import build_generation_structural_proof +from .structural_source_proof import ( + StructuralProofCache, + build_generation_structural_proof, +) from .validation import ValidationPointRecord, validation_point_map if TYPE_CHECKING: @@ -1151,6 +1154,7 @@ def write_schema_v3_artifact( staged_root = builder.root if staged_root is None: # pragma: no cover - builder invariant raise RuntimeError("artifact builder is not active") + proof_cache = StructuralProofCache() for process_record in process_records: process_id = str(process_record["id"]) execution_path = f"processes/{process_id}/execution.json" @@ -1174,6 +1178,7 @@ def write_schema_v3_artifact( execution=execution_payload, evaluator_container_path=evaluator_container_path, evaluator_container_index_sha256=(evaluator_container_index_sha256), + cache=proof_cache, ) builder.add_json( f"processes/{process_id}/structural-source-proof.json", diff --git a/pyamplicol/src/pyamplicol/generation/recurrence_schedule_sharing.py b/pyamplicol/src/pyamplicol/generation/recurrence_schedule_sharing.py index 5ea6aed790..abb5b74581 100644 --- a/pyamplicol/src/pyamplicol/generation/recurrence_schedule_sharing.py +++ b/pyamplicol/src/pyamplicol/generation/recurrence_schedule_sharing.py @@ -650,21 +650,49 @@ def recurrence_schedule_semantic_digest( return _canonical_digest(payload) -def build_recurrence_process_executor_pack( - *, - direct_catalog: object, - kernel_pack: object, - required_executor_ids: Sequence[int], - runtime_layout_digest: str, -) -> RecurrenceProcessExecutorPack: - """Project one validated model catalog to its exact fixed-plan executor set.""" +@dataclass(frozen=True, slots=True) +class _ExecutorCatalog: + """The validated model-wide part of every process executor pack.""" + + templates: Mapping[int, Mapping[str, object]] + kernels: Mapping[int, Mapping[str, object]] + backend: str + target_triple: str + portable: bool + cpu_features: tuple[str, ...] + compiled_model_digest: str + recurrence_template_catalog_digest: str + prepared_kernel_pack_digest: str + direct_template_catalog_digest: str + + +# by object identity; the entries keep their keys alive so ids are not reused +_EXECUTOR_CATALOGS: dict[tuple[int, int], tuple[object, object, _ExecutorCatalog]] = {} +_EXECUTOR_CATALOGS_LOCK = Lock() +_EXECUTOR_CATALOGS_LIMIT = 4 + +def _executor_catalog(direct_catalog: object, kernel_pack: object) -> _ExecutorCatalog: + key = (id(direct_catalog), id(kernel_pack)) + with _EXECUTOR_CATALOGS_LOCK: + entry = _EXECUTOR_CATALOGS.get(key) + if entry is None: + entry = (direct_catalog, kernel_pack, + _build_executor_catalog(direct_catalog, kernel_pack)) + if len(_EXECUTOR_CATALOGS) >= _EXECUTOR_CATALOGS_LIMIT: + _EXECUTOR_CATALOGS.pop(next(iter(_EXECUTOR_CATALOGS))) + _EXECUTOR_CATALOGS[key] = entry + return entry[2] + + +def _build_executor_catalog( + direct_catalog: object, kernel_pack: object +) -> _ExecutorCatalog: direct_mapping = _object_mapping(direct_catalog, "direct-template catalog") pack_mapping = _object_mapping(kernel_pack, "prepared-kernel pack") raw_templates = _object_sequence( direct_mapping.get("templates"), "direct-template catalog templates" ) - catalog_count = len(raw_templates) templates: dict[int, Mapping[str, object]] = {} for expected_id, raw_template in enumerate(raw_templates): template = _require_mapping(raw_template, "direct-template descriptor") @@ -677,19 +705,6 @@ def build_recurrence_process_executor_pack( ) templates[executor_id] = template - selected_ids = tuple( - _wire_u32(value, "required direct-executor ID") - for value in required_executor_ids - ) - if selected_ids != tuple(sorted(set(selected_ids))): - raise RecurrenceScheduleSharingError( - "required direct-executor IDs must be sorted and unique" - ) - if any(executor_id not in templates for executor_id in selected_ids): - raise RecurrenceScheduleSharingError( - "required direct-executor ID is absent from its complete catalog" - ) - raw_kernels = _object_sequence( pack_mapping.get("kernels"), "prepared-kernel pack kernels" ) @@ -731,16 +746,13 @@ def build_recurrence_process_executor_pack( raise RecurrenceScheduleSharingError( "target CPU features must be sorted and unique" ) - - descriptors = tuple( - _encode_process_executor_descriptor( - templates[executor_id], - backend=backend, - kernels=kernels, - ) - for executor_id in selected_ids - ) - return RecurrenceProcessExecutorPack( + return _ExecutorCatalog( + templates=templates, + kernels=kernels, + backend=backend, + target_triple=target_triple, + portable=portable, + cpu_features=cpu_features, compiled_model_digest=_wire_sha256( direct_mapping.get("compiled_model_digest"), "compiled-model digest" ), @@ -755,14 +767,54 @@ def build_recurrence_process_executor_pack( direct_template_catalog_digest=_wire_sha256( direct_mapping.get("catalog_digest"), "direct-template catalog digest" ), + ) + + +def build_recurrence_process_executor_pack( + *, + direct_catalog: object, + kernel_pack: object, + required_executor_ids: Sequence[int], + runtime_layout_digest: str, +) -> RecurrenceProcessExecutorPack: + """Project one validated model catalog to its exact fixed-plan executor set.""" + + catalog = _executor_catalog(direct_catalog, kernel_pack) + templates = catalog.templates + selected_ids = tuple( + _wire_u32(value, "required direct-executor ID") + for value in required_executor_ids + ) + if selected_ids != tuple(sorted(set(selected_ids))): + raise RecurrenceScheduleSharingError( + "required direct-executor IDs must be sorted and unique" + ) + if any(executor_id not in templates for executor_id in selected_ids): + raise RecurrenceScheduleSharingError( + "required direct-executor ID is absent from its complete catalog" + ) + + descriptors = tuple( + _encode_process_executor_descriptor( + templates[executor_id], + backend=catalog.backend, + kernels=catalog.kernels, + ) + for executor_id in selected_ids + ) + return RecurrenceProcessExecutorPack( + compiled_model_digest=catalog.compiled_model_digest, + recurrence_template_catalog_digest=catalog.recurrence_template_catalog_digest, + prepared_kernel_pack_digest=catalog.prepared_kernel_pack_digest, + direct_template_catalog_digest=catalog.direct_template_catalog_digest, runtime_layout_digest=_wire_sha256( runtime_layout_digest, "runtime-layout digest" ), - backend=backend, - target_triple=target_triple, - portable=portable, - cpu_features=cpu_features, - catalog_executor_count=catalog_count, + backend=catalog.backend, + target_triple=catalog.target_triple, + portable=catalog.portable, + cpu_features=catalog.cpu_features, + catalog_executor_count=len(templates), executor_ids=selected_ids, descriptor_payloads=descriptors, ) diff --git a/pyamplicol/src/pyamplicol/generation/structural_source_proof.py b/pyamplicol/src/pyamplicol/generation/structural_source_proof.py index ee865f980a..f5e97d2dc7 100644 --- a/pyamplicol/src/pyamplicol/generation/structural_source_proof.py +++ b/pyamplicol/src/pyamplicol/generation/structural_source_proof.py @@ -11,6 +11,7 @@ import hashlib import json from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -30,6 +31,24 @@ _DIAGNOSTIC_PATH_FIELDS = frozenset({"compiler_path"}) +@dataclass +class StructuralProofCache: + """Payload records shared by the proofs of one staged artifact. + + Only valid while the referenced files do not change; existing paths are + cached, missing ones are checked again. + """ + + # loose path -> (object record, nested payload paths) + loose: dict[str, tuple[dict[str, object], tuple[str, ...]]] = field( + default_factory=dict + ) + # (container path, index sha256) -> (member records, container record) + containers: dict[ + tuple[str, str], tuple[dict[str, dict[str, object]], dict[str, object]] + ] = field(default_factory=dict) + + def canonical_sha256(domain: str, value: object) -> str: """Hash canonical JSON under a distinct semantic domain.""" @@ -219,6 +238,29 @@ def add(record: object, path: str) -> None: return tuple(lanes) +def _loose_object( + root: Path, + loose: str, +) -> tuple[dict[str, object], tuple[str, ...]]: + """The inventory record of a loose payload and the paths it references.""" + + path = root / loose + record: dict[str, object] = { + "path": loose, + "sha256": sha256_file(path), + "size_bytes": path.stat().st_size, + } + if path.suffix != ".json": + return record, () + try: + nested = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ValueError( + f"referenced lane manifest is not valid JSON: {loose}" + ) from error + return record, _payload_paths(nested) + + def _resolve_loose_payload( root: Path, process_id: str, @@ -243,13 +285,25 @@ def build_generation_structural_proof( execution: Mapping[str, object], evaluator_container_path: str | None, evaluator_container_index_sha256: str | None, + cache: StructuralProofCache | None = None, ) -> dict[str, object]: - """Build exact source/lane evidence from the private artifact staging tree.""" + """Build exact source/lane evidence from the private artifact staging tree. + + ``cache`` shares the records of payloads common to the processes of one + staged artifact. + """ lanes = execution_lanes(execution) member_records: dict[str, dict[str, object]] = {} container_record: dict[str, object] | None = None - if evaluator_container_path is not None: + container_key = ( + None + if evaluator_container_path is None or evaluator_container_index_sha256 is None + else (evaluator_container_path, evaluator_container_index_sha256) + ) + if cache is not None and container_key in cache.containers: + member_records, container_record = cache.containers[container_key] + elif evaluator_container_path is not None: container = artifact_root / evaluator_container_path if not container.is_file() or container.is_symlink(): raise ValueError("evaluator container is absent during structural proof") @@ -275,6 +329,8 @@ def build_generation_structural_proof( "index_sha256": reader.index.index_sha256, "member_count": len(reader.members), } + if cache is not None and container_key is not None: + cache.containers[container_key] = (member_records, container_record) loose_objects: dict[str, dict[str, object]] = { execution_path: { @@ -299,7 +355,10 @@ def include( lane_members.add(relative) used_members.add(relative) return - loose = _resolve_loose_payload(artifact_root, process_id, relative) + if cache is not None and relative in cache.loose: + loose: str | None = relative + else: + loose = _resolve_loose_payload(artifact_root, process_id, relative) if loose is None: # Some manifest paths are logical roots interpreted by the # runtime rather than physical files. Preserve them in the @@ -308,24 +367,15 @@ def include( if loose in lane_loose: return lane_loose.add(loose) - path = artifact_root / loose - loose_objects.setdefault( - loose, - { - "path": loose, - "sha256": sha256_file(path), - "size_bytes": path.stat().st_size, - }, - ) - if path.suffix == ".json": - try: - nested = json.loads(path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError) as error: - raise ValueError( - f"referenced lane manifest is not valid JSON: {loose}" - ) from error - for nested_relative in _payload_paths(nested): - include(nested_relative, lane_loose, lane_members) + cached = None if cache is None else cache.loose.get(loose) + if cached is None: + cached = _loose_object(artifact_root, loose) + if cache is not None: + cache.loose[loose] = cached + record, nested_paths = cached + loose_objects.setdefault(loose, dict(record)) + for nested_relative in nested_paths: + include(nested_relative, lane_loose, lane_members) for relative in _payload_paths(lane): include(relative, referenced_loose, referenced_members) @@ -641,6 +691,7 @@ def validate_generation_structural_proof( "ROLE", "SCHEMA", "SEMANTIC_MAP_DOMAINS", + "StructuralProofCache", "build_generation_structural_proof", "canonical_sha256", "execution_lanes", diff --git a/pyamplicol/src/pyamplicol/models/prepared.py b/pyamplicol/src/pyamplicol/models/prepared.py index 5e17fd9790..9fc0a9b52f 100644 --- a/pyamplicol/src/pyamplicol/models/prepared.py +++ b/pyamplicol/src/pyamplicol/models/prepared.py @@ -1151,6 +1151,12 @@ class PreparedKernelPack: default=None, ) ) + _recurrence_template_catalog: RecurrenceTemplateCatalog | None = field( + init=False, + repr=False, + compare=False, + default=None, + ) def __post_init__(self) -> None: if self.backend not in _BACKENDS: @@ -1412,12 +1418,18 @@ def recurrence_template_catalog(self) -> RecurrenceTemplateCatalog | None: recurrence_template = self.recurrence_template if recurrence_template is None: return None + # The catalog is immutable and model-wide, but rebuilding it is costly + cached = self._recurrence_template_catalog + if cached is not None: + return cached from .recurrence_template import RecurrenceTemplateCatalog thawed = _thaw_json(recurrence_template) - return RecurrenceTemplateCatalog.from_dict( + catalog = RecurrenceTemplateCatalog.from_dict( _mapping(thawed, "kernel_pack.recurrence_template") ) + object.__setattr__(self, "_recurrence_template_catalog", catalog) + return catalog @property def recurrence_direct_template_catalog( From 19c2a26a3d02772e8b660c5d9eea7988fe043ffc Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 26 Oct 2026 13:59:17 +0100 Subject: [PATCH 40/88] madspace: pt^2, not m_T^2, bounds the momentum transfer to the beam In the 2->3 step off the opposite beam (a side of two particles in ColorOrderedMapping), any cut bounded |(p2 - k)^2| from below by etmin_2^2, the m_T floor of p2. For a massless beam k along z, |(p2 - k)^2| >= m_T^2 - m^2 = pt^2, which is the same only for a massless p2. A top with no cut of its own has the floor m_T >= m_t, and the momentum transfers below m_t^2 were never sampled: about 6% of the volume of g g > t t~ g at sqrt(s) = 600 GeV and 20% of it at small top pt. The AmpliCol channels with the top pair on one side came out 25% low, and test_mg7_amplicol::test_gg_ttxg failed. --- .../madspace/phasespace/three_particle.hpp | 5 +- madspace/src/phasespace/three_particle.cpp | 15 +++- .../tests/test_color_ordered_massive_side.py | 74 +++++++++++++++++++ 3 files changed, 88 insertions(+), 6 deletions(-) create mode 100644 madspace/tests/test_color_ordered_massive_side.py diff --git a/madspace/include/madspace/phasespace/three_particle.hpp b/madspace/include/madspace/phasespace/three_particle.hpp index 2ba277ccd3..f0a3cc3cdb 100644 --- a/madspace/include/madspace/phasespace/three_particle.hpp +++ b/madspace/include/madspace/phasespace/three_particle.hpp @@ -179,7 +179,7 @@ class TwoToThreeParticleScattering : public Mapping { * the momentum transfer @f$(p_2 - k)^2 \le 0@f$ * instead of @f$(p_2 + p_3)^2@f$; its absolute * value is sampled. With @p has_cut it is bounded - * by @f$|(p_2 - k)^2| \ge \texttt{etmin\_2}^2@f$, + * by @f$|(p_2 - k)^2| \ge \texttt{etmin\_2}^2 - m_2^2@f$, * the `drcut` and `s23_min_cut` conditions are * not used. Requires @p p12_condition. AmpliCol * builds a side of two particles this way, with @@ -225,11 +225,12 @@ class TwoToThreeParticleScattering : public Mapping { split_conditions(FunctionBuilder& fb, const NamedVector& conditions) const; // Sampled and kinematic range of the second invariant: s23 itself, or // |(p2 - k)^2| for an incoming recoil. Arguments: the kinematic range of - // s23 and the conditions. + // s23, the squared mass of p2 and the conditions. std::array sampling_ranges( FunctionBuilder& fb, Value s23_min, Value s23_max, + Value m2_squared, const NamedVector& conditions ) const; Invariant _t_invariant; diff --git a/madspace/src/phasespace/three_particle.cpp b/madspace/src/phasespace/three_particle.cpp index 98b0f4ac9c..a3871414bc 100644 --- a/madspace/src/phasespace/three_particle.cpp +++ b/madspace/src/phasespace/three_particle.cpp @@ -178,6 +178,7 @@ std::array TwoToThreeParticleScattering::sampling_ranges( FunctionBuilder& fb, Value s23_min, Value s23_max, + Value m2_squared, const NamedVector& conditions ) const { // returns the sampled and the kinematic range of the sampled variable @@ -188,8 +189,11 @@ std::array TwoToThreeParticleScattering::sampling_ranges( auto v_phys_min = negated(fb, s23_max), v_phys_max = negated(fb, s23_min); auto v_min = v_phys_min, v_max = v_phys_max; if (_has_cut) { - // |(p2 - k)^2| >= m_T(p2)^2 for a massless k along the beam axis - v_min = fb.max(v_min, fb.square(conditions.at(4))); + // |(p2 - k)^2| >= m_T(p2)^2 - m2^2 = p_T(p2)^2 for a massless k along + // the beam axis: the ETmin floor of a massive particle bounds m_T + v_min = fb.max( + v_min, fb.max(fb.sub(fb.square(conditions.at(4)), m2_squared), Value(0.)) + ); // an empty range keeps a small width; the weight is zero there, as // the point then lies outside the kinematic range auto width = fb.max(fb.mul(v_min, Value(1e-10)), Value(1e-300)); @@ -217,7 +221,7 @@ Mapping::Result TwoToThreeParticleScattering::build_forward_impl( std::array ranges; if (_incoming_recoil) { auto [lo, hi] = fb.s23_min_max(p_a, p_12, p_3, t1, m1, m2); - ranges = sampling_ranges(fb, lo, hi, conditions); + ranges = sampling_ranges(fb, lo, hi, fb.square(m2), conditions); } else if (_has_cut) { auto [s_min, s_max] = fb.s23_min_max_cut( p_a, @@ -303,7 +307,10 @@ Mapping::Result TwoToThreeParticleScattering::build_inverse_impl( if (_incoming_recoil) { auto [s23, lo, hi] = fb.s23_value_and_min_max(p_a, p_3, t1_abs, p1, p2); v = negated(fb, s23); - ranges = sampling_ranges(fb, lo, hi, conditions); + nested_vector2 factors_p2{{1.0}}; + auto m2_squared = + fb.unstack(fb.invariants_from_momenta(fb.stack({p2}), factors_p2)).at(0); + ranges = sampling_ranges(fb, lo, hi, m2_squared, conditions); } else if (_has_cut) { auto [s23, s_min, s_max] = fb.s23_value_and_min_max_cut( p_a, diff --git a/madspace/tests/test_color_ordered_massive_side.py b/madspace/tests/test_color_ordered_massive_side.py new file mode 100644 index 0000000000..19b42d7346 --- /dev/null +++ b/madspace/tests/test_color_ordered_massive_side.py @@ -0,0 +1,74 @@ +"""A side of two massive particles in ColorOrderedMapping with a cut elsewhere. + +g g > t t~ g in the color order t g g g t~ (0 2 4 1 3 in particle numbers): the +gluon is alone on one side, the top pair is the other, built by the 2->3 step +off the opposite beam. Any cut turns on the transverse-energy floors of that +step, which are m_T >= m for a top without a pt cut. The floor on the momentum +transfer to the beam was etmin^2 = m_T^2, right only for massless particles: +the bound is m_T^2 - m^2 = pt^2. It removed the tops with |t| < m_t^2, about 6% +of the volume and 20% of it at small top pt, and the AmpliCol channels of +g g > t t~ g came out 25% low. +""" + +import math + +import pytest +import torch +from madspace_testing import generator, rand, randint + +import madspace as ms + +N = 400_000 +CM_ENERGY = 600.0 +ORDER = [2, 0, 4, 1, 3] +MASSES = [173.0, 173.0, 0.0] +PT_MIN = 20.0 + + +def sample(pt_min, seed, masses=MASSES): + mapping = ms.ColorOrderedMapping( + ms.ColorOrderedMapping.Options(ORDER), pt_min=pt_min + ) + g = generator(seed) + r = rand((N, mapping.random_dim()), generator=g) + d = randint(0, 2, (N, mapping.discrete_dim()), generator=g, dtype=torch.int32) + inputs = [r[:, i].contiguous() for i in range(r.shape[1])] + inputs += [d[:, j].contiguous() for j in range(d.shape[1])] + conditions = [torch.full((N,), CM_ENERGY)] + [torch.full((N,), m) for m in masses] + *momenta, det = mapping.map_forward(inputs, conditions) + det = torch.nan_to_num(det) * 2.0 ** mapping.discrete_dim() + return mapping, inputs, conditions, momenta, r, det + + +def pt(p): + return torch.sqrt(p[:, 1] ** 2 + p[:, 2] ** 2) + + +@pytest.mark.parametrize("top_pt_max", [math.inf, 50.0], ids=["all", "top pt < 50"]) +def test_cut_handed_to_mapping_keeps_volume(top_pt_max): + """The volume passing the gluon cut is the same whether the mapping knows + the cut or it is applied afterwards.""" + _, _, _, p_free, _, det_free = sample([0.0, 0.0, 0.0], 1) + _, _, _, p_cut, _, det_cut = sample([0.0, 0.0, PT_MIN], 2) + # beams first: the top is momenta[2], the gluon momenta[4] + external = torch.where( + (pt(p_free[4]) >= PT_MIN) & (pt(p_free[2]) < top_pt_max), det_free, 0.0 + ) + piped = torch.where( + (pt(p_cut[4]) >= PT_MIN) & (pt(p_cut[2]) < top_pt_max), det_cut, 0.0 + ) + error = math.sqrt((external.var(correction=0) + piped.var(correction=0)) / N) + assert error > 0.0 + assert abs(piped.mean() - external.mean()).item() < 5.0 * error + + +def test_round_trip_massive_side_with_cut(): + mapping, _, conditions, momenta, r, det = sample([0.0, 0.0, PT_MIN], 3) + out = mapping.map_inverse(momenta, conditions) + r_inv = torch.stack(out[: mapping.random_dim()], dim=1) + ok = det > 0 + rt = torch.abs(det / 2.0 ** mapping.discrete_dim() * out[-1] - 1.0)[ok] + assert torch.mean((rt > 1e-2).double()) < 1e-4 + assert torch.quantile(rt, 0.999) < 1e-6 + r_err = torch.max(torch.abs(r_inv - r), dim=1).values[ok] + assert torch.mean((r_err > 1e-6).double()) < 1e-4 From bed7c498d31b4b9aef6c1633f6429413bea75859 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 26 Oct 2026 14:06:17 +0100 Subject: [PATCH 41/88] pyamplicol: stage artifact payloads without reading them back add_stream hashes and counts a payload while writing it instead of opening it again to fsync and hash it, and a structural source proof resolves each referenced path once. For p p > j j j [amplicol] the worker takes 65 s instead of 77 s, with the same artifacts. --- pyamplicol/src/pyamplicol/artifacts/writer.py | 37 +++++++++++++++++-- .../generation/recurrence_schedule_sharing.py | 7 +++- .../generation/structural_source_proof.py | 12 ++++-- 3 files changed, 48 insertions(+), 8 deletions(-) diff --git a/pyamplicol/src/pyamplicol/artifacts/writer.py b/pyamplicol/src/pyamplicol/artifacts/writer.py index e7dc464711..65e2047768 100644 --- a/pyamplicol/src/pyamplicol/artifacts/writer.py +++ b/pyamplicol/src/pyamplicol/artifacts/writer.py @@ -3,6 +3,7 @@ from __future__ import annotations +import hashlib import io import json import os @@ -114,13 +115,36 @@ def register_staged_file( if executable is not None: path.chmod(0o755 if executable else 0o644) fsync_file(path) - record = PayloadRecord( - path=path_value, + return self._record_payload( + path_value, role=role, media_type=media_type, size_bytes=path.stat().st_size, sha256=sha256_file(path), executable=executable_bit(path), + target=target, + process_id=process_id, + ) + + def _record_payload( + self, + path_value: str, + *, + role: str, + media_type: str, + size_bytes: int, + sha256: str, + executable: bool, + target: Mapping[str, object] | None, + process_id: str | None, + ) -> PayloadRecord: + record = PayloadRecord( + path=path_value, + role=role, + media_type=media_type, + size_bytes=size_bytes, + sha256=sha256, + executable=executable, target=dict(target) if target is not None else None, process_id=process_id, ) @@ -178,6 +202,9 @@ def add_stream( ) temporary = Path(temporary_name) descriptor_open = True + # hashed while written: the staged file is not read back + digest = hashlib.sha256() + size_bytes = 0 try: with os.fdopen(descriptor, "wb") as stream: descriptor_open = False @@ -192,6 +219,8 @@ def add_stream( "artifact stream source returned more bytes than requested" ) view = memoryview(chunk) + digest.update(view) + size_bytes += len(view) written = 0 while written < len(view): count = stream.write(view[written:]) @@ -207,10 +236,12 @@ def add_stream( os.close(descriptor) with suppress(FileNotFoundError): temporary.unlink() - return self.register_staged_file( + return self._record_payload( path_value, role=role, media_type=media_type, + size_bytes=size_bytes, + sha256=digest.hexdigest(), executable=executable, target=target, process_id=process_id, diff --git a/pyamplicol/src/pyamplicol/generation/recurrence_schedule_sharing.py b/pyamplicol/src/pyamplicol/generation/recurrence_schedule_sharing.py index abb5b74581..dd82dfe8f1 100644 --- a/pyamplicol/src/pyamplicol/generation/recurrence_schedule_sharing.py +++ b/pyamplicol/src/pyamplicol/generation/recurrence_schedule_sharing.py @@ -677,8 +677,11 @@ def _executor_catalog(direct_catalog: object, kernel_pack: object) -> _ExecutorC with _EXECUTOR_CATALOGS_LOCK: entry = _EXECUTOR_CATALOGS.get(key) if entry is None: - entry = (direct_catalog, kernel_pack, - _build_executor_catalog(direct_catalog, kernel_pack)) + entry = ( + direct_catalog, + kernel_pack, + _build_executor_catalog(direct_catalog, kernel_pack), + ) if len(_EXECUTOR_CATALOGS) >= _EXECUTOR_CATALOGS_LIMIT: _EXECUTOR_CATALOGS.pop(next(iter(_EXECUTOR_CATALOGS))) _EXECUTOR_CATALOGS[key] = entry diff --git a/pyamplicol/src/pyamplicol/generation/structural_source_proof.py b/pyamplicol/src/pyamplicol/generation/structural_source_proof.py index f5e97d2dc7..a27409830a 100644 --- a/pyamplicol/src/pyamplicol/generation/structural_source_proof.py +++ b/pyamplicol/src/pyamplicol/generation/structural_source_proof.py @@ -340,6 +340,8 @@ def build_generation_structural_proof( } } used_members: set[str] = set() + # the staged files do not change while this proof is built + resolved: dict[str, str | None] = {} lane_records: list[dict[str, object]] = [] roles: list[dict[str, str]] = [] for lane_id, lane in lanes: @@ -355,10 +357,14 @@ def include( lane_members.add(relative) used_members.add(relative) return - if cache is not None and relative in cache.loose: - loose: str | None = relative + if relative in resolved: + loose = resolved[relative] + elif cache is not None and relative in cache.loose: + loose = resolved[relative] = relative else: - loose = _resolve_loose_payload(artifact_root, process_id, relative) + loose = resolved[relative] = _resolve_loose_payload( + artifact_root, process_id, relative + ) if loose is None: # Some manifest paths are logical roots interpreted by the # runtime rather than physical files. Preserve them in the From ce5db24be21bd615b615ed14b33aaf5d86fb8936 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 26 Oct 2026 14:47:37 +0100 Subject: [PATCH 42/88] pyamplicol: strict flavour-blind generation process.flavor_blind_strict (--flavor-blind-strict) implies flavor_blind but routes a member only through a representative whose relation passes the vertex-symmetry check, trying every flavour pairing, conjugated ones last. A member no representative of its class accepts is generated and becomes a representative itself, so the written process-groups.json has no warnings. --- .../docs/development/CONFIG_CONTRACT.md | 6 +- .../docs/user/command-line-interface.md | 4 +- pyamplicol/docs/user/configuration.md | 4 +- pyamplicol/docs/user/process-grouping.md | 14 +++ pyamplicol/examples/all_options.toml | 4 + pyamplicol/src/pyamplicol/cli/parser.py | 10 +++ pyamplicol/src/pyamplicol/config/models.py | 2 + .../src/pyamplicol/generation/service.py | 90 +++++++++++++++---- pyamplicol/tests/unit/test_config_models.py | 2 +- .../tests/unit/test_flavor_blind_expansion.py | 28 ++++++ 10 files changed, 144 insertions(+), 20 deletions(-) diff --git a/pyamplicol/docs/development/CONFIG_CONTRACT.md b/pyamplicol/docs/development/CONFIG_CONTRACT.md index c568afb0ee..58db22bbd3 100644 --- a/pyamplicol/docs/development/CONFIG_CONTRACT.md +++ b/pyamplicol/docs/development/CONFIG_CONTRACT.md @@ -39,6 +39,7 @@ overrides in order, then recorded license/resource clamping. - `selected_color_sector_ids: list[int] = []` - `selected_source_helicities: table[int, int] = {}` - `flavor_blind: bool = false` +- `flavor_blind_strict: bool = false` - `consider_charge_conjugation: bool = true` - `only_unit_multiplier: bool = false` - `analyze_process_groups: none | numerical = none` @@ -52,9 +53,10 @@ The explicit color-sector, reference-order, and source-helicity controls are developer-facing generation constraints. Ordinary users should leave them at their defaults and use the runtime selectors below. -The four process-group fields decide what `process-groups.json` records: +The five process-group fields decide what `process-groups.json` records: `flavor_blind` generates one representative per flavour-relabelling class and -records the others as unverified members; `analyze_process_groups = "numerical"` +records the others as unverified members; `flavor_blind_strict` implies it but +routes only members whose relabelling is a symmetry of the allowed vertices; `analyze_process_groups = "numerical"` merges generated processes whose matrix elements agree numerically up to a rational factor, restricted to factor 1 by `only_unit_multiplier`. `consider_charge_conjugation` also admits charge-conjugate relations in both diff --git a/pyamplicol/docs/user/command-line-interface.md b/pyamplicol/docs/user/command-line-interface.md index c536673d71..63608e8379 100644 --- a/pyamplicol/docs/user/command-line-interface.md +++ b/pyamplicol/docs/user/command-line-interface.md @@ -109,7 +109,9 @@ Useful generation options include: - `--name NAME` to name a request; - `--multiparticle 'p=d,d~,g'` to define a multiparticle label; - `--flavor-scheme N` and `--max-quark-lines N` to constrain expansion; -- `--flavor-blind` to generate one subprocess per flavour-relabelling class +- `--flavor-blind` to generate one subprocess per flavour-relabelling class, + or `--flavor-blind-strict` to route only the relations that are symmetries + of the allowed vertices (see [Group equivalent subprocesses](#group-equivalent-subprocesses)); - `--analyze-process-groups {none,numerical}`, `--only-unit-multiplier`, and `--[no-]consider-charge-conjugation` to search the generated subprocesses diff --git a/pyamplicol/docs/user/configuration.md b/pyamplicol/docs/user/configuration.md index 786c117eef..192e1a714a 100644 --- a/pyamplicol/docs/user/configuration.md +++ b/pyamplicol/docs/user/configuration.md @@ -166,12 +166,13 @@ assign arbitrary names fixed QCD or electroweak meanings. ### Process groups -Four `process` fields control which subprocesses are recorded as sharing one +Five `process` fields control which subprocesses are recorded as sharing one matrix element in the artifact's `process-groups.json`: ```toml [process] flavor_blind = false # generate one subprocess per flavour class +flavor_blind_strict = false # same, routing only vertex-symmetric relations consider_charge_conjugation = true # also relate charge-conjugate subprocesses only_unit_multiplier = false # accept numerical relations only with factor 1 analyze_process_groups = "none" # or "numerical", run after generation @@ -179,6 +180,7 @@ analyze_process_groups = "none" # or "numerical", run after generation `flavor_blind` is an assumption, not a check: members whose flavour relabelling is not a symmetry of the model couplings are reported with a warning. +`flavor_blind_strict` generates those members instead of routing them. `consider_charge_conjugation` matters only with `flavor_blind` or numerical analysis. See [Group equivalent subprocesses](command-line-interface.md#group-equivalent-subprocesses) diff --git a/pyamplicol/docs/user/process-grouping.md b/pyamplicol/docs/user/process-grouping.md index 89917daedd..7e5c2c189a 100644 --- a/pyamplicol/docs/user/process-grouping.md +++ b/pyamplicol/docs/user/process-grouping.md @@ -272,6 +272,20 @@ one fixed way, so it can warn about a relation that holds. In the same run, `d <-> s` alone moves the allowed `W` vertex `[d w+ u]`, although also exchanging `u <-> c` would restore it. +### Strict flavour-blind generation + +`--flavor-blind-strict` (`process.flavor_blind_strict = true`, which implies +`flavor_blind`) routes only the relations that pass the check above. Within a +flavour class, each member tries every flavour pairing, conjugated ones last, +against each representative kept so far, and joins the first one without a +warning. A member that matches none is generated and becomes a representative +of its own, which later members may join. No member of the written map has a +warning. For `p p > j j` with explicit electroweak orders, `u u~ > g g` is then +generated next to `d d~ > g g`, while `s c > s c` still joins `d u > d u`. + +The check is local and conservative, so strict generation may generate a +member whose relation holds; it never routes a member the check rejects. + ### When to use it Equal matrix elements within a class are **your responsibility**. The flag is diff --git a/pyamplicol/examples/all_options.toml b/pyamplicol/examples/all_options.toml index c0eb70f3a1..b4d4313648 100644 --- a/pyamplicol/examples/all_options.toml +++ b/pyamplicol/examples/all_options.toml @@ -51,6 +51,10 @@ selected_source_helicities = { "1" = -1, "2" = 1 } # class and record the others as unverified members in process-groups.json. # Default false. flavor_blind = false +# process.flavor_blind_strict: Like flavor_blind, but route only the relations +# that are symmetries of the allowed model vertices and generate the members +# that would draw a warning. Implies flavor_blind. Default false. +flavor_blind_strict = false # process.consider_charge_conjugation: Also relate processes by charge # conjugation for flavor_blind and numerical grouping. Default true. consider_charge_conjugation = true diff --git a/pyamplicol/src/pyamplicol/cli/parser.py b/pyamplicol/src/pyamplicol/cli/parser.py index f6699b0868..3cfb7ad0cb 100644 --- a/pyamplicol/src/pyamplicol/cli/parser.py +++ b/pyamplicol/src/pyamplicol/cli/parser.py @@ -276,6 +276,16 @@ def _add_process_grouping_options( "record the others as its group members (assumed, not verified)" ), ) + parser.add_argument( + "--flavor-blind-strict", + dest="process.flavor_blind_strict", + action="store_true", + default=argparse.SUPPRESS, + help=( + "like --flavor-blind, but route only the relations that are " + "symmetries of the allowed model vertices and generate the others" + ), + ) parser.add_argument( "--consider-charge-conjugation", dest="process.consider_charge_conjugation", diff --git a/pyamplicol/src/pyamplicol/config/models.py b/pyamplicol/src/pyamplicol/config/models.py index 40bc9719da..8728f002fd 100644 --- a/pyamplicol/src/pyamplicol/config/models.py +++ b/pyamplicol/src/pyamplicol/config/models.py @@ -345,6 +345,7 @@ class ProcessConfig: metadata=_setting("map_int", dynamic_kind="int"), ) flavor_blind: bool = field(default=False, metadata=_setting("bool")) + flavor_blind_strict: bool = field(default=False, metadata=_setting("bool")) consider_charge_conjugation: bool = field(default=True, metadata=_setting("bool")) only_unit_multiplier: bool = field(default=False, metadata=_setting("bool")) analyze_process_groups: ProcessGroupAnalysis = field( @@ -431,6 +432,7 @@ def __post_init__(self) -> None: ) for name in ( "flavor_blind", + "flavor_blind_strict", "consider_charge_conjugation", "only_unit_multiplier", ): diff --git a/pyamplicol/src/pyamplicol/generation/service.py b/pyamplicol/src/pyamplicol/generation/service.py index 55dc5f2729..3a0df17f71 100644 --- a/pyamplicol/src/pyamplicol/generation/service.py +++ b/pyamplicol/src/pyamplicol/generation/service.py @@ -162,6 +162,7 @@ flavor_blind_classes, flavor_blind_mapping, particle_table, + profile_mappings, ) from .progress import GenerationPhaseReporter, PhaseHandle from .recurrence_color import ( @@ -4488,6 +4489,7 @@ def _expand_process_set( table, consider_charge_conjugation=self._consider_charge_conjugation, relation_warning=self._flavor_blind_relation_warning(resolved, table), + strict=self._flavor_blind_strict, ) return tuple(result) @@ -7561,7 +7563,14 @@ def _coupling_order_limits(self) -> dict[str, int]: @property def _flavor_blind(self) -> bool: run = self._run_config - return run is not None and run.process.flavor_blind + return run is not None and ( + run.process.flavor_blind or run.process.flavor_blind_strict + ) + + @property + def _flavor_blind_strict(self) -> bool: + run = self._run_config + return run is not None and run.process.flavor_blind_strict @property def _consider_charge_conjugation(self) -> bool: @@ -8655,42 +8664,86 @@ def _grouping_process(entry: _ExpandedProcess) -> GroupingProcess: ) +def _flavor_blind_relation( + process: GroupingProcess, + representative: GroupingProcess, + table: Mapping[int, GroupingParticle], + *, + consider_charge_conjugation: bool, + relation_warning: _RelationWarning, + representative_ir: CanonicalProcessIR, + strict: bool, +) -> tuple[LegMapping, str | None]: + """The assumed mapping and its warning; ``strict`` tries every flavour + pairing, conjugated ones last, for one that draws no warning.""" + + assumed = flavor_blind_mapping( + process, + representative, + table, + consider_charge_conjugation=consider_charge_conjugation, + ) + warning = relation_warning(process, representative, assumed, representative_ir) + if warning is None or not strict: + return assumed, warning + candidates = profile_mappings(process, representative, table) + if consider_charge_conjugation: + candidates += profile_mappings(process, representative, table, conjugated=True) + for candidate in candidates: + mapping = replace(candidate, method="flavor-blind") + if ( + mapping != assumed + and relation_warning(process, representative, mapping, representative_ir) + is None + ): + return mapping, None + return assumed, warning + + def _collapse_flavor_blind_classes( entries: Sequence[_ExpandedProcess], table: Mapping[int, GroupingParticle], *, consider_charge_conjugation: bool, relation_warning: _RelationWarning, + strict: bool = False, ) -> list[_ExpandedProcess]: """Keep the representative of each flavour class and record the others on it. - Processes with explicit aliases keep their own artifact entry. + Processes with explicit aliases keep their own artifact entry. With + ``strict``, a member joins the first representative of its class whose + relation draws no warning, and becomes a generated representative otherwise. """ by_name = {entry.request.name: entry for entry in entries} members: dict[str, list[GroupMember]] = {} warnings: list[str] = [] + split = 0 for flavour_class in flavor_blind_classes( [_grouping_process(entry) for entry in entries], table, consider_charge_conjugation=consider_charge_conjugation, ): - representative = flavour_class[0] + representatives = [flavour_class[0]] for process in flavour_class[1:]: if by_name[process.process_id].aliases: continue - mapping = flavor_blind_mapping( - process, - representative, - table, - consider_charge_conjugation=consider_charge_conjugation, - ) - warning = relation_warning( - process, - representative, - mapping, - by_name[representative.process_id].process_ir, - ) + for representative in representatives: + mapping, warning = _flavor_blind_relation( + process, + representative, + table, + consider_charge_conjugation=consider_charge_conjugation, + relation_warning=relation_warning, + representative_ir=by_name[representative.process_id].process_ir, + strict=strict, + ) + if warning is None or not strict: + break + else: + representatives.append(process) + split += 1 + continue if warning is not None: warnings.append(warning) members.setdefault(representative.process_id, []).append( @@ -8727,6 +8780,13 @@ def _collapse_flavor_blind_classes( len(warnings), "" if len(warnings) == 1 else "s", ) + if strict and split: + _LOGGER.info( + "Strict flavor-blind grouping: %d subprocess%s generated because " + "their relabelling is not a symmetry of the allowed vertices", + split, + "" if split == 1 else "es", + ) return [ replace(entry, group_members=tuple(members.get(entry.request.name, ()))) for entry in entries diff --git a/pyamplicol/tests/unit/test_config_models.py b/pyamplicol/tests/unit/test_config_models.py index dd19a57651..6be37cefc1 100644 --- a/pyamplicol/tests/unit/test_config_models.py +++ b/pyamplicol/tests/unit/test_config_models.py @@ -40,7 +40,7 @@ def test_schema_v1_registry_contains_every_contract_leaf() -> None: - assert len(FIELD_REGISTRY) == 83 + assert len(FIELD_REGISTRY) == 84 assert "generation.umami_grouping" not in FIELD_REGISTRY assert "evaluator.jit.direct_translation" not in FIELD_REGISTRY assert FIELD_REGISTRY["action"].required diff --git a/pyamplicol/tests/unit/test_flavor_blind_expansion.py b/pyamplicol/tests/unit/test_flavor_blind_expansion.py index 3fa20c4429..bbdf43c122 100644 --- a/pyamplicol/tests/unit/test_flavor_blind_expansion.py +++ b/pyamplicol/tests/unit/test_flavor_blind_expansion.py @@ -251,6 +251,34 @@ def test_flavor_blind_warns_when_couplings_distinguish_the_flavours( assert dijet["s c > s c"][1].warning is None +def test_strict_flavor_blind_generates_the_members_it_cannot_verify( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(licensing, "detect_symbolica_license", _restricted_license) + + assumed = _expanded("p p > j j", flavor_scheme=4, coupling_order_policy="explicit") + strict = _expanded( + "p p > j j", + flavor_scheme=4, + coupling_order_policy="explicit", + flavor_blind_strict=True, + ) + + members = _members(strict) + assert all(member.warning is None for _, member in members.values()) + generated = {entry.request.expression for entry in strict} + assert "u u~ > g g" in generated + assert "u u~ > g g" not in {entry.request.expression for entry in assumed} + assert members["s c > s c"][0] == "d u > d u" + # Every subprocess is still generated or routed exactly once. + assert len(generated) + len(members) == len(assumed) + len(_members(assumed)) + # Pure QCD needs no split. + qcd = _expanded("p p > j j", flavor_scheme=4, flavor_blind_strict=True) + assert [e.request.name for e in qcd] == [ + e.request.name for e in _expanded("p p > j j", flavor_scheme=4) + ] + + def test_flavor_blind_keeps_processes_with_explicit_aliases( monkeypatch: pytest.MonkeyPatch, ) -> None: From 587e4f9933f12d41799586c84133ca0076f38d6b Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 26 Oct 2026 14:47:57 +0100 Subject: [PATCH 43/88] madspace: LHE flavour options drawn in proportion to their weights SubprocArgs.pdg_id_weights gives each PDG-id option of a flavour a relative weight. Flavours with unequal weights store cumulative probabilities, saved with the completer; equal or absent weights keep the uniform draw and its random numbers. Completers saved before load as uniform. --- .../include/madspace/driver/lhe_output.hpp | 7 +++ madspace/src/driver/lhe_output.cpp | 63 +++++++++++++++++-- madspace/src/python/madspace.cpp | 11 +++- madspace/tests/test_lhe.py | 58 +++++++++++++++++ 4 files changed, 133 insertions(+), 6 deletions(-) diff --git a/madspace/include/madspace/driver/lhe_output.hpp b/madspace/include/madspace/driver/lhe_output.hpp index d580554edb..c13309ae33 100644 --- a/madspace/include/madspace/driver/lhe_output.hpp +++ b/madspace/include/madspace/driver/lhe_output.hpp @@ -183,6 +183,10 @@ class LHECompleter { /// then by `Topology::Decay::flat_propagator_index`. Falls back to /// `Topology::Decay::pdg_id` if empty. nested_vector3 diagram_propagator_pdgs; + /// Relative weight of each PDG-id option of each flavor: an option is + /// written with probability proportional to its weight. Uniform if empty, + /// or for a flavor whose weights are empty. + nested_vector2 pdg_id_weights; }; /// @param subproc_args Bookkeeping for every subprocess, indexed like @@ -244,6 +248,9 @@ class LHECompleter { std::vector _helicities; std::vector> _pdg_id_and_count; std::vector _pdg_ids; + // flavor slot -> offset of its cumulative option probabilities, if not uniform + std::unordered_map _pdg_id_weight_offsets; + std::vector _pdg_id_cumulative; std::unordered_map> _propagator_index_and_count; std::vector _propagators; diff --git a/madspace/src/driver/lhe_output.cpp b/madspace/src/driver/lhe_output.cpp index 6064f1d7ea..fc9f66cbda 100644 --- a/madspace/src/driver/lhe_output.cpp +++ b/madspace/src/driver/lhe_output.cpp @@ -3,6 +3,7 @@ #include "madspace/util.hpp" #include +#include #include #include @@ -190,12 +191,47 @@ LHECompleter::append_colors(const SubprocArgs& args, std::size_t particle_count) } void LHECompleter::append_pdg_ids(const SubprocArgs& args, std::size_t particle_count) { - for (auto& pdg_id_options : args.pdg_ids) { + if (args.pdg_id_weights.size() != 0 && + args.pdg_id_weights.size() != args.pdg_ids.size()) { + throw std::invalid_argument("Need PDG-id option weights for every flavor"); + } + for (std::size_t flavor = 0; auto& pdg_id_options : args.pdg_ids) { if (pdg_id_options.size() == 0) { throw std::invalid_argument( "Must provide at least one option per flavor index" ); } + std::span weights; + if (args.pdg_id_weights.size() != 0) { + weights = args.pdg_id_weights.at(flavor); + } + ++flavor; + if (weights.size() != 0) { + if (weights.size() != pdg_id_options.size()) { + throw std::invalid_argument("Need one weight per PDG-id option"); + } + double total = 0.; + for (double weight : weights) { + if (!(weight > 0.) || !std::isfinite(weight)) { + throw std::invalid_argument( + "PDG-id option weights must be positive and finite" + ); + } + total += weight; + } + // equal weights keep the uniform draw + if (std::any_of(weights.begin(), weights.end(), [&](double weight) { + return weight != weights[0]; + })) { + _pdg_id_weight_offsets[_pdg_id_and_count.size()] = + _pdg_id_cumulative.size(); + double sum = 0.; + for (double weight : weights) { + sum += weight; + _pdg_id_cumulative.push_back(sum / total); + } + } + } _pdg_id_and_count.push_back({_pdg_ids.size(), pdg_id_options.size()}); for (auto& pdg_ids : pdg_id_options) { if (pdg_ids.size() != particle_count) { @@ -560,9 +596,20 @@ void LHECompleter::complete_event_data( subproc_data.helicity_offset + subproc_data.particle_count * helicity_index; std::size_t mass_offset = subproc_data.mass_offset; - auto [pdg_index, pdg_count] = - _pdg_id_and_count.at(subproc_data.pdg_id_offset + flavor_index); - std::size_t pdg_random = rand_gen.generate_int(pdg_count); + std::size_t flavor_slot = subproc_data.pdg_id_offset + flavor_index; + auto [pdg_index, pdg_count] = _pdg_id_and_count.at(flavor_slot); + std::size_t pdg_random; + if (auto weights = _pdg_id_weight_offsets.find(flavor_slot); + weights != _pdg_id_weight_offsets.end()) { + auto begin = _pdg_id_cumulative.begin() + weights->second; + pdg_random = std::min( + std::upper_bound(begin, begin + pdg_count, rand_gen.generate_double()) - + begin, + pdg_count - 1 + ); + } else { + pdg_random = rand_gen.generate_int(pdg_count); + } std::size_t pdg_offset = pdg_index + subproc_data.particle_count * pdg_random; for (std::size_t particle_index = 0; auto& particle : event.particles) { @@ -744,6 +791,8 @@ void madspace::to_json(nlohmann::json& j, const LHECompleter& lhe_completer) { {"helicities", lhe_completer._helicities}, {"pdg_id_and_count", lhe_completer._pdg_id_and_count}, {"pdg_ids", lhe_completer._pdg_ids}, + {"pdg_id_weight_offsets", lhe_completer._pdg_id_weight_offsets}, + {"pdg_id_cumulative", lhe_completer._pdg_id_cumulative}, {"propagator_index_and_count", propagator_index_and_count}, {"propagators", lhe_completer._propagators}, {"propagator_colors", lhe_completer._propagator_colors}, @@ -762,6 +811,12 @@ void madspace::from_json(const nlohmann::json& j, LHECompleter& lhe_completer) { lhe_completer._pdg_id_and_count = j.at("pdg_id_and_count").get>>(); lhe_completer._pdg_ids = j.at("pdg_ids").get>(); + // absent from completers saved before weighted options: all uniform + lhe_completer._pdg_id_weight_offsets = j.value( + "pdg_id_weight_offsets", std::unordered_map{} + ); + lhe_completer._pdg_id_cumulative = + j.value("pdg_id_cumulative", std::vector{}); lhe_completer._propagator_index_and_count = {}; for (auto& item : j.at("propagator_index_and_count")) { lhe_completer._propagator_index_and_count[item.at(0).get()] = diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 0293f72782..c93b6d3a13 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -3193,7 +3193,8 @@ PYBIND11_MODULE(_madspace_py, m) { std::unordered_map, nested_vector2, nested_vector3, - nested_vector3>(), + nested_vector3, + nested_vector2>(), py::arg("process_id") = 0, py::arg("topologies") = std::vector{}, py::arg("permutations") = nested_vector3{}, @@ -3203,7 +3204,8 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("pdg_color_types") = std::unordered_map{}, py::arg("helicities") = nested_vector2{}, py::arg("pdg_ids") = nested_vector3{}, - py::arg("diagram_propagator_pdgs") = nested_vector3{} + py::arg("diagram_propagator_pdgs") = nested_vector3{}, + py::arg("pdg_id_weights") = nested_vector2{} ) .def_readwrite( "process_id", @@ -3254,6 +3256,11 @@ PYBIND11_MODULE(_madspace_py, m) { "diagram_propagator_pdgs", &LHECompleter::SubprocArgs::diagram_propagator_pdgs, pydoc::doc("LHECompleter::SubprocArgs::diagram_propagator_pdgs") + ) + .def_readwrite( + "pdg_id_weights", + &LHECompleter::SubprocArgs::pdg_id_weights, + pydoc::doc("LHECompleter::SubprocArgs::pdg_id_weights") ); py::classh(m, "MixMaxRandom", pydoc::doc("MixMaxRandom")) .def(py::init<>(), pydoc::doc("MixMaxRandom::MixMaxRandom")) diff --git a/madspace/tests/test_lhe.py b/madspace/tests/test_lhe.py index 9d6827c7a6..9f83841b0e 100644 --- a/madspace/tests/test_lhe.py +++ b/madspace/tests/test_lhe.py @@ -535,3 +535,61 @@ def test_subprocess_without_diagrams(): assert [p.spin for p in event.particles] == [-1, -1, 1, 1, -1] with pytest.raises(RuntimeError, match="Diagram index out of range"): lhe_completer.complete_event_data(event, 0, 0, 0, 0, 0, ms.MixMaxRandom(1)) + + +def weighted_flavor_completer(pdg_id_weights): + """e+ e- > l+ l- with two flavors of two options each.""" + options = [ + [[11, -11, 13, -13], [11, -11, 15, -15]], + [[-11, 11, 13, -13], [-11, 11, 15, -15]], + ] + subproc_args = ms.SubprocArgs( + topologies=[chain_topology([0.0] * 4)], + permutations=[[]], + diagram_indices=[[]], + diagram_color_indices=[[]], + color_flows=[[[0, 0]] * 4], + pdg_color_types={pdg: 1 for pdg in (11, -11, 13, -13, 15, -15)}, + helicities=[[1, -1, 1, -1]], + pdg_ids=options, + pdg_id_weights=pdg_id_weights, + ) + return ms.LHECompleter([subproc_args], bw_cutoff=BW_CUTOFF) + + +def drawn_taus(lhe_completer, flavor_index, count, seed): + rand_gen = ms.MixMaxRandom(seed) + taus = 0 + for _ in range(count): + event = build_event(np.ones((4, 4)) * [100.0, 0.0, 0.0, 0.0]) + lhe_completer.complete_event_data(event, 0, -1, 0, flavor_index, 0, rand_gen) + taus += event.particles[2].pdg_id == 15 + return taus + + +def test_pdg_id_options_are_drawn_in_proportion_to_their_weights(): + """Options are written with probability proportional to their weight; + equal weights and flavors without weights draw them uniformly.""" + lhe_completer = weighted_flavor_completer([[1.0, 3.0], []]) + count = 4000 + assert drawn_taus(lhe_completer, 0, count, 5) / count == approx(0.75, abs=0.03) + assert drawn_taus(lhe_completer, 1, count, 5) / count == approx(0.5, abs=0.03) + + with tempfile.TemporaryDirectory() as tmpdir: + save_path = os.path.join(tmpdir, "lhe_completer.json") + lhe_completer.save(save_path) + loaded = ms.LHECompleter.load(save_path) + assert drawn_taus(loaded, 0, 200, 9) == drawn_taus(lhe_completer, 0, 200, 9) + + # equal weights consume the random numbers like the uniform draw + equal = weighted_flavor_completer([[2.0, 2.0], [2.0, 2.0]]) + uniform = weighted_flavor_completer([]) + assert drawn_taus(equal, 0, 200, 3) == drawn_taus(uniform, 0, 200, 3) + + +@pytest.mark.parametrize( + "pdg_id_weights", [[[1.0, 0.0], []], [[1.0, -1.0], []], [[1.0], []], [[1.0, 1.0]]] +) +def test_invalid_pdg_id_weights_raise(pdg_id_weights): + with pytest.raises(ValueError): + weighted_flavor_completer(pdg_id_weights) From daf23c97de0bca5af4b4e402e034b060c5e63e5c Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 26 Oct 2026 14:47:59 +0100 Subject: [PATCH 44/88] mg7: [amplicol] generates one subprocess per verified flavour class The worker generates with pyamplicol's strict flavour-blind grouping, so subprocesses related by a flavour relabelling that is a symmetry of the allowed vertices are evaluated through their representative. The rows of a UMAMI selector become the options of one flavour per mirror flag, written in proportion to their factors, whose sum is the flavour factor. --- madgraph/iolibs/amplicol_worker.py | 1 + madgraph/iolibs/export_amplicol.py | 22 +++++++---- .../template_files/mg7/subprocess_amplicol.py | 1 + .../interface/test_amplicol_syntax.py | 38 +++++++++++++++++++ 4 files changed, 55 insertions(+), 7 deletions(-) diff --git a/madgraph/iolibs/amplicol_worker.py b/madgraph/iolibs/amplicol_worker.py index b546a6284a..34fb303f6b 100644 --- a/madgraph/iolibs/amplicol_worker.py +++ b/madgraph/iolibs/amplicol_worker.py @@ -45,6 +45,7 @@ def generate_artifact(pyamplicol, artifact, accuracy, process, model): 'action': 'generate', 'color': {'accuracy': accuracy}, 'process': {'multiparticles': process['multiparticles'], + 'flavor_blind_strict': True, 'coupling_order_policy': 'explicit', 'max_coupling_orders': process['orders']}, 'generation': {'emit_api_bundle': True, diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index 84141aa7f7..ade077d8f6 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -19,7 +19,10 @@ (full), and the sources of the UMAMI library returning their ratio (reweight). Its makefile builds the libraries into lib/. Every colour ordering (UMAMI channel) of every provider becomes one entry of subprocesses.json, integrated -on its own with a colour-ordered phase-space mapping. +on its own with a colour-ordered phase-space mapping. Subprocesses related by a +flavour relabelling that is a symmetry of the allowed vertices are generated +once (pyamplicol's strict flavour-blind grouping) and evaluated through their +representative. """ import json @@ -278,14 +281,19 @@ def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power, raise madgraph.MadGraph5Error( 'AmpliCol contributions shared by several channels are not ' 'supported yet') + # the rows of a selector share its evaluation and incoming legs: one + # flavour per mirror flag, its options written in proportion to factor + by_mirror = {} for term in process['matrix_elements']: pdgs = pdg_ids[term['pdg_ids']] - flavors.append({ - 'index': flavor_index, - 'options': [pdgs], - 'factor': term['factor'], - 'mirror': mirrored(pdgs, incoming, initial_legs, exported), - }) + mirror = mirrored(pdgs, incoming, initial_legs, exported) + flavor = by_mirror.setdefault(mirror, { + 'index': flavor_index, 'options': [], 'weights': [], + 'factor': 0, 'mirror': mirror}) + flavor['options'].append(pdgs) + flavor['weights'].append(term['factor']) + flavor['factor'] += term['factor'] + flavors.extend(by_mirror.values()) first = flavors[0]['options'][0] entries.append({ 'backend': 'amplicol', diff --git a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py index 28e8967a4a..965b26c5e0 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py +++ b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py @@ -69,6 +69,7 @@ def amplicol_subproc_args(meta: dict, process: MadgraphProcess) -> ms.SubprocArg }, helicities=lhe_helicities(meta, process.run_card), pdg_ids=[flavor["options"] for flavor in meta["flavors"]], + pdg_id_weights=[flavor.get("weights", []) for flavor in meta["flavors"]], ) diff --git a/tests/unit_tests/interface/test_amplicol_syntax.py b/tests/unit_tests/interface/test_amplicol_syntax.py index 854b7bb3ff..c08d9f8660 100644 --- a/tests/unit_tests/interface/test_amplicol_syntax.py +++ b/tests/unit_tests/interface/test_amplicol_syntax.py @@ -151,3 +151,41 @@ def test_mirrored_beams(self): self.assertFalse(mirrored([2, -2, 2, -2, 21], 2, [partons, partons], exported)) # g u > u g g only self.assertFalse(mirrored([21, 2, 21, 21, 2], 2, [{21}, {2}], exported)) + + def test_selector_rows_are_folded_into_weighted_options(self): + """The rows of a UMAMI selector share its evaluation: they become the + options of one flavour per mirror flag, written in proportion to their + factors, which add up to the flavour factor.""" + import madgraph.iolibs.export_amplicol as export_amplicol + metadata = { + 'pdg_ids': [[1, -1, 2, -2, 21], [1, -1, 3, -3, 21], + [1, -1, 4, -4, 21], [-1, 1, 3, -3, 21]], + 'provider': {'incoming_count': 2, 'masses': [0.] * 5, + 'supports_alpha_s': True}, + 'channels': [{ + 'phasespace_order': [0, 2, 3, 4, 1], + 'processes': [ + {'multichannels': [0], 'matrix_elements': [ + {'pdg_ids': 0, 'factor': 1}, {'pdg_ids': 2, 'factor': 2}, + {'pdg_ids': 1, 'factor': 1}]}, + {'multichannels': [0], 'matrix_elements': [ + {'pdg_ids': 3, 'factor': 1}]}, + ], + }], + 'color_flows': [], 'helicities': [], 'runtime_processes': [], + } + partons = {21, 1, 2, 3, 4, -1, -2, -3, -4} + entries = export_amplicol.subprocess_entries( + metadata, self.cmd._curr_model, 'SubProcesses/A1', 'me.so', 'ratio.so', + 3, [partons, partons]) + self.assertEqual(len(entries), 1) + flavors = entries[0]['flavors'] + # d d~ > s s~ g has its mirror exported, the other rows are mirrored + self.assertEqual(flavors, [ + {'index': 0, 'options': [[1, -1, 2, -2, 21], [1, -1, 4, -4, 21]], + 'weights': [1, 2], 'factor': 3, 'mirror': True}, + {'index': 0, 'options': [[1, -1, 3, -3, 21]], 'weights': [1], + 'factor': 1, 'mirror': False}, + {'index': 1, 'options': [[-1, 1, 3, -3, 21]], 'weights': [1], + 'factor': 1, 'mirror': False}, + ]) From e467b9a60f9ba401fac9ec16aae5fb4627b795cd Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 26 Oct 2026 17:15:31 +0100 Subject: [PATCH 45/88] pyamplicol: pack per-process payloads into payloads.pacbin With generation.pack_payloads, the payloads under processes/ (except nested pacbin containers and the recurrence bootstrap), recurrence/schedules/ and model/eager-kernels/ are stored in one payloads.pacbin, declared with the new payload-container role. Their manifest records are kept, so the artifact ID does not change. The runtime opens the container with positioned reads of its index and checks every member it reads against its indexed SHA-256 (pacbin's PacbinFileReader). read_payload, payload_path, the evaluator payload store and the schedule loader serve packed payloads from it; the recurrence-bootstrap path uses the members its inventory declares. The builder removes the directories left empty by payloads moved into containers. The [amplicol] worker packs: p p > j j j writes 73 files instead of 2033, and a runtime load opens 4 artifact files instead of 207, with the same matrix elements. --- madgraph/iolibs/amplicol_worker.py | 3 +- pyamplicol/examples/all_options.toml | 2 + .../rust/crates/rusticol-core/src/artifact.rs | 181 +++++++++++++- .../src/engine/recurrence_load.rs | 6 +- .../rust/crates/rusticol-core/src/pacbin.rs | 230 +++++++++++++++--- .../crates/rusticol-core/src/pacbin_tests.rs | 30 +++ .../src/pyamplicol/artifacts/manifest.py | 1 + pyamplicol/src/pyamplicol/artifacts/writer.py | 10 + pyamplicol/src/pyamplicol/cli/parser.py | 6 + pyamplicol/src/pyamplicol/config/models.py | 4 + .../pyamplicol/generation/artifact_writer.py | 121 ++++++++- .../generation/evaluator_container.py | 2 + pyamplicol/tests/unit/test_artifacts.py | 96 ++++++++ pyamplicol/tests/unit/test_config_models.py | 2 +- 14 files changed, 643 insertions(+), 51 deletions(-) diff --git a/madgraph/iolibs/amplicol_worker.py b/madgraph/iolibs/amplicol_worker.py index b546a6284a..73f870a345 100644 --- a/madgraph/iolibs/amplicol_worker.py +++ b/madgraph/iolibs/amplicol_worker.py @@ -47,7 +47,8 @@ def generate_artifact(pyamplicol, artifact, accuracy, process, model): 'process': {'multiparticles': process['multiparticles'], 'coupling_order_policy': 'explicit', 'max_coupling_orders': process['orders']}, - 'generation': {'emit_api_bundle': True, + # pack_payloads: one payloads.pacbin instead of files per process + 'generation': {'emit_api_bundle': True, 'pack_payloads': True, 'relation_discovery': {'mode': 'off'}}, 'evaluator': {'backend': 'jit', 'execution_mode': 'recurrence', 'jit': {'optimization_level': 2}}, diff --git a/pyamplicol/examples/all_options.toml b/pyamplicol/examples/all_options.toml index c0eb70f3a1..bb7fbabc20 100644 --- a/pyamplicol/examples/all_options.toml +++ b/pyamplicol/examples/all_options.toml @@ -102,6 +102,8 @@ mode = "error" workers = "auto" # generation.emit_api_bundle: Emit Python/Rust/C/C++/Fortran/UMAMI API bundle. Default true. emit_api_bundle = true +# generation.pack_payloads: Store per-process payloads and recurrence schedules in one payloads.pacbin. Default false. +pack_payloads = false [generation.validation] # generation.validation.enabled: Run generation validation. Default true. enabled = true diff --git a/pyamplicol/rust/crates/rusticol-core/src/artifact.rs b/pyamplicol/rust/crates/rusticol-core/src/artifact.rs index d7b6d61a29..75e7a2fe0c 100644 --- a/pyamplicol/rust/crates/rusticol-core/src/artifact.rs +++ b/pyamplicol/rust/crates/rusticol-core/src/artifact.rs @@ -1,6 +1,6 @@ // SPDX-License-Identifier: 0BSD -use crate::pacbin::{PacbinMemberKind, PacbinReader}; +use crate::pacbin::{PacbinFileReader, PacbinMemberKind, PacbinReader}; use crate::{ ARTIFACT_MANIFEST_FILE, C_ABI_VERSION, COMPILED_MODEL_SCHEMA_VERSION, PROCESS_ARTIFACT_SCHEMA_VERSION, PYTHON_API_VERSION, RUNTIME_PHYSICS_SCHEMA_VERSION, @@ -40,6 +40,8 @@ const EVALUATOR_PAYLOAD_CONTAINER_KIND: &str = "pyamplicol-evaluator-payload-con const EVALUATOR_PAYLOAD_CONTAINER_STORAGE_ABI: &str = "pacbin-v1"; #[cfg(any(feature = "f64-compiled", feature = "f64-symjit"))] const RECURRENCE_BOOTSTRAP_FILENAME: &str = "recurrence-bootstrap.bin"; +/// Holds declared payloads that are not stored as loose files. +pub const PAYLOAD_CONTAINER_PATH: &str = "payloads.pacbin"; const SUPPORTED_ARTIFACT_TARGETS: [&str; 3] = [ "aarch64-apple-darwin", "x86_64-apple-darwin", @@ -209,6 +211,7 @@ pub enum PayloadRole { ApiBuildFile, SdkMetadata, ProcessGroups, + PayloadContainer, } #[derive(Clone, Debug, Decode, Deserialize, Eq, PartialEq, Serialize)] @@ -317,6 +320,7 @@ pub struct VerifiedArtifact { manifest: ArtifactManifest, payloads: Arc>, evaluator_payload_container: Option>, + payload_container: Option>, #[cfg(feature = "f64-compiled")] native_library_cache: Arc>>>, } @@ -363,6 +367,11 @@ pub(crate) enum EvaluatorPayloadSource { container: Arc, logical_path: String, }, + /// A declared payload stored in the artifact's payload container. + PackedDeclared { + container: Arc, + payload: Payload, + }, } /// A dynamic library loaded only from process-private snapshot bytes. @@ -633,6 +642,9 @@ impl EvaluatorPayloadSource { container, logical_path, } => container.member_bytes(logical_path).map(Cow::Borrowed), + Self::PackedDeclared { container, payload } => { + read_packed_payload(container, payload).map(Cow::Owned) + } } } @@ -643,6 +655,9 @@ impl EvaluatorPayloadSource { root.join(&payload.path).display().to_string() } Self::Packed { logical_path, .. } => format!("evaluators.pacbin:{logical_path}"), + Self::PackedDeclared { payload, .. } => { + format!("{PAYLOAD_CONTAINER_PATH}:{}", payload.path) + } } } } @@ -654,6 +669,7 @@ pub(crate) struct EvaluatorPayloadStore { relative_root: PathBuf, container: Option>, payloads: Option>>, + payload_container: Option>, #[cfg(feature = "f64-compiled")] native_library_cache: Arc>>>, } @@ -665,6 +681,7 @@ impl EvaluatorPayloadStore { relative_root: root.to_path_buf(), container: None, payloads: None, + payload_container: None, #[cfg(feature = "f64-compiled")] native_library_cache: Arc::new(Mutex::new(BTreeMap::new())), } @@ -698,6 +715,14 @@ impl EvaluatorPayloadStore { payload.role ))); } + if let Some(container) = &self.payload_container + && container.contains(&logical_path) + { + return Ok(EvaluatorPayloadSource::PackedDeclared { + container: container.clone(), + payload: payload.clone(), + }); + } return Ok(EvaluatorPayloadSource::AuthenticatedFile { root: self.artifact_root.clone(), payload: payload.clone(), @@ -914,6 +939,7 @@ impl VerifiedArtifact { &decoded.evaluator_payload_container, &payloads, )?; + let payload_container = load_bootstrap_payload_container(&root, &payloads)?; let artifact = Self { root, manifest_path, @@ -923,6 +949,7 @@ impl VerifiedArtifact { // leaves. Avoid opening and retaining the model-wide evaluator // PACBIN for source/intrinsic-only recurrence processes. evaluator_payload_container: None, + payload_container, #[cfg(feature = "f64-compiled")] native_library_cache: Arc::new(Mutex::new(BTreeMap::new())), }; @@ -1004,6 +1031,7 @@ impl VerifiedArtifact { validate_references(&manifest, &payloads)?; let evaluator_payload_container = load_evaluator_payload_container(&root, &manifest, &payloads)?; + let payload_container = load_payload_container(&root, &payloads)?; let payloads = Arc::new(payloads); Ok(Self { root, @@ -1011,6 +1039,7 @@ impl VerifiedArtifact { manifest, payloads, evaluator_payload_container, + payload_container, #[cfg(feature = "f64-compiled")] native_library_cache: Arc::new(Mutex::new(BTreeMap::new())), }) @@ -1050,7 +1079,7 @@ impl VerifiedArtifact { "artifact declares more than one process-groups payload", )); } - let bytes = read_declared_payload(&self.root, payload)?; + let bytes = self.read_declared(payload)?; if format!("{:x}", Sha256::digest(&bytes)) != payload.sha256 { return Err(RusticolError::integrity(format!( "process-groups payload {:?} does not match its declared sha256", @@ -1068,7 +1097,42 @@ impl VerifiedArtifact { pub fn read_payload(&self, path: &str) -> RusticolResult> { let payload = self.payload(path)?; - read_declared_payload(&self.root, payload) + self.read_declared(payload) + } + + fn read_declared(&self, payload: &Payload) -> RusticolResult> { + match &self.payload_container { + Some(container) if container.contains(&payload.path) => { + read_packed_payload(container, payload) + } + _ => read_declared_payload(&self.root, payload), + } + } + + fn is_packed(&self, path: &str) -> bool { + self.payload_container + .as_ref() + .is_some_and(|container| container.contains(path)) + } + + /// Open a declared payload that is itself a pacbin container, loose or packed, + /// and authenticate it against `expected_sha256`. + pub(crate) fn open_payload_pacbin( + &self, + path: &str, + expected_sha256: &[u8; 32], + ) -> RusticolResult { + if self.is_packed(path) { + let bytes = self.read_payload(path)?; + if <[u8; 32]>::from(Sha256::digest(&bytes)) != *expected_sha256 { + return Err(RusticolError::integrity(format!( + "pacbin container {path:?} payload digest mismatch" + ))); + } + return PacbinReader::from_bytes(bytes); + } + let file = self.open_payload_file(path)?; + PacbinReader::open_file_with_sha256(file, &self.root.join(path), expected_sha256) } /// Open one declared loose payload and pin the checked file description. @@ -1078,12 +1142,20 @@ impl VerifiedArtifact { /// the bytes between declaration checks and use. pub(crate) fn open_payload_file(&self, path: &str) -> RusticolResult { let payload = self.payload(path)?; + if self.is_packed(path) { + return Err(RusticolError::compatibility(format!( + "payload {path:?} is stored in {PAYLOAD_CONTAINER_PATH} and cannot be opened as a file" + ))); + } open_checked_payload(&self.root, payload) } + /// The location of a declared payload; packed payloads have no file there. pub(crate) fn payload_path(&self, path: &str) -> RusticolResult { let payload = self.payload(path)?; - drop(open_checked_payload(&self.root, payload)?); + if !self.is_packed(path) { + drop(open_checked_payload(&self.root, payload)?); + } Ok(self.root.join(path)) } @@ -1101,6 +1173,7 @@ impl VerifiedArtifact { relative_root: relative_root.to_path_buf(), container: self.evaluator_payload_container.clone(), payloads: Some(self.payloads.clone()), + payload_container: self.payload_container.clone(), #[cfg(feature = "f64-compiled")] native_library_cache: self.native_library_cache.clone(), }) @@ -1524,6 +1597,7 @@ fn payload_role_name(role: PayloadRole) -> &'static str { PayloadRole::ApiBuildFile => "api-build-file", PayloadRole::SdkMetadata => "sdk-metadata", PayloadRole::ProcessGroups => "process-groups", + PayloadRole::PayloadContainer => "payload-container", } } @@ -1771,11 +1845,110 @@ pub(crate) fn open_staged_recurrence_bootstrap_artifact( payloads: Arc::new(payloads), // Process-ready lowering reads only loose plan/binding/color payloads. evaluator_payload_container: None, + payload_container: None, #[cfg(feature = "f64-compiled")] native_library_cache: Arc::new(Mutex::new(BTreeMap::new())), }) } +/// Open the payload container, checking that every member is a declared payload +/// with the declared size and digest. +fn load_payload_container( + root: &Path, + payloads: &BTreeMap, +) -> RusticolResult>> { + let mut declared = payloads + .values() + .filter(|payload| payload.role == PayloadRole::PayloadContainer); + let Some(payload) = declared.next() else { + return Ok(None); + }; + if declared.next().is_some() || payload.path != PAYLOAD_CONTAINER_PATH { + return Err(RusticolError::artifact(format!( + "an artifact declares at most one payload container, at {PAYLOAD_CONTAINER_PATH:?}" + ))); + } + open_payload_container(root, payload, payloads, true).map(Some) +} + +/// The payload container next to a recurrence bootstrap. Its inventory predates +/// the container, so only the members it declares are checked and used. +fn load_bootstrap_payload_container( + root: &Path, + payloads: &BTreeMap, +) -> RusticolResult>> { + let metadata = match fs::symlink_metadata(root.join(PAYLOAD_CONTAINER_PATH)) { + Ok(metadata) => metadata, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => { + return Err(RusticolError::security(format!( + "could not inspect {PAYLOAD_CONTAINER_PATH:?}: {error}" + ))); + } + }; + let container = Payload { + path: PAYLOAD_CONTAINER_PATH.to_string(), + role: PayloadRole::PayloadContainer, + media_type: "application/octet-stream".to_string(), + size_bytes: metadata.len(), + sha256: String::new(), + executable: false, + target: None, + process_id: None, + }; + open_payload_container(root, &container, payloads, false).map(Some) +} + +fn open_payload_container( + root: &Path, + payload: &Payload, + payloads: &BTreeMap, + require_declared: bool, +) -> RusticolResult> { + let file = open_checked_payload(root, payload)?; + let reader = PacbinFileReader::open_file(file, &root.join(&payload.path))?; + if reader.index().file_size() != payload.size_bytes { + return Err(RusticolError::integrity( + "payload container size disagrees with its declaration", + )); + } + for member in reader.index().members() { + let Some(packed) = payloads.get(member.logical_path()) else { + if require_declared { + return Err(RusticolError::integrity(format!( + "payload container member {:?} is not a declared payload", + member.logical_path() + ))); + } + continue; + }; + if member.kind() != PacbinMemberKind::DeclaredPayload + || packed.role == PayloadRole::PayloadContainer + || packed.executable + || member.length() != packed.size_bytes + || hex_digest(member.sha256()) != packed.sha256 + { + return Err(RusticolError::integrity(format!( + "invalid payload container member {:?}", + member.logical_path() + ))); + } + } + Ok(Arc::new(reader)) +} + +fn read_packed_payload(container: &PacbinFileReader, payload: &Payload) -> RusticolResult> { + validate_payload_declaration(payload)?; + let bytes = container.read_member(&payload.path)?; + if u64::try_from(bytes.len()).ok() != Some(payload.size_bytes) { + return Err(RusticolError::integrity(format!( + "packed payload {:?} disagrees with its declared size", + payload.path + ))); + } + Ok(bytes) +} + fn load_evaluator_payload_container( root: &Path, manifest: &ArtifactManifest, diff --git a/pyamplicol/rust/crates/rusticol-core/src/engine/recurrence_load.rs b/pyamplicol/rust/crates/rusticol-core/src/engine/recurrence_load.rs index 943c4722c7..2982579969 100644 --- a/pyamplicol/rust/crates/rusticol-core/src/engine/recurrence_load.rs +++ b/pyamplicol/rust/crates/rusticol-core/src/engine/recurrence_load.rs @@ -27,7 +27,7 @@ use super::recurrence_process_pack::{ decode_process_executor_descriptors, read_string, semantic_digest_from_bytes, }; use super::*; -use crate::pacbin::{PacbinMemberKind, PacbinReader}; +use crate::pacbin::PacbinMemberKind; #[cfg(feature = "on-the-fly-test-support")] use crate::recurrence::direct_backend::{ DirectCurrentObservation, begin_direct_current_observation, take_direct_current_observation, @@ -1248,7 +1248,6 @@ fn load_plan_reference( )?; validate_ready_plan_reference(&spec)?; let container = spec.plan.runtime_schedule; - let path = artifact.root().join(&container.path); let payload = artifact.payload(&container.path)?; if payload.role != PayloadRole::EvaluatorState || payload.media_type != "application/octet-stream" @@ -1265,8 +1264,7 @@ fn load_plan_reference( let process_binding = validate_process_binding(artifact, evaluator_root, &spec)?; let expected_file_sha = decode_sha256(&container.sha256)?; - let container_file = artifact.open_payload_file(&container.path)?; - let reader = PacbinReader::open_file_with_sha256(container_file, &path, &expected_file_sha)?; + let reader = artifact.open_payload_pacbin(&container.path, &expected_file_sha)?; let index = reader.index(); if index.file_size() != container.size_bytes || reader.container_size() as u64 != container.size_bytes diff --git a/pyamplicol/rust/crates/rusticol-core/src/pacbin.rs b/pyamplicol/rust/crates/rusticol-core/src/pacbin.rs index cf72999a47..7a7ebe7080 100644 --- a/pyamplicol/rust/crates/rusticol-core/src/pacbin.rs +++ b/pyamplicol/rust/crates/rusticol-core/src/pacbin.rs @@ -53,6 +53,8 @@ pub enum PacbinMemberKind { RecurrenceColorProjectionCertificate = 8, OnTheFlyProcessSeed = 9, ColorContraction = 10, + /// A payload declared by the artifact manifest, stored in its payload container. + DeclaredPayload = 11, } impl PacbinMemberKind { @@ -67,6 +69,7 @@ impl PacbinMemberKind { 8 => Ok(Self::RecurrenceColorProjectionCertificate), 9 => Ok(Self::OnTheFlyProcessSeed), 10 => Ok(Self::ColorContraction), + 11 => Ok(Self::DeclaredPayload), _ => Err(compatibility(format!( "unknown pacbin member kind: {value}" ))), @@ -590,6 +593,130 @@ impl PacbinReader { } } +/// A container read member by member from its open file. +/// +/// Only the frame and index are read when opening, so a large container costs +/// neither memory nor a full digest. Each member is copied out and checked +/// against its indexed SHA-256 before use. +#[derive(Debug)] +pub struct PacbinFileReader { + file: std::sync::Mutex, + display_path: PathBuf, + index: PacbinIndex, +} + +impl PacbinFileReader { + pub(crate) fn open_file(mut file: File, display_path: &Path) -> RusticolResult { + let file_size = file + .metadata() + .map_err(|error| { + RusticolError::artifact(format!( + "could not inspect pacbin container {}: {error}", + display_path.display() + )) + })? + .len(); + let minimum_size = (HEADER_SIZE + INDEX_HEADER_SIZE + FOOTER_SIZE) as u64; + if file_size < minimum_size { + return Err(integrity("truncated pacbin container")); + } + let header = read_file_range(&mut file, display_path, 0, HEADER_SIZE as u64)?; + let footer = read_file_range( + &mut file, + display_path, + file_size - FOOTER_SIZE as u64, + FOOTER_SIZE as u64, + )?; + let frame = parse_frame(&header, &footer, file_size)?; + let index_bytes = read_file_range( + &mut file, + display_path, + frame.index_offset, + frame.index_length, + )?; + let (members, index_sha256) = parse_index(&index_bytes, &frame)?; + // Padding is never handed out, so it is not read back here. + validate_canonical_payload_layout(&members, frame.index_offset, |_, _| Ok(()))?; + Ok(Self { + file: std::sync::Mutex::new(file), + display_path: display_path.to_path_buf(), + index: PacbinIndex { + version: frame.version, + index_offset: frame.index_offset, + index_sha256, + file_size, + members, + }, + }) + } + + pub fn index(&self) -> &PacbinIndex { + &self.index + } + + pub fn member(&self, logical_path: &str) -> RusticolResult<&PacbinMember> { + let normalized = normalize_logical_path(logical_path)?; + self.index + .members + .binary_search_by(|member| member.logical_path.as_bytes().cmp(normalized.as_bytes())) + .map(|index| &self.index.members[index]) + .map_err(|_| { + RusticolError::invalid_argument(format!("unknown pacbin member: {normalized}")) + }) + } + + pub fn contains(&self, logical_path: &str) -> bool { + self.member(logical_path).is_ok() + } + + /// Copy one member out of the container and authenticate the copy. + pub fn read_member(&self, logical_path: &str) -> RusticolResult> { + let member = self.member(logical_path)?; + let bytes = { + let mut file = self + .file + .lock() + .map_err(|_| RusticolError::artifact("pacbin container file handle is poisoned"))?; + read_file_range(&mut file, &self.display_path, member.offset, member.length)? + }; + let digest: [u8; 32] = Sha256::digest(&bytes).into(); + if digest != member.sha256 { + return Err(integrity(format!( + "pacbin member digest mismatch: {}", + member.logical_path + ))); + } + Ok(bytes) + } +} + +fn read_file_range( + file: &mut File, + display_path: &Path, + offset: u64, + length: u64, +) -> RusticolResult> { + let length = usize::try_from(length) + .map_err(|_| integrity("pacbin read length exceeds platform bounds"))?; + let mut bytes = Vec::new(); + bytes.try_reserve_exact(length).map_err(|error| { + RusticolError::artifact(format!( + "could not allocate {length} bytes for pacbin container {}: {error}", + display_path.display() + )) + })?; + bytes.resize(length, 0); + file.seek(SeekFrom::Start(offset)) + .and_then(|_| file.read_exact(&mut bytes)) + .map_err(|error| { + RusticolError::artifact(format!( + "could not read pacbin container {}: {error}", + display_path.display() + )) + })?; + Ok(bytes) +} + /// Atomically stream and publish one deterministic `pacbin-v1` container. /// /// Member metadata is retained in memory, while payload bytes are copied and @@ -1040,8 +1167,48 @@ fn parse_and_validate(bytes: &[u8], verify_payloads: bool) -> RusticolResult RusticolResult { if &header[0..8] != HEADER_MAGIC { return Err(integrity("invalid pacbin header magic")); } @@ -1066,17 +1233,13 @@ fn parse_and_validate(bytes: &[u8], verify_payloads: bool) -> RusticolResult= footer_offset_u64 { return Err(integrity("pacbin index offset is out of bounds")); } - let footer = checked_slice(bytes, footer_offset, FOOTER_SIZE, "pacbin footer")?; if &footer[0..8] != FOOTER_MAGIC { return Err(integrity("invalid pacbin footer magic")); } @@ -1104,17 +1267,23 @@ fn parse_and_validate(bytes: &[u8], verify_payloads: bool) -> RusticolResult RusticolResult<(Vec, [u8; 32])> { + let member_count = frame.member_count; let actual_index_digest: [u8; 32] = Sha256::digest(index_bytes).into(); let index_header = checked_slice(index_bytes, 0, INDEX_HEADER_SIZE, "pacbin index header")?; @@ -1228,21 +1397,11 @@ fn parse_and_validate(bytes: &[u8], verify_payloads: bool) -> RusticolResult Rusti } fn validate_canonical_payload_layout( - bytes: &[u8], members: &[PacbinMember], index_offset: u64, + mut validate_padding: impl FnMut(u64, u64) -> RusticolResult<()>, ) -> RusticolResult<()> { let mut expected_offset = HEADER_SIZE as u64; for member in members { @@ -1323,7 +1482,7 @@ fn validate_canonical_payload_layout( member.logical_path ))); } - validate_zero_region(bytes, unaligned_offset, padding, "pacbin payload padding")?; + validate_padding(unaligned_offset, padding)?; expected_offset = end; } @@ -1339,12 +1498,7 @@ fn validate_canonical_payload_layout( "pacbin payload region has a non-canonical trailing gap", )); } - validate_zero_region( - bytes, - expected_offset, - trailing_padding, - "pacbin payload padding", - )?; + validate_padding(expected_offset, trailing_padding)?; Ok(()) } diff --git a/pyamplicol/rust/crates/rusticol-core/src/pacbin_tests.rs b/pyamplicol/rust/crates/rusticol-core/src/pacbin_tests.rs index 4afba1b5c4..3f07319f20 100644 --- a/pyamplicol/rust/crates/rusticol-core/src/pacbin_tests.rs +++ b/pyamplicol/rust/crates/rusticol-core/src/pacbin_tests.rs @@ -198,6 +198,36 @@ fn open_keeps_one_immutable_snapshot_after_atomic_source_replacement() { fs::remove_file(path).unwrap(); } +#[test] +fn file_reader_authenticates_each_member_it_reads() { + static NEXT_PATH: AtomicU64 = AtomicU64::new(0); + let path = std::env::temp_dir().join(format!( + "rusticol-pacbin-file-reader-{}-{}.pacbin", + std::process::id(), + NEXT_PATH.fetch_add(1, Ordering::Relaxed) + )); + let mut bytes = decode_hex(PYTHON_GOLDEN_HEX); + fs::write(&path, &bytes).unwrap(); + let reader = PacbinFileReader::open_file(File::open(&path).unwrap(), &path).unwrap(); + assert_eq!(reader.index().members().len(), 2); + assert!(reader.contains("a/jit.symjit") && !reader.contains("a/missing")); + assert_eq!(reader.read_member("a/jit.symjit").unwrap(), b"jit-v1"); + assert_eq!(reader.read_member("z/exact.bin").unwrap(), b"exact-state"); + + let offset = reader.member("a/jit.symjit").unwrap().offset() as usize; + bytes[offset] ^= 1; + fs::write(&path, &bytes).unwrap(); + let error = reader.read_member("a/jit.symjit").unwrap_err(); + assert_eq!(error.kind(), RusticolErrorKind::Integrity); + assert!(error.to_string().contains("member digest mismatch")); + + let truncated = &decode_hex(PYTHON_GOLDEN_HEX)[..100]; + fs::write(&path, truncated).unwrap(); + let error = PacbinFileReader::open_file(File::open(&path).unwrap(), &path).unwrap_err(); + assert_eq!(error.kind(), RusticolErrorKind::Integrity); + fs::remove_file(path).unwrap(); +} + #[test] fn expected_container_digest_authenticates_an_immutable_snapshot() { static NEXT_PATH: AtomicU64 = AtomicU64::new(0); diff --git a/pyamplicol/src/pyamplicol/artifacts/manifest.py b/pyamplicol/src/pyamplicol/artifacts/manifest.py index 6b430c636a..c8e19ada7d 100644 --- a/pyamplicol/src/pyamplicol/artifacts/manifest.py +++ b/pyamplicol/src/pyamplicol/artifacts/manifest.py @@ -70,6 +70,7 @@ "api-build-file", "sdk-metadata", "process-groups", + "payload-container", } _RUNTIME_IDENTITY_PAYLOAD_ROLES = frozenset( { diff --git a/pyamplicol/src/pyamplicol/artifacts/writer.py b/pyamplicol/src/pyamplicol/artifacts/writer.py index 65e2047768..1185fa31d1 100644 --- a/pyamplicol/src/pyamplicol/artifacts/writer.py +++ b/pyamplicol/src/pyamplicol/artifacts/writer.py @@ -371,6 +371,7 @@ def finalize( ) -> Path: if not self._payloads: raise ValueError("an artifact must contain at least one payload") + self._remove_empty_directories() manifest_extensions = dict(extensions or {}) supplied_identity = manifest_extensions.get(ARTIFACT_IDENTITY_EXTENSION) expected_identity = artifact_identity_extension() @@ -409,6 +410,15 @@ def finalize( fsync_file(path) return path + def _remove_empty_directories(self) -> None: + """Remove the directories left empty by payloads moved into containers.""" + + root = self._root() + for directory, _, _ in os.walk(root, topdown=False): + if Path(directory) != root: + with suppress(OSError): + os.rmdir(directory) + def __exit__(self, *exception: object) -> bool: if exception[0] is None and not (self._root() / MANIFEST_NAME).is_file(): raise RuntimeError("artifact builder exited without finalize()") diff --git a/pyamplicol/src/pyamplicol/cli/parser.py b/pyamplicol/src/pyamplicol/cli/parser.py index f6699b0868..e5f7bf2cc9 100644 --- a/pyamplicol/src/pyamplicol/cli/parser.py +++ b/pyamplicol/src/pyamplicol/cli/parser.py @@ -366,6 +366,12 @@ def _add_generation_options(parser: argparse.ArgumentParser) -> None: action=argparse.BooleanOptionalAction, default=argparse.SUPPRESS, ) + parser.add_argument( + "--pack-payloads", + dest="generation.pack_payloads", + action=argparse.BooleanOptionalAction, + default=argparse.SUPPRESS, + ) parser.add_argument( "--validation", dest="generation.validation.enabled", diff --git a/pyamplicol/src/pyamplicol/config/models.py b/pyamplicol/src/pyamplicol/config/models.py index 40bc9719da..5bc5cb98e8 100644 --- a/pyamplicol/src/pyamplicol/config/models.py +++ b/pyamplicol/src/pyamplicol/config/models.py @@ -655,6 +655,8 @@ class GenerationConfig: ) workers: AutoInt = field(default="auto", metadata=_setting("auto_int")) emit_api_bundle: bool = field(default=True, metadata=_setting("bool")) + # store per-process payloads and schedules in one payloads.pacbin + pack_payloads: bool = field(default=False, metadata=_setting("bool")) validation: GenerationValidationConfig = field( default_factory=GenerationValidationConfig, metadata=_section() ) @@ -673,6 +675,8 @@ def __post_init__(self) -> None: ) if not isinstance(self.emit_api_bundle, bool): raise ConfigurationError("generation.emit_api_bundle must be a boolean") + if not isinstance(self.pack_payloads, bool): + raise ConfigurationError("generation.pack_payloads must be a boolean") if not isinstance(self.validation, GenerationValidationConfig): raise ConfigurationError( "generation.validation must be a GenerationValidationConfig" diff --git a/pyamplicol/src/pyamplicol/generation/artifact_writer.py b/pyamplicol/src/pyamplicol/generation/artifact_writer.py index c81c620edc..1a6eb4c1ef 100644 --- a/pyamplicol/src/pyamplicol/generation/artifact_writer.py +++ b/pyamplicol/src/pyamplicol/generation/artifact_writer.py @@ -149,6 +149,14 @@ _EVALUATOR_PAYLOAD_CONTAINER_EXTENSION = "evaluator_payload_container" _RECURRENCE_BOOTSTRAP_FILENAME = "recurrence-bootstrap.bin" _EVALUATOR_PAYLOAD_CONTAINER_PATH = "evaluators.pacbin" +_PAYLOAD_CONTAINER_PATH = "payloads.pacbin" +_PAYLOAD_CONTAINER_ROLE = "payload-container" +# trees whose payloads generation.pack_payloads moves into payloads.pacbin +_PACKED_PAYLOAD_PREFIXES = ( + "processes/", + "recurrence/schedules/", + "model/eager-kernels/", +) _COMPILED_COLOR_CONTRACTION_MEMBER_PATH = "compiled-color.pacrclr3" _EVALUATOR_PAYLOAD_CONTAINER_KIND = "pyamplicol-evaluator-payload-container" _EVALUATOR_PAYLOAD_CONTAINER_SCHEMA_VERSION = 1 @@ -945,6 +953,11 @@ def write_schema_v3_artifact( existing = load_manifest(output) if mode == "append" else None if existing is not None: _reject_legacy_eager_append(existing) + if any(record.role == _PAYLOAD_CONTAINER_ROLE for record in existing.payloads): + raise ValueError( + "cannot append to an artifact with packed payloads; regenerate it " + "with mode='replace'" + ) hook = api_bundle_hook or _default_api_bundle_hook() requested_config = _config_payload(configuration.requested) effective_config = _config_payload(configuration.effective) @@ -1230,6 +1243,8 @@ def write_schema_v3_artifact( compiled_model=compiled_model, process_groups=group_map, ) + if _payloads_packed(configuration.effective): + _pack_payloads(builder) builder.finalize( kind=( "pyamplicol-process" @@ -1269,9 +1284,13 @@ def write_schema_v3_artifact( ) manifest = load_manifest(output) - _validate_artifact_references(manifest) + packed = _validate_artifact_references(manifest) files = tuple( - [output / record.path for record in manifest.payloads] + [ + output / record.path + for record in manifest.payloads + if record.path not in packed + ] + [output / "artifact.json"] ) return ArtifactWriteResult( @@ -5366,8 +5385,11 @@ def _call_api_bundle_hook( return "API" -def _validate_artifact_references(manifest: ArtifactManifest) -> None: +def _validate_artifact_references(manifest: ArtifactManifest) -> frozenset[str]: + """Check the payload references and files; return the packed payload paths.""" + declared = {record.path for record in manifest.payloads} + packed = _validate_payload_container(manifest) _validate_evaluator_payload_container(manifest) _validate_recurrence_schedule_sharing(manifest) required = { @@ -5400,6 +5422,99 @@ def _validate_artifact_references(manifest: ArtifactManifest) -> None: raise ValueError( "artifact contains undeclared files: " + ", ".join(sorted(undeclared)) ) + if actual & packed: + raise ValueError( + "packed payloads are also stored as files: " + + ", ".join(sorted(actual & packed)) + ) + return packed + + +def _payloads_packed(config: GenerationConfig | RunConfig) -> bool: + generation = config if isinstance(config, GenerationConfig) else config.generation + return bool(generation.pack_payloads) + + +def _is_packable_payload(record: PayloadRecord) -> bool: + path = record.path + return ( + path.startswith(_PACKED_PAYLOAD_PREFIXES) + and not record.executable + # nested containers below processes/ are opened as files by their loaders + and not (path.startswith("processes/") and path.endswith(".pacbin")) + # a loose bootstrap sidecar selects the fast recurrence loading path + and not path.endswith("/recurrence-bootstrap.bin") + ) + + +def _pack_payloads(builder: ArtifactBuilder) -> None: + """Move the packable staged payloads into payloads.pacbin. + + Their manifest records are kept, so the artifact identity does not change. + """ + + records = [ + record for record in builder.payload_records() if _is_packable_payload(record) + ] + if not records: + return + index = write_pacbin_atomic( + builder.staged_path(_PAYLOAD_CONTAINER_PATH), + ( + PacbinMemberSource( + record.path, + PacbinMemberKind.DECLARED_PAYLOAD, + builder.staged_path(record.path), + ) + for record in records + ), + ) + members = {member.logical_path: member for member in index.members} + for record in records: + member = members[record.path] + if member.length != record.size_bytes or member.sha256 != record.sha256: + raise ValueError(f"packed payload {record.path!r} changed while packing") + for record in records: + builder.staged_path(record.path).unlink() + builder.register_staged_file( + _PAYLOAD_CONTAINER_PATH, + role=_PAYLOAD_CONTAINER_ROLE, + media_type="application/octet-stream", + ) + + +def _validate_payload_container(manifest: ArtifactManifest) -> frozenset[str]: + """Check that the payload container holds exactly declared payloads.""" + + containers = [ + record for record in manifest.payloads if record.role == _PAYLOAD_CONTAINER_ROLE + ] + if not containers: + return frozenset() + if len(containers) != 1 or containers[0].path != _PAYLOAD_CONTAINER_PATH: + raise ValueError( + f"an artifact declares at most one payload container, at " + f"{_PAYLOAD_CONTAINER_PATH!r}" + ) + records = {record.path: record for record in manifest.payloads} + with PacbinReader.open( + manifest.root / _PAYLOAD_CONTAINER_PATH, verify_payloads=False + ) as reader: + members = reader.members + for member in members: + record = records.get(member.logical_path) + if ( + record is None + or member.kind is not PacbinMemberKind.DECLARED_PAYLOAD + or record.role == _PAYLOAD_CONTAINER_ROLE + or record.executable + or member.length != record.size_bytes + or member.sha256 != record.sha256 + ): + raise ValueError( + f"invalid payload container member {member.logical_path!r}" + ) + return frozenset(member.logical_path for member in members) def _validate_recurrence_schedule_sharing(manifest: ArtifactManifest) -> None: diff --git a/pyamplicol/src/pyamplicol/generation/evaluator_container.py b/pyamplicol/src/pyamplicol/generation/evaluator_container.py index e8a5a0fffe..e708b304ce 100644 --- a/pyamplicol/src/pyamplicol/generation/evaluator_container.py +++ b/pyamplicol/src/pyamplicol/generation/evaluator_container.py @@ -66,6 +66,8 @@ class PacbinMemberKind(IntEnum): RECURRENCE_COLOR_PROJECTION_CERTIFICATE = 8 ON_THE_FLY_PROCESS_SEED = 9 COLOR_CONTRACTION = 10 + # a payload declared by the artifact manifest, stored in its payload container + DECLARED_PAYLOAD = 11 @dataclass(frozen=True, slots=True) diff --git a/pyamplicol/tests/unit/test_artifacts.py b/pyamplicol/tests/unit/test_artifacts.py index ef8a7d537d..8073e7140a 100644 --- a/pyamplicol/tests/unit/test_artifacts.py +++ b/pyamplicol/tests/unit/test_artifacts.py @@ -1095,3 +1095,99 @@ def test_process_group_update_waits_for_the_generation_lock( writer.join() assert read_process_groups(root) is not None load_manifest(root, verify_payloads=True) + + +def _build_packable(path: Path, *, pack: bool) -> None: + with ArtifactBuilder(path) as builder: + builder.add_json( + "physics/process.json", + {"schema_version": 1}, + role="runtime-physics", + process_id="dd_to_z", + ) + builder.add_json( + "processes/dd_to_z/execution.json", + {"plan": [1, 2, 3]}, + role="evaluator-manifest", + process_id="dd_to_z", + ) + for relative, content in ( + ("processes/dd_to_z/recurrence-binding.bin", b"binding"), + ("processes/dd_to_z/eager-runtime.pacbin", b"nested container"), + ("recurrence/schedules/ab/recurrence-runtime.pacbin", b"schedule"), + ): + builder.add_bytes( + relative, + content, + role="evaluator-manifest", + media_type="application/octet-stream", + process_id="dd_to_z", + ) + if pack: + generation_artifact_writer._pack_payloads(builder) + builder.finalize( + kind="pyamplicol-process", + producer=_producer(), + model=_model(), + configuration=_configuration(), + processes=[_process()], + default_process_id="dd_to_z", + runtime=_runtime(), + ) + + +def test_packed_payloads_keep_records_and_artifact_identity(tmp_path: Path) -> None: + _build_packable(tmp_path / "loose", pack=False) + _build_packable(tmp_path / "packed", pack=True) + loose = load_manifest(tmp_path / "loose") + packed = load_manifest(tmp_path / "packed") + + assert packed.artifact_id == loose.artifact_id + assert {record.path for record in packed.payloads} == { + record.path for record in loose.payloads + } | {"payloads.pacbin"} + assert generation_artifact_writer._validate_payload_container(packed) == { + "processes/dd_to_z/execution.json", + "processes/dd_to_z/recurrence-binding.bin", + "recurrence/schedules/ab/recurrence-runtime.pacbin", + } + root = tmp_path / "packed" + assert not [ + path for path in root.rglob("*") if path.is_dir() and not any(path.iterdir()) + ] + assert not (root / "recurrence").exists() + assert not (root / "processes/dd_to_z/execution.json").exists() + # a nested container stays a file for the loaders that open it + assert (root / "processes/dd_to_z/eager-runtime.pacbin").is_file() + with PacbinReader.open(root / "payloads.pacbin") as reader: + assert ( + reader.read_member( + "processes/dd_to_z/recurrence-binding.bin", length=len(b"binding") + ) + == b"binding" + ) + + +def test_payload_container_must_match_its_records(tmp_path: Path) -> None: + from pyamplicol.generation.evaluator_container import ( + PacbinMemberKind, + PacbinMemberSource, + write_pacbin_atomic, + ) + + root = tmp_path / "packed" + _build_packable(root, pack=True) + container = root / "payloads.pacbin" + container.chmod(0o644) + write_pacbin_atomic( + container, + [ + PacbinMemberSource( + "processes/dd_to_z/recurrence-binding.bin", + PacbinMemberKind.DECLARED_PAYLOAD, + io.BytesIO(b"tampered"), + ) + ], + ) + with pytest.raises(ValueError, match="invalid payload container member"): + generation_artifact_writer._validate_payload_container(load_manifest(root)) diff --git a/pyamplicol/tests/unit/test_config_models.py b/pyamplicol/tests/unit/test_config_models.py index dd19a57651..6be37cefc1 100644 --- a/pyamplicol/tests/unit/test_config_models.py +++ b/pyamplicol/tests/unit/test_config_models.py @@ -40,7 +40,7 @@ def test_schema_v1_registry_contains_every_contract_leaf() -> None: - assert len(FIELD_REGISTRY) == 83 + assert len(FIELD_REGISTRY) == 84 assert "generation.umami_grouping" not in FIELD_REGISTRY assert "evaluator.jit.direct_translation" not in FIELD_REGISTRY assert FIELD_REGISTRY["action"].required From 2e0064404a70ccff9366e21ed7dbcf22d40e4400 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 26 Oct 2026 17:42:28 +0100 Subject: [PATCH 46/88] pyamplicol: spool packed payloads instead of staging them as files An ArtifactBuilder with a packing policy appends the packable payloads to one staging spool and writes payloads.pacbin from it, so they are never created as files; read_staged serves them to the staging-time readers (the structural proofs, the evaluator container and the UMAMI and recurrence bootstrap writers). The container is sealed before the recurrence bootstrap, whose staged artifact now reads it. For p p > j j j [amplicol], staging creates about 37 files per artifact instead of about 2000, with the same artifacts. --- .../rust/crates/rusticol-core/src/artifact.rs | 5 +- pyamplicol/src/pyamplicol/artifacts/writer.py | 177 +++++++++++++++++- .../pyamplicol/generation/artifact_writer.py | 82 +++----- .../generation/structural_source_proof.py | 57 ++++-- pyamplicol/tests/unit/test_artifacts.py | 23 ++- 5 files changed, 265 insertions(+), 79 deletions(-) diff --git a/pyamplicol/rust/crates/rusticol-core/src/artifact.rs b/pyamplicol/rust/crates/rusticol-core/src/artifact.rs index 75e7a2fe0c..2f773c1492 100644 --- a/pyamplicol/rust/crates/rusticol-core/src/artifact.rs +++ b/pyamplicol/rust/crates/rusticol-core/src/artifact.rs @@ -1838,14 +1838,15 @@ pub(crate) fn open_staged_recurrence_bootstrap_artifact( } } validate_evaluator_payload_container_declaration(&evaluator_payload_container, &payloads)?; + let payload_container = load_payload_container(&root, &payloads)?; Ok(VerifiedArtifact { manifest_path: root.join(ARTIFACT_MANIFEST_FILE), root, manifest, payloads: Arc::new(payloads), - // Process-ready lowering reads only loose plan/binding/color payloads. + // Process-ready lowering reads only plan/binding/color payloads. evaluator_payload_container: None, - payload_container: None, + payload_container, #[cfg(feature = "f64-compiled")] native_library_cache: Arc::new(Mutex::new(BTreeMap::new())), }) diff --git a/pyamplicol/src/pyamplicol/artifacts/writer.py b/pyamplicol/src/pyamplicol/artifacts/writer.py index 1185fa31d1..ce2b22ddeb 100644 --- a/pyamplicol/src/pyamplicol/artifacts/writer.py +++ b/pyamplicol/src/pyamplicol/artifacts/writer.py @@ -8,7 +8,7 @@ import json import os import tempfile -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from contextlib import suppress from datetime import UTC, datetime from pathlib import Path @@ -32,6 +32,11 @@ ) from .transaction import ArtifactTransaction, ArtifactWriteMode +PAYLOAD_CONTAINER_PATH = "payloads.pacbin" +PAYLOAD_CONTAINER_ROLE = "payload-container" +# packable payloads are appended here while staging, then moved into the container +_PAYLOAD_SPOOL_NAME = ".payloads.spool" + class ArtifactBuilder: """Build one artifact in a private staging directory.""" @@ -42,13 +47,24 @@ def __init__( *, mode: ArtifactWriteMode = "error", expected_artifact_id: str | None = None, + packs: Callable[[str], bool] | None = None, ) -> None: + """``packs`` selects the payloads stored in payloads.pacbin; they are + spooled into one staging file instead of being written as files.""" if expected_artifact_id is not None and mode != "append": raise ValueError("expected_artifact_id is valid only for append mode") + if packs is not None and mode == "append": + raise ValueError("packed payloads cannot be appended") self._transaction = ArtifactTransaction(destination, mode=mode) self._expected_artifact_id = expected_artifact_id self.root: Path | None = None self._payloads: dict[str, PayloadRecord] = {} + self._packs = packs + self._spool: BinaryIO | None = None + # path -> (offset, length) in the spool + self._spooled: dict[str, tuple[int, int]] = {} + self._sealed = False + self._packed: frozenset[str] = frozenset() def __enter__(self) -> ArtifactBuilder: self.root = self._transaction.__enter__() @@ -196,6 +212,16 @@ def add_stream( if not callable(read): raise TypeError("artifact stream source must provide read()") path_value = normalize_relative_path(relative) + if self._packs is not None and not executable and self._packs(path_value): + return self._spool_stream( + path_value, + source, + role=role, + media_type=media_type, + target=target, + process_id=process_id, + chunk_size=chunk_size, + ) path = self.staged_path(path_value, create_parent=True) descriptor, temporary_name = tempfile.mkstemp( prefix=f".{path.name}.", suffix=".tmp", dir=path.parent @@ -319,6 +345,9 @@ def discard_payloads(self, relative: str, *, recursive: bool = False) -> None: if path == path_value or (recursive and path.startswith(prefix)) ) for owned_path in owned: + if self._spooled.pop(owned_path, None) is not None: + self._payloads.pop(owned_path, None) + continue path = confined_path(self._root(), owned_path, must_exist=False) if path.is_symlink() or path.is_dir(): raise ValueError( @@ -343,6 +372,124 @@ def discard_payloads(self, relative: str, *, recursive: bool = False) -> None: with suppress(OSError): directory.rmdir() + def _spool_stream( + self, + path_value: str, + source: BinaryIO, + *, + role: str, + media_type: str, + target: Mapping[str, object] | None, + process_id: str | None, + chunk_size: int, + ) -> PayloadRecord: + if self._sealed: + raise RuntimeError( + f"payload {path_value!r} arrived after the payload container was sealed" + ) + if path_value == MANIFEST_NAME: + raise ValueError(f"{MANIFEST_NAME} is reserved for the artifact manifest") + if self._spool is None: + self._spool = (self._root() / _PAYLOAD_SPOOL_NAME).open("w+b") + spool = self._spool + offset = spool.seek(0, os.SEEK_END) + digest = hashlib.sha256() + while True: + chunk = source.read(chunk_size) + if not isinstance(chunk, bytes | bytearray | memoryview): + raise TypeError("artifact stream source must return bytes") + if not chunk: + break + digest.update(chunk) + spool.write(chunk) + length = spool.tell() - offset + self._spooled[path_value] = (offset, length) + return self._record_payload( + path_value, + role=role, + media_type=media_type, + size_bytes=length, + sha256=digest.hexdigest(), + executable=False, + target=target, + process_id=process_id, + ) + + def is_spooled(self, relative: str) -> bool: + """Whether a staged payload is held in the spool rather than a file.""" + + return normalize_relative_path(relative) in self._spooled + + def has_payload(self, relative: str) -> bool: + """Whether ``relative`` is a payload staged so far.""" + + return normalize_relative_path(relative) in self._payloads + + def read_staged(self, relative: str) -> bytes: + """The bytes of a staged payload, wherever it is stored.""" + + path_value = normalize_relative_path(relative) + spooled = self._spooled.get(path_value) + if spooled is not None and self._spool is not None: + offset, length = spooled + self._spool.flush() + data = os.pread(self._spool.fileno(), length, offset) + if len(data) != length: + raise OSError(f"short read of spooled payload {path_value!r}") + return data + if self._sealed and path_value in self._packed: + from pyamplicol.generation.evaluator_container import PacbinReader + + record = self._payloads[path_value] + with PacbinReader.open( + self._root() / PAYLOAD_CONTAINER_PATH, verify_payloads=False + ) as reader: + return reader.read_member(path_value, length=record.size_bytes) + return confined_path(self._root(), path_value).read_bytes() + + def seal_payload_container(self) -> PayloadRecord | None: + """Move the spooled payloads into payloads.pacbin; later packable + payloads are rejected. Their records are kept.""" + + if self._sealed: + raise RuntimeError("the payload container is already sealed") + self._sealed = True + self._packed = frozenset(self._spooled) + if self._spool is None: + return None + from pyamplicol.generation.evaluator_container import ( + PacbinMemberKind, + PacbinMemberSource, + write_pacbin_atomic, + ) + + spool = self._spool + spool.flush() + sources = [ + PacbinMemberSource( + path, + PacbinMemberKind.DECLARED_PAYLOAD, + _SpoolSlice(spool.fileno(), offset, length), + ) + for path, (offset, length) in self._spooled.items() + ] + index = write_pacbin_atomic(self.staged_path(PAYLOAD_CONTAINER_PATH), sources) + for member in index.members: + record = self._payloads[member.logical_path] + if member.length != record.size_bytes or member.sha256 != record.sha256: + raise ValueError( + f"packed payload {member.logical_path!r} changed while packing" + ) + spool.close() + self._spool = None + self._spooled = {} + (self._root() / _PAYLOAD_SPOOL_NAME).unlink() + return self.register_staged_file( + PAYLOAD_CONTAINER_PATH, + role=PAYLOAD_CONTAINER_ROLE, + media_type="application/octet-stream", + ) + def payload_records(self) -> tuple[PayloadRecord, ...]: """Return the immutable, path-ordered staged payload inventory. @@ -369,6 +516,8 @@ def finalize( default_process_id: str | None = None, extensions: Mapping[str, object] | None = None, ) -> Path: + if not self._sealed and self._packs is not None: + self.seal_payload_container() if not self._payloads: raise ValueError("an artifact must contain at least one payload") self._remove_empty_directories() @@ -420,9 +569,33 @@ def _remove_empty_directories(self) -> None: os.rmdir(directory) def __exit__(self, *exception: object) -> bool: + if self._spool is not None: + self._spool.close() + self._spool = None if exception[0] is None and not (self._root() / MANIFEST_NAME).is_file(): raise RuntimeError("artifact builder exited without finalize()") return self._transaction.__exit__(*exception) # type: ignore[arg-type] -__all__ = ["ArtifactBuilder"] +class _SpoolSlice: + """One spooled payload, read with positioned reads from the open spool.""" + + def __init__(self, descriptor: int, offset: int, length: int) -> None: + self._descriptor = descriptor + self._position = offset + self._end = offset + length + + def read(self, size: int = -1) -> bytes: + remaining = self._end - self._position + if size < 0 or size > remaining: + size = remaining + if size == 0: + return b"" + data = os.pread(self._descriptor, size, self._position) + if not data: + raise OSError("short read of a spooled payload") + self._position += len(data) + return data + + +__all__ = ["PAYLOAD_CONTAINER_PATH", "PAYLOAD_CONTAINER_ROLE", "ArtifactBuilder"] diff --git a/pyamplicol/src/pyamplicol/generation/artifact_writer.py b/pyamplicol/src/pyamplicol/generation/artifact_writer.py index 1a6eb4c1ef..3179ebb6d6 100644 --- a/pyamplicol/src/pyamplicol/generation/artifact_writer.py +++ b/pyamplicol/src/pyamplicol/generation/artifact_writer.py @@ -6,6 +6,7 @@ import hashlib import importlib import importlib.metadata +import io import json import math import os @@ -33,7 +34,7 @@ process_groups_payload, validate_process_groups, ) -from pyamplicol.artifacts.security import sha256_file +from pyamplicol.artifacts.writer import PAYLOAD_CONTAINER_PATH, PAYLOAD_CONTAINER_ROLE from pyamplicol.config import ( ConfigClamp, ConfigResolution, @@ -149,8 +150,8 @@ _EVALUATOR_PAYLOAD_CONTAINER_EXTENSION = "evaluator_payload_container" _RECURRENCE_BOOTSTRAP_FILENAME = "recurrence-bootstrap.bin" _EVALUATOR_PAYLOAD_CONTAINER_PATH = "evaluators.pacbin" -_PAYLOAD_CONTAINER_PATH = "payloads.pacbin" -_PAYLOAD_CONTAINER_ROLE = "payload-container" +_PAYLOAD_CONTAINER_PATH = PAYLOAD_CONTAINER_PATH +_PAYLOAD_CONTAINER_ROLE = PAYLOAD_CONTAINER_ROLE # trees whose payloads generation.pack_payloads moves into payloads.pacbin _PACKED_PAYLOAD_PREFIXES = ( "processes/", @@ -612,8 +613,12 @@ def _register_staged_source( relative: str, kind: PacbinMemberKind, ) -> None: - path = self._builder.staged_path(relative) - self._new_sources[relative] = PacbinMemberSource(relative, kind, path) + source: Path | BinaryIO = ( + io.BytesIO(self._builder.read_staged(relative)) + if self._builder.is_spooled(relative) + else self._builder.staged_path(relative) + ) + self._new_sources[relative] = PacbinMemberSource(relative, kind, source) self._staged_loose_paths.add(relative) def _is_discarded(self, relative: str) -> bool: @@ -859,11 +864,7 @@ def _write_umami_bundle( physics = ( process.umami_physics if process is not None and process.umami_physics is not None - else json.loads( - builder.staged_path(str(record["physics_path"])).read_text( - encoding="utf-8" - ) - ) + else json.loads(builder.read_staged(str(record["physics_path"]))) ) umami_input = UmamiProcessInput( process_id=identifier, @@ -1045,6 +1046,7 @@ def write_schema_v3_artifact( output, mode=write_mode, expected_artifact_id=(existing.artifact_id if existing is not None else None), + packs=_packs_payload if _payloads_packed(configuration.effective) else None, ) as builder: evaluator_payloads = _EvaluatorPayloadCollector( builder, @@ -1171,10 +1173,8 @@ def write_schema_v3_artifact( for process_record in process_records: process_id = str(process_record["id"]) execution_path = f"processes/{process_id}/execution.json" - execution_file = builder.staged_path(execution_path) - execution_payload = json.loads( - execution_file.read_text(encoding="utf-8") - ) + execution_bytes = builder.read_staged(execution_path) + execution_payload = json.loads(execution_bytes.decode("utf-8")) if not isinstance(execution_payload, Mapping): raise ValueError( f"execution payload for {process_id!r} is not an object" @@ -1187,11 +1187,12 @@ def write_schema_v3_artifact( source_revision=source_revision, native_build_inputs_sha256=native_build_inputs_sha256, execution_path=execution_path, - execution_sha256=sha256_file(execution_file), + execution_sha256=hashlib.sha256(execution_bytes).hexdigest(), execution=execution_payload, evaluator_container_path=evaluator_container_path, evaluator_container_index_sha256=(evaluator_container_index_sha256), cache=proof_cache, + payloads=builder, ) builder.add_json( f"processes/{process_id}/structural-source-proof.json", @@ -1200,6 +1201,9 @@ def write_schema_v3_artifact( process_id=process_id, compact=True, ) + if _payloads_packed(configuration.effective): + # the native recurrence bootstrap reads the packed payloads + builder.seal_payload_container() _write_recurrence_bootstraps( builder, processes=recurrence_processes, @@ -1243,8 +1247,6 @@ def write_schema_v3_artifact( compiled_model=compiled_model, process_groups=group_map, ) - if _payloads_packed(configuration.effective): - _pack_payloads(builder) builder.finalize( kind=( "pyamplicol-process" @@ -2167,8 +2169,8 @@ def _write_recurrence_bootstraps( separators=(",", ":"), allow_nan=False, ).encode("utf-8"), - builder.staged_path(execution_path).read_bytes(), - builder.staged_path(physics_path).read_bytes(), + builder.read_staged(execution_path), + builder.read_staged(physics_path), ) if not isinstance(image, bytes) or not image: raise RuntimeError( @@ -5435,11 +5437,11 @@ def _payloads_packed(config: GenerationConfig | RunConfig) -> bool: return bool(generation.pack_payloads) -def _is_packable_payload(record: PayloadRecord) -> bool: - path = record.path +def _packs_payload(path: str) -> bool: + """Whether generation.pack_payloads stores ``path`` in payloads.pacbin.""" + return ( path.startswith(_PACKED_PAYLOAD_PREFIXES) - and not record.executable # nested containers below processes/ are opened as files by their loaders and not (path.startswith("processes/") and path.endswith(".pacbin")) # a loose bootstrap sidecar selects the fast recurrence loading path @@ -5447,42 +5449,6 @@ def _is_packable_payload(record: PayloadRecord) -> bool: ) -def _pack_payloads(builder: ArtifactBuilder) -> None: - """Move the packable staged payloads into payloads.pacbin. - - Their manifest records are kept, so the artifact identity does not change. - """ - - records = [ - record for record in builder.payload_records() if _is_packable_payload(record) - ] - if not records: - return - index = write_pacbin_atomic( - builder.staged_path(_PAYLOAD_CONTAINER_PATH), - ( - PacbinMemberSource( - record.path, - PacbinMemberKind.DECLARED_PAYLOAD, - builder.staged_path(record.path), - ) - for record in records - ), - ) - members = {member.logical_path: member for member in index.members} - for record in records: - member = members[record.path] - if member.length != record.size_bytes or member.sha256 != record.sha256: - raise ValueError(f"packed payload {record.path!r} changed while packing") - for record in records: - builder.staged_path(record.path).unlink() - builder.register_staged_file( - _PAYLOAD_CONTAINER_PATH, - role=_PAYLOAD_CONTAINER_ROLE, - media_type="application/octet-stream", - ) - - def _validate_payload_container(manifest: ArtifactManifest) -> frozenset[str]: """Check that the payload container holds exactly declared payloads.""" diff --git a/pyamplicol/src/pyamplicol/generation/structural_source_proof.py b/pyamplicol/src/pyamplicol/generation/structural_source_proof.py index a27409830a..7c6c5fe5ce 100644 --- a/pyamplicol/src/pyamplicol/generation/structural_source_proof.py +++ b/pyamplicol/src/pyamplicol/generation/structural_source_proof.py @@ -13,7 +13,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from pathlib import Path -from typing import Any +from typing import Any, Protocol from pyamplicol.artifacts.security import sha256_file @@ -238,23 +238,45 @@ def add(record: object, path: str) -> None: return tuple(lanes) +class StagedPayloads(Protocol): + """The payloads of a staged artifact, wherever they are stored.""" + + def has_payload(self, relative: str) -> bool: ... + + def read_staged(self, relative: str) -> bytes: ... + + +class _DirectoryPayloads: + """Payloads stored as files below ``root``.""" + + def __init__(self, root: Path) -> None: + self._root = root + + def has_payload(self, relative: str) -> bool: + path = self._root / relative + return path.is_file() and not path.is_symlink() + + def read_staged(self, relative: str) -> bytes: + return (self._root / relative).read_bytes() + + def _loose_object( - root: Path, + payloads: StagedPayloads, loose: str, ) -> tuple[dict[str, object], tuple[str, ...]]: """The inventory record of a loose payload and the paths it references.""" - path = root / loose + data = payloads.read_staged(loose) record: dict[str, object] = { "path": loose, - "sha256": sha256_file(path), - "size_bytes": path.stat().st_size, + "sha256": hashlib.sha256(data).hexdigest(), + "size_bytes": len(data), } - if path.suffix != ".json": + if not loose.endswith(".json"): return record, () try: - nested = json.loads(path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError) as error: + nested = json.loads(data.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as error: raise ValueError( f"referenced lane manifest is not valid JSON: {loose}" ) from error @@ -262,14 +284,13 @@ def _loose_object( def _resolve_loose_payload( - root: Path, + payloads: StagedPayloads, process_id: str, relative: str, ) -> str | None: candidates = (relative, f"processes/{process_id}/{relative}") for candidate in candidates: - path = root / candidate - if path.is_file() and not path.is_symlink(): + if payloads.has_payload(candidate): return candidate return None @@ -286,13 +307,18 @@ def build_generation_structural_proof( evaluator_container_path: str | None, evaluator_container_index_sha256: str | None, cache: StructuralProofCache | None = None, + payloads: StagedPayloads | None = None, ) -> dict[str, object]: """Build exact source/lane evidence from the private artifact staging tree. ``cache`` shares the records of payloads common to the processes of one - staged artifact. + staged artifact. ``payloads`` reads the staged payloads, by default the + files below ``artifact_root``. """ + if payloads is None: + payloads = _DirectoryPayloads(artifact_root) + lanes = execution_lanes(execution) member_records: dict[str, dict[str, object]] = {} container_record: dict[str, object] | None = None @@ -336,7 +362,7 @@ def build_generation_structural_proof( execution_path: { "path": execution_path, "sha256": execution_sha256, - "size_bytes": (artifact_root / execution_path).stat().st_size, + "size_bytes": len(payloads.read_staged(execution_path)), } } used_members: set[str] = set() @@ -363,7 +389,7 @@ def include( loose = resolved[relative] = relative else: loose = resolved[relative] = _resolve_loose_payload( - artifact_root, process_id, relative + payloads, process_id, relative ) if loose is None: # Some manifest paths are logical roots interpreted by the @@ -375,7 +401,7 @@ def include( lane_loose.add(loose) cached = None if cache is None else cache.loose.get(loose) if cached is None: - cached = _loose_object(artifact_root, loose) + cached = _loose_object(payloads, loose) if cache is not None: cache.loose[loose] = cached record, nested_paths = cached @@ -697,6 +723,7 @@ def validate_generation_structural_proof( "ROLE", "SCHEMA", "SEMANTIC_MAP_DOMAINS", + "StagedPayloads", "StructuralProofCache", "build_generation_structural_proof", "canonical_sha256", diff --git a/pyamplicol/tests/unit/test_artifacts.py b/pyamplicol/tests/unit/test_artifacts.py index 8073e7140a..08537bdde1 100644 --- a/pyamplicol/tests/unit/test_artifacts.py +++ b/pyamplicol/tests/unit/test_artifacts.py @@ -1098,7 +1098,8 @@ def test_process_group_update_waits_for_the_generation_lock( def _build_packable(path: Path, *, pack: bool) -> None: - with ArtifactBuilder(path) as builder: + packs = generation_artifact_writer._packs_payload if pack else None + with ArtifactBuilder(path, packs=packs) as builder: builder.add_json( "physics/process.json", {"schema_version": 1}, @@ -1123,8 +1124,26 @@ def _build_packable(path: Path, *, pack: bool) -> None: media_type="application/octet-stream", process_id="dd_to_z", ) + assert builder.root is not None if pack: - generation_artifact_writer._pack_payloads(builder) + # packable payloads are spooled, never written as files + assert not (builder.root / "processes/dd_to_z/execution.json").exists() + assert not (builder.root / "recurrence").exists() + assert builder.read_staged("processes/dd_to_z/recurrence-binding.bin") == ( + b"binding" + ) + if pack: + builder.seal_payload_container() + assert builder.read_staged("processes/dd_to_z/recurrence-binding.bin") == ( + b"binding" + ) + with pytest.raises(RuntimeError, match="sealed"): + builder.add_bytes( + "processes/dd_to_z/late.bin", + b"late", + role="evaluator-manifest", + media_type="application/octet-stream", + ) builder.finalize( kind="pyamplicol-process", producer=_producer(), From 19767762e178aa4c08ffc2b3d1361e295edf9585 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 26 Oct 2026 18:21:31 +0100 Subject: [PATCH 47/88] madspace: keep the fixed channel of weighting-step matrix elements in JSON weighting_groups_to_json dropped MatrixElement::channel(), so a reloaded step expected channel_in as an input instead of its fixed channel, as the AmpliCol ratio of a colour ordering needs. --- madspace/src/driver/weighting_pipeline.cpp | 4 +++- madspace/tests/test_weighting_steps.py | 9 +++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/madspace/src/driver/weighting_pipeline.cpp b/madspace/src/driver/weighting_pipeline.cpp index 46b61dc0d8..6cda38adb4 100644 --- a/madspace/src/driver/weighting_pipeline.cpp +++ b/madspace/src/driver/weighting_pipeline.cpp @@ -95,6 +95,7 @@ json matrix_element_to_json(const MatrixElement& matrix_element) { {"sample_random_inputs", matrix_element.sample_random_inputs()}, {"me_frame", me_frame}, {"incoming_count", incoming_count}, + {"channel", matrix_element.channel()}, }; } @@ -115,7 +116,8 @@ MatrixElement matrix_element_from_json(const json& j) { j.at("diagram_count").get(), j.at("sample_random_inputs").get(), j.at("me_frame").get>(), - j.at("incoming_count").get() + j.at("incoming_count").get(), + j.at("channel").get() ); } diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py index 1967cc74f8..8a3cc78651 100644 --- a/madspace/tests/test_weighting_steps.py +++ b/madspace/tests/test_weighting_steps.py @@ -334,6 +334,15 @@ def test_ratio_matrix_element_rules(tmp_path): loaded = ms.weighting_groups_from_json(text) assert loaded[0].steps[0].returns_ratio() assert ms.weighting_groups_to_json(loaded) == text + # a fixed channel (AmpliCol colour ordering) survives the round trip + channel_me = python_matrix_element( + context, true_over_surrogate, ("matrix_element_ratio",), 3 + ) + text = ms.weighting_groups_to_json( + [ms.WeightingGroup([0], [ms.WeightingStep([channel_me])])] + ) + loaded = ms.weighting_groups_from_json(text) + assert loaded[0].steps[0].matrix_elements[0].channel() == 3 def test_invalid_groups(tmp_path): From c2183443d304bffda6d497bb5bd52557f5af316a Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 26 Oct 2026 18:21:33 +0100 Subject: [PATCH 48/88] mg7: gridpacks of [amplicol] processes Every matrix-element library is loaded through MadgraphProcess.load_matrix_element, which records the path, init path (param card or AmpliCol artifact) and parameters of each load; the gridpack replays that list, so the indices match the run. The weighting groups of the second pass are stored in data/groups.json and the lc/full artifacts are copied to SubProcesses/A*. generate_events applies --max_overweight_truncation and --freeze_max_weight_after to the final pass, and the new --first_pass_* options (defaults from [amplicol] in the grid card) to the leading-colour unweighting. Gridpacks with include_source, and the systematics of [amplicol] processes mixing coupling orders, are rejected for now. --- docs/source/gridpacks.rst | 10 +++++ .../iolibs/template_files/mg7/event_stream.py | 8 +++- .../iolibs/template_files/mg7/gridpack.py | 27 ++++++++++++- .../template_files/mg7/gridpack_setup.py | 27 +++++++++---- .../template_files/mg7/madgraph_process.py | 39 ++++++++++++++++--- .../template_files/mg7/subprocess_amplicol.py | 20 ++++++---- .../mg7/subprocess_madmatrix.py | 28 ++++++------- madgraph/various/banner.py | 4 +- tests/acceptance_tests/test_mg7_amplicol.py | 34 ++++++++++++++-- 9 files changed, 151 insertions(+), 46 deletions(-) diff --git a/docs/source/gridpacks.rst b/docs/source/gridpacks.rst index a65e7d9063..55e8571c32 100644 --- a/docs/source/gridpacks.rst +++ b/docs/source/gridpacks.rst @@ -154,6 +154,9 @@ gridpack is created. * - - ``cut_efficiency_threshold``, ``max_cut_repetitions`` - Card only. + * - ``[amplicol]`` + - ``first_pass_max_overweight_truncation``, ``first_pass_freeze_max_weight_after`` + - ``--first_pass_max_overweight_truncation``, ``--first_pass_freeze_max_weight_after`` * - ``[gridpack]`` - ``output_dir``, ``temp_output_dir`` - ``--output_dir``, ``--temp_output_dir`` @@ -167,6 +170,13 @@ the main run card, see :doc:`run_card`. ``./bin/generate_events --help`` lists a Relative ``output_dir`` and ``temp_output_dir`` paths in the card refer to the gridpack directory. Both are reset to ``""`` when the gridpack is created. +With ``[amplicol]`` processes unweighted to full colour, ``max_overweight_truncation`` and +``freeze_max_weight_after`` set the second pass, which gives the events, and the +``first_pass_`` options the leading-colour unweighting, as in the main run. The colour +accuracy and the mode are fixed when the gridpack is created. The pyamplicol artifacts +the libraries read are copied to ``SubProcesses/``, and ``include_source = true`` is not +supported for these processes. + The card names the ``cpu_mode`` that the gridpack was built with, which can differ from the ``"auto"`` of the original run card. The matrix element libraries carry this name. If you change ``cpu_mode`` or ``precision``, you have to rebuild them. diff --git a/madgraph/iolibs/template_files/mg7/event_stream.py b/madgraph/iolibs/template_files/mg7/event_stream.py index 3a24ab9c97..b4f8c71564 100644 --- a/madgraph/iolibs/template_files/mg7/event_stream.py +++ b/madgraph/iolibs/template_files/mg7/event_stream.py @@ -39,6 +39,7 @@ load_lhe_completer, load_madspace_data, load_run_card, + load_weighting_groups, make_context, resolve_seed, source_hash_matches, @@ -188,8 +189,13 @@ def pick(value, default): self.seed = resolve_seed(pick(seed, run_args.get("seed", -1))) load_context_data([self._context], [backend], madspace_data) channel_generators = load_channels([self._context], config) + # the maxima of all passes are fixed, so their settings do not matter + groups, _ = load_weighting_groups( + gen_args["max_overweight_truncation"], gen_args["freeze_max_weight_after"] + ) self._stream = ms.EventStream( - self._context, channel_generators, self.seed, load_lhe_completer(), config + self._context, channel_generators, self.seed, load_lhe_completer(), config, + groups=groups, ) def __iter__(self) -> Iterator["ms.LHEEvent"]: diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 85fb33b5ef..171e05ddb1 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -21,6 +21,7 @@ load_madspace_data, load_run_card, load_systematics, + load_weighting_groups, make_context, resolve_seed, source_hash_matches, @@ -76,6 +77,7 @@ def main() -> None: run_args = run_card["run"] gen_args = run_card["generation"] grid_args = run_card.get("gridpack", {}) + amplicol_args = run_card["amplicol"] madspace_data = load_madspace_data() # parse command line arguments @@ -123,8 +125,21 @@ def main() -> None: choices=["lhe", "lhe_npy", "compact_npy"] ) parser.add_argument("--events", type=int, default=gen_args["events"]) + # these two set the final pass, which gives the events parser.add_argument("--max_overweight_truncation", type=float, default=gen_args["max_overweight_truncation"]) parser.add_argument("--freeze_max_weight_after", type=int, default=gen_args["freeze_max_weight_after"]) + parser.add_argument( + "--first_pass_max_overweight_truncation", type=float, + default=amplicol_args["first_pass_max_overweight_truncation"], + help="max_overweight_truncation of the leading-colour unweighting of " + "[amplicol] processes unweighted to full colour" + ) + parser.add_argument( + "--first_pass_freeze_max_weight_after", type=int, + default=amplicol_args["first_pass_freeze_max_weight_after"], + help="freeze_max_weight_after of the leading-colour unweighting of " + "[amplicol] processes unweighted to full colour" + ) parser.add_argument("--cpu_batch_size", type=int, default=gen_args["cpu_batch_size"]) parser.add_argument("--gpu_batch_size", type=int, default=gen_args["gpu_batch_size"]) args = parser.parse_args() @@ -198,8 +213,15 @@ def resolve_dir(cli_value, card_value): # set up generator configuration config = ms.GeneratorConfig() config.target_count = args.events - config.max_overweight_truncation = args.max_overweight_truncation - config.freeze_max_weight_after = args.freeze_max_weight_after + groups, first_pass = load_weighting_groups( + args.max_overweight_truncation, args.freeze_max_weight_after + ) + if first_pass: + config.max_overweight_truncation = args.first_pass_max_overweight_truncation + config.freeze_max_weight_after = args.first_pass_freeze_max_weight_after + else: + config.max_overweight_truncation = args.max_overweight_truncation + config.freeze_max_weight_after = args.freeze_max_weight_after config.cpu_batch_size = args.cpu_batch_size config.gpu_batch_size = args.gpu_batch_size config.verbosity = resolve_verbosity(args.verbosity) @@ -216,6 +238,7 @@ def resolve_dir(cli_value, card_value): status_file=ms.StatusFile(os.path.join(run_path, "info.json")), config=config, seed=seed, + groups=groups, ) # scale/PDF systematics (as configured when the gridpack was made) diff --git a/madgraph/iolibs/template_files/mg7/gridpack_setup.py b/madgraph/iolibs/template_files/mg7/gridpack_setup.py index 49dfc93a8e..74976c7baa 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack_setup.py +++ b/madgraph/iolibs/template_files/mg7/gridpack_setup.py @@ -99,9 +99,6 @@ def make_context(device_name: str, cpu_mode: str, cpu_thread_pool_size: int, def load_context_data(contexts, backends, madspace_data: dict) -> None: """Load the globals and the matrix elements of the gridpack into every context.""" - # run-time matrix-element parameters (bwcutoff) of the run that made the - # gridpack; gridpacks written before they were recorded used the default - me_parameters = madspace_data.get("me_parameters", {}) for context in contexts: context.load_globals(gridpack_path("data", "globals.tar")) # custom functions come first, as in the run that made the gridpack @@ -114,12 +111,13 @@ def load_context_data(contexts, backends, madspace_data: dict) -> None: load_custom_function( ms, contexts, device_types, kind, _gridpack_relative(path), parameters ) + # in the order of the run, which gave them their indices for context, backend in zip(contexts, backends): - for me_path in madspace_data["matrix_elements"]: + for path_format, init_path, parameters in madspace_data["matrix_elements"]: context.load_matrix_element( - _gridpack_relative(me_path.format(device=backend)), - gridpack_path("Cards", "param_card.dat"), - me_parameters, + _gridpack_relative(path_format.format(device=backend)), + _gridpack_relative(init_path), + parameters, ) @@ -141,6 +139,21 @@ def load_channels(contexts, config, file_dir=None): ] +def load_weighting_groups(max_overweight_truncation, freeze_max_weight_after): + """The re- and unweighting steps after the first unweighting stored in + data/groups.json, with the unweighting settings of the final pass, and + whether a step unweights: the first unweighting is then a first pass.""" + with open(gridpack_path("data", "groups.json")) as f: + groups = json.load(f) + unweights = False + for group in groups: + for step in group["steps"]: + step["max_overweight_truncation"] = max_overweight_truncation + step["freeze_max_weight_after"] = freeze_max_weight_after + unweights |= step["mode"] == "unweight" + return ms.weighting_groups_from_json(json.dumps(groups)), unweights + + def load_lhe_completer(): return ms.LHECompleter.load(gridpack_path("data", "lhe.json")) diff --git a/madgraph/iolibs/template_files/mg7/madgraph_process.py b/madgraph/iolibs/template_files/mg7/madgraph_process.py index 2ed303eb61..ca44e13cbc 100644 --- a/madgraph/iolibs/template_files/mg7/madgraph_process.py +++ b/madgraph/iolibs/template_files/mg7/madgraph_process.py @@ -302,6 +302,8 @@ def init_context(self) -> None: self.contexts = [] # AmpliCol libraries, shared by all channels of their process self.amplicol_apis = {} + # [path format, init path, parameters] of every load_matrix_element, in order + self.loaded_matrix_elements = [] self.device_types = [] self.devices = [] self.pool_sizes = [] @@ -326,6 +328,18 @@ def init_context(self) -> None: self.pool_sizes.append(pool_size) self.contexts.append(ms.Context(device=device, thread_count=pool_size)) + def load_matrix_element(self, path_format: str, init_path: str, parameters: dict): + """Load a matrix-element library into every context, recorded for the + gridpack to load the same ones at the same indices. The init path is + the param card, or the artifact of an AmpliCol library.""" + for context, backend in zip(self.contexts, self.backends): + api = context.load_matrix_element( + path_format.format(device=backend), os.path.abspath(init_path), + parameters + ) + self.loaded_matrix_elements.append([path_format, init_path, parameters]) + return api + CUSTOM_FUNCTION_KINDS = ("cut", "scale", "pdf", "bias") def init_custom_functions(self) -> None: @@ -1367,13 +1381,16 @@ def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenera subprocess_channels[i].append(len(channel_generators)) channel_generators.append(channel.event_generator) + groups = self.weighting_groups(subprocess_channels) + # the groups of the last generator built, the one a gridpack stores + self.weighting_groups_json = ms.weighting_groups_to_json(groups) event_generator = ms.EventGenerator( contexts=self.contexts, channels=channel_generators, status_file=self.status_file, config=self.event_generator_config, seed=self.run_seed, - groups=self.weighting_groups(subprocess_channels), + groups=groups, ) unused_globals = ( set(self.contexts[0].global_names()) - event_generator.used_globals() @@ -1997,6 +2014,8 @@ def save_gridpack(self) -> None: ] with open(os.path.join(data_path, "channels.json"), "w") as f: f.write("{" + ",".join(channel_entries) + "}") + with open(os.path.join(data_path, "groups.json"), "w") as f: + f.write(self.weighting_groups_json) lib_path = os.path.join(gridpack_path, "lib") if self.run_card["gridpack"]["include_source"]: @@ -2032,10 +2051,18 @@ def save_gridpack(self) -> None: ignore=shutil.ignore_patterns("include"), ) - matrix_elements = [] - for subproc in self.subprocess_data: - me_path = subproc["me_path"] - matrix_elements.append(me_path) + # the AmpliCol artifacts are read at run time; the ratio library finds + # them next to its own, otherwise empty, reweight directory + for artifact_dir in sorted({ + os.path.dirname(meta["artifact"]) + for meta in self.subprocess_data if is_amplicol(meta) + }): + for name in ("lc", "full"): + shutil.copytree( + os.path.join(artifact_dir, name), + os.path.join(gridpack_path, artifact_dir, name), + ) + os.makedirs(os.path.join(gridpack_path, artifact_dir, "reweight")) cards_path = os.path.join(gridpack_path, "Cards") os.mkdir(cards_path) @@ -2100,7 +2127,7 @@ def save_gridpack(self) -> None: shutil.copy(os.path.join(template_dir, "event_stream.py"), bin_path) data = { - "matrix_elements": matrix_elements, + "matrix_elements": self.loaded_matrix_elements, # loaded before the matrix elements, as in the run "custom_functions": gridpack_custom_paths, "me_parameters": me_parameters(self.run_card), diff --git a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py index 28e8967a4a..c220a22455 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py +++ b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py @@ -9,7 +9,6 @@ from __future__ import annotations -import os from typing import TYPE_CHECKING from madgraph.iolibs.template_files.mg7.common import lhe_helicities, ms @@ -94,14 +93,22 @@ def check_flavor_blind_cuts(self): f"subprocess {self.meta['incoming']} > {self.meta['outgoing']}" ) - def load_matrix_elements(self, all_api_paths, api_path_formats, member_metas): + def load_matrix_elements(self, api_path_formats, member_metas): run_args = self.process.run_card["run"] if any(device != "cpu" for device in self.process.device_types): raise RuntimeError("[amplicol] processes are only supported on CPUs") if run_args["dummy_matrix_element"]: raise RuntimeError("[amplicol] processes need their matrix elements") - if self.process.run_card["gridpack"]["save_gridpack"]: - raise RuntimeError("[amplicol] processes do not support gridpacks yet") + gridpack_args = self.process.run_card["gridpack"] + if gridpack_args["save_gridpack"] and gridpack_args["include_source"]: + raise RuntimeError( + "gridpacks of [amplicol] processes cannot include the sources" + ) + if self.meta["qcd_power"] < 0 and self.process.systematics_enabled(): + raise RuntimeError( + "the scale and PDF systematics of [amplicol] processes that mix " + "coupling orders are not supported yet: disable [systematics]" + ) self.supports_alpha_s = self.meta.get("supports_alpha_s", True) if not self.supports_alpha_s and not self.process.run_card["beam"]["fixed_ren_scale"]: raise RuntimeError( @@ -125,10 +132,7 @@ def load(self, path: str, artifact: str, parameters: dict): # loaded once for all channels; each load reads the whole pyamplicol artifact apis = self.process.amplicol_apis if path not in apis: - for context in self.process.contexts: - apis[path] = context.load_matrix_element( - path, os.path.abspath(artifact), parameters - ) + apis[path] = self.process.load_matrix_element(path, artifact, parameters) return apis[path] def build_multi_channel_data(self): diff --git a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py index 7d17692695..dd1779eeb9 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py +++ b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py @@ -250,12 +250,6 @@ def __init__( else: self.unmerged_meta = unmerged_meta - # The libraries were all built up front by MadgraphProcess.compile_matrix_elements - all_api_paths = [ - [api_path_format.format(device=backend) for backend in self.process.backends] - for api_path_format in api_path_formats - ] - self.incoming_masses = [ self.process.get_mass(pid) for pid in clean_pids(self.meta["incoming"]) ] @@ -339,21 +333,21 @@ def cut_pids(observable_kwargs): **self.process.scale_kwargs ) - self.load_matrix_elements(all_api_paths, api_path_formats, member_metas) + # the libraries were all built up front by MadgraphProcess.compile_matrix_elements + self.load_matrix_elements(api_path_formats, member_metas) - def load_matrix_elements(self, all_api_paths, api_path_formats, member_metas): + def load_matrix_elements(self, api_path_formats, member_metas): """The matrix element and [surrogate] of every member subprocess.""" if self.process.run_card["run"]["dummy_matrix_element"]: - self.matrix_elements = [None] * len(all_api_paths) + self.matrix_elements = [None] * len(api_path_formats) else: - self.matrix_elements = [] - for api_paths in all_api_paths: - for context, api_path in zip(self.process.contexts, api_paths): - mat = context.load_matrix_element( - api_path, self.process.param_card_path, - me_parameters(self.process.run_card) - ) - self.matrix_elements.append(mat) + self.matrix_elements = [ + self.process.load_matrix_element( + path_format, self.process.param_card_path, + me_parameters(self.process.run_card) + ) + for path_format in api_path_formats + ] self.surrogates = [] if self.process.surrogate is not None: for path_format, member_meta in zip(api_path_formats, member_metas): diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 9800dffec9..0dcedf277d 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6957,9 +6957,9 @@ def default_setup(self): self.add_toml_param('amplicol', 'mode', "reweight", allowed=["unweight", "reweight"], comment="unweight or reweight the events to the colour accuracy") self.add_toml_param('amplicol', 'first_pass_max_overweight_truncation', 0.01, - comment="max_overweight_truncation of the leading-colour unweighting") + gridpack=True, comment="max_overweight_truncation of the leading-colour unweighting") self.add_toml_param('amplicol', 'first_pass_freeze_max_weight_after', 10000, - comment="freeze_max_weight_after of the leading-colour unweighting") + gridpack=True, comment="freeze_max_weight_after of the leading-colour unweighting") # ----------------------------- [madnis] ----------------------- self.add_toml_param('madnis', 'enable', False, allowed=["auto", True, False], auto=True, diff --git a/tests/acceptance_tests/test_mg7_amplicol.py b/tests/acceptance_tests/test_mg7_amplicol.py index e5f5559cd1..701984c6b5 100644 --- a/tests/acceptance_tests/test_mg7_amplicol.py +++ b/tests/acceptance_tests/test_mg7_amplicol.py @@ -45,8 +45,13 @@ _PDF = 'NNPDF23_lo_as_0130_qed' -def _edit_run_card(path, **amplicol): +def _edit_run_card(path, gridpack=False, **amplicol): card = open(path).read() + if gridpack: + card = re.sub(r'(?m)^save_gridpack = false', 'save_gridpack = true', card) + card = re.sub(r'(?m)^compress = true', 'compress = false', card) + # with events, for the cross section of the run + card = re.sub(r'(?m)^run_mode = ".*"', 'run_mode = "regular"', card) card = card.replace('fixed_ren_scale = false', 'fixed_ren_scale = true') card = card.replace('fixed_fact_scale = false', 'fixed_fact_scale = true') card = re.sub(r'(?m)^(pdf[12]) = ".*"$', r'\1 = "%s"' % _PDF, card) @@ -101,7 +106,7 @@ def setUp(self): def tearDown(self): shutil.rmtree(self.path, ignore_errors=True) - def run_mg7(self, name, processes, **amplicol): + def run_mg7(self, name, processes, gridpack=False, **amplicol): """Output and run the processes, returns the cross section, its error and the run directory.""" run_dir = pjoin(self.path, name) @@ -112,7 +117,7 @@ def run_mg7(self, name, processes, **amplicol): mg.exec_cmd('%s %s' % ('generate' if index == 0 else 'add process', process)) mg.exec_cmd('output mg7 %s' % run_dir) - _edit_run_card(pjoin(run_dir, 'Cards', 'run_card.toml'), **amplicol) + _edit_run_card(pjoin(run_dir, 'Cards', 'run_card.toml'), gridpack, **amplicol) log = pjoin(run_dir, 'generate_events.log') with open(log, 'w') as stream: status = subprocess.call( @@ -167,3 +172,26 @@ def test_mixed(self): mixed = self.run_mg7('mixed', ['p p > j j', 'g g > g g g [amplicol]'])[:2] self.assert_compatible( mixed, (jj[0] + ggg[0], (jj[1]**2 + ggg[1]**2)**0.5)) + + def test_gridpack(self): + """A gridpack of [amplicol] and madmatrix processes, with the second + pass unweighting to full colour, reproduces the cross section of its run.""" + mean, error, run_dir = self.run_mg7( + 'grid', ['p p > j j', 'g g > g g g [amplicol]'], gridpack=True, + mode='unweight') + gridpack = sorted(glob.glob(pjoin(run_dir, 'Events', '*', 'gridpack')))[-1] + self.assertTrue(glob.glob(pjoin(gridpack, 'SubProcesses', 'A*', 'lc'))) + log = pjoin(self.path, 'gridpack.log') + with open(log, 'w') as stream: + status = subprocess.call( + [sys.executable, pjoin(gridpack, 'bin', 'generate_events'), + '--seed', '11', '--output_dir', pjoin(self.path, 'grid_events')], + cwd=self.path, stdout=stream, stderr=subprocess.STDOUT) + self.assertEqual(status, 0, open(log).read()[-3000:]) + with open(pjoin(self.path, 'grid_events', 'info.json')) as stream: + result = json.load(stream)['process'] + self.assertTrue(result['done'], open(log).read()[-3000:]) + self.assert_compatible((float(result['mean']), float(result['error'])), + (mean, error)) + self.assertTrue(_color_balanced( + pjoin(self.path, 'grid_events', 'events.lhe.gz'))) From 7ee06f1504649e26e3be8e1b8fae5eca2df3cfae Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 26 Oct 2026 18:43:28 +0100 Subject: [PATCH 49/88] madspace: multi-channel weights from the densities of the channel mappings MappingChannelWeights weights channel c of an event of flavour f by m_fc g_c / sum_k m_fk g_k, with g_k the density of the PhaseSpaceMapping of channel k from its inverse (discrete choices sampled uniformly, adaptive samplers excluded) and m the channels that integrate the flavour. A mapping that cannot produce the event gets no weight. It also draws a diagram in proportion to the weights, for the resonances of the event record. Integrand takes it in place of PropagatorChannelWeights, so the remaps, the channel weight network and MadNIS training apply to its channels, and uses its diagram when the matrix element returns none. --- .../madspace/phasespace/channel_weights.hpp | 80 ++++++++++ .../include/madspace/phasespace/integrand.hpp | 17 ++- madspace/instruction_set.yaml | 17 +++ madspace/src/kernels/multichannel.hpp | 29 ++++ madspace/src/phasespace/channel_weights.cpp | 137 +++++++++++++++++ madspace/src/phasespace/integrand.cpp | 79 ++++++++-- madspace/src/python/madspace.cpp | 40 ++++- .../tests/test_mapping_channel_weights.py | 144 ++++++++++++++++++ 8 files changed, 530 insertions(+), 13 deletions(-) create mode 100644 madspace/tests/test_mapping_channel_weights.py diff --git a/madspace/include/madspace/phasespace/channel_weights.hpp b/madspace/include/madspace/phasespace/channel_weights.hpp index e058bad6d5..e66af441e6 100644 --- a/madspace/include/madspace/phasespace/channel_weights.hpp +++ b/madspace/include/madspace/phasespace/channel_weights.hpp @@ -1,6 +1,7 @@ #pragma once #include "madspace/phasespace/base.hpp" +#include "madspace/phasespace/phasespace.hpp" #include "madspace/phasespace/topology.hpp" #include "madspace/util.hpp" @@ -113,4 +114,83 @@ class SubchannelWeights : public FunctionGenerator { std::vector _subchannel_indices; }; +/** + * Local multi-channel weights from the sampling densities of the channel + * mappings. + * + * The weight of channel @f$c@f$ for an event of flavour @f$f@f$ is + * @f$\alpha_c(x) = m_{fc} g_c(x) / \sum_k m_{fk} g_k(x)@f$ [1, 2], where + * @f$g_k@f$ is the density with which the @ref PhaseSpaceMapping of channel + * @f$k@f$ samples the event, read off its inverse with uniformly sampled + * discrete choices, and @f$m_{fk}@f$ is 1 if channel @f$k@f$ integrates + * flavour @f$f@f$, else 0. The adaptive samplers do not enter, so the weights + * of one channel do not depend on the state of the others. A channel whose + * mapping cannot produce the event gets the weight 0. + * + * Optionally, it also samples a diagram for the event record: diagram + * @f$d@f$ belongs to channel @f$c_d@f$ and is drawn with a probability + * proportional to @f$\alpha_{c_d}@f$ among the diagrams of the flavour. + * + * `batch` is the leading batch dimension. `c` indexes the channels. + * + * **Arguments** + * - `momenta` – `float`, shape `(batch, n, 4)` – the lab-frame momenta. + * - `x1`, `x2` – `float`, shape `(batch,)` – the momentum fractions returned + * by the mapping that generated the event. + * - `flavor` – `int`, shape `(batch,)` – the flavour index. + * - `channel_index` – `int`, shape `(batch,)` – the channel that generated the + * event. + * - `mirror_index` – `int`, shape `(batch,)` – the orientation of the beams. + * Present only if the mappings are built with `mirror_beams`. + * + * **Returns** + * - `channel_weights` – `float`, shape `(batch, c)` – the normalized + * per-channel weights. + * - `diagram_index` – `int`, shape `(batch,)` – the sampled diagram. Present + * only if diagrams are given. + * + * **References** + * - [1] R. Kleiss, R. Pittau, "Weight optimization in multichannel Monte + * Carlo", https://arxiv.org/abs/hep-ph/9405257 + * - [2] R. Frederix, T. Vitos, "Leading-colour-based unweighted event + * generation for multi-parton tree-level processes", + * https://arxiv.org/abs/2409.12128 + */ +class MappingChannelWeights : public FunctionGenerator { +public: + /** + * @param mappings The mapping of each channel, without + * permutations. They take the same momenta. + * @param flavor_masks Per flavour, 1 for each channel that + * integrates it, else 0. + * @param diagram_channels The channel of each diagram. + * @param flavor_diagram_masks Per flavour, 1 for each diagram it can be + * written with, else 0. + */ + MappingChannelWeights( + const std::vector& mappings, + const nested_vector2& flavor_masks, + const std::vector& diagram_channels = {}, + const nested_vector2& flavor_diagram_masks = {} + ); + + /// Number of channels. + std::size_t channel_count() const { return _mappings.size(); } + /// Whether the beams' orientation is an argument. + bool mirror_beams() const { return _mirror_beams; } + /// Whether a diagram is sampled. + bool has_diagrams() const { return !_diagram_channels.empty(); } + +private: + NamedVector build_function_impl( + FunctionBuilder& fb, const NamedVector& args + ) const override; + + std::vector _mappings; + nested_vector2 _flavor_masks; + std::vector _diagram_channels; + nested_vector2 _flavor_diagram_masks; + bool _mirror_beams; +}; + } // namespace madspace diff --git a/madspace/include/madspace/phasespace/integrand.hpp b/madspace/include/madspace/phasespace/integrand.hpp index 9768e47f34..b3883431df 100644 --- a/madspace/include/madspace/phasespace/integrand.hpp +++ b/madspace/include/madspace/phasespace/integrand.hpp @@ -57,7 +57,8 @@ namespace madspace { * * Only `matrix_element_out` is required from the matrix element. * `diagram_amp2_out` is needed for diagram-based channel weights, i.e. with - * several channels and no @p prop_chan_weights. + * several channels and neither @p prop_chan_weights nor + * @p mapping_chan_weights. * * **References** * - [1] T. Heimel, O. Mattelaer, R. Winterhalder, "MadSpace", @@ -155,6 +156,12 @@ class Integrand : public FunctionGenerator { * the prefix @ref pdf2_prefix. * Not allowed with a @ref CustomPdf, * which receives the beam index. + * @param mapping_chan_weights Optional + * @ref MappingChannelWeights, in + * place of @p prop_chan_weights. + * Its sampled diagram is the + * diagram index of the event when + * the matrix element returns none. */ Integrand( const PhaseSpaceMapping& mapping, @@ -186,7 +193,8 @@ class Integrand : public FunctionGenerator { const std::vector& flavor_per_subproc_remap = {}, std::size_t compressed_channel_weight_count = 50, const std::optional& custom_bias = std::nullopt, - const std::optional& pdf_grid2 = std::nullopt + const std::optional& pdf_grid2 = std::nullopt, + const std::optional& mapping_chan_weights = std::nullopt ); /// Globals prefix under which the second beam's PDF grid is registered /// (PdfGrid::initialize_globals) when `pdf_grid2` is given. @@ -235,6 +243,10 @@ class Integrand : public FunctionGenerator { const std::optional& prop_chan_weights() const { return _prop_chan_weights; } + /// The mapping-density channel weights, if any. + const std::optional& mapping_chan_weights() const { + return _mapping_chan_weights; + } /// The channel-weight network, if any. const std::optional& chan_weight_net() const { return _chan_weight_net; @@ -273,6 +285,7 @@ class Integrand : public FunctionGenerator { std::optional _running_coupling; std::optional _energy_scale; std::optional _prop_chan_weights; + std::optional _mapping_chan_weights; std::optional _subchan_weights; std::optional _chan_weight_net; nested_vector2 _first_chan_weight_remap; diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index 29155d2e8e..8c3b55b552 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -2868,6 +2868,23 @@ collect_channel_weights: desc: "weight per channel" desc: "Sums squared amplitudes into per-channel weights. Building block of @ref PropagatorChannelWeights." +density_channel_weights: + inputs: + - name: densities + type: [float, n] + desc: "sampling density of each channel at the event, from its inverse mapping" + - name: mask + type: [float, n] + desc: "1 for the channels that integrate the flavor of the event, else 0" + - name: channel_index + type: [int] + desc: "channel that generated the event" + outputs: + - name: channel_weights + type: [float, n] + desc: "weight per channel" + desc: "Normalised channel weights proportional to the masked channel densities, all of the weight on the generating channel if none is positive. Building block of @ref MappingChannelWeights." + interpolate_pdf: inputs: - name: x diff --git a/madspace/src/kernels/multichannel.hpp b/madspace/src/kernels/multichannel.hpp index fd6c10e2f5..cb3c22f2b2 100644 --- a/madspace/src/kernels/multichannel.hpp +++ b/madspace/src/kernels/multichannel.hpp @@ -44,6 +44,35 @@ KERNELSPEC void kernel_collect_channel_weights( } } +template +KERNELSPEC void kernel_density_channel_weights( + FIn densities, + FIn mask, + IIn channel_index, + FOut channel_weights +) { + // A density that is not a finite positive number (an inverse mapping that + // failed) counts as zero: that channel cannot have produced the event. + FVal norm(0.); + for (std::size_t i = 0; i < channel_weights.size(); ++i) { + auto density = densities[i]; + auto valid = (density > 0.) & (density < 1e300) & (mask[i] > 0.); + channel_weights[i] = where(valid, density, FVal(0.)); + norm = norm + channel_weights[i]; + } + // the generating channel always has a positive density; if its inverse + // failed numerically, it keeps the whole weight + auto degenerate = norm == 0.; + auto safe_norm = where(degenerate, FVal(1.), norm); + for (std::size_t i = 0; i < channel_weights.size(); ++i) { + channel_weights[i] = where( + degenerate, + where(i == channel_index, FVal(1.), 0.), + channel_weights[i] / safe_norm + ); + } +} + template KERNELSPEC void kernel_sde2_channel_weights( FIn invariants, diff --git a/madspace/src/phasespace/channel_weights.cpp b/madspace/src/phasespace/channel_weights.cpp index dff2cdfe1a..0907c4614f 100644 --- a/madspace/src/phasespace/channel_weights.cpp +++ b/madspace/src/phasespace/channel_weights.cpp @@ -2,6 +2,9 @@ #include "madspace/util.hpp" +#include +#include + using namespace madspace; PropagatorChannelWeights::PropagatorChannelWeights( @@ -211,3 +214,137 @@ NamedVector SubchannelWeights::build_function_impl( ); return {{"channel_weights_out", channel_weights_out}}; } + +MappingChannelWeights::MappingChannelWeights( + const std::vector& mappings, + const nested_vector2& flavor_masks, + const std::vector& diagram_channels, + const nested_vector2& flavor_diagram_masks +) : + FunctionGenerator( + "MappingChannelWeights", + [&] { + if (mappings.empty()) { + throw std::invalid_argument("MappingChannelWeights needs a channel"); + } + NamedVector args{ + {"momenta", batch_four_vec_array(mappings.at(0).particle_count())}, + {"x1", batch_float}, + {"x2", batch_float}, + {"flavor", batch_int}, + {"channel_index", batch_int}, + }; + if (mappings.at(0).mirror_beams()) { + args.push_back("mirror_index", batch_int); + } + return args; + }(), + [&] { + NamedVector ret{ + {"channel_weights", batch_float_array(mappings.size())} + }; + if (!diagram_channels.empty()) { + ret.push_back("diagram_index", batch_int); + } + return ret; + }() + ), + _mappings(mappings), + _flavor_masks(flavor_masks), + _diagram_channels(diagram_channels), + _flavor_diagram_masks(flavor_diagram_masks), + _mirror_beams(mappings.at(0).mirror_beams()) { + for (auto& mapping : mappings) { + if (mapping.channel_count() > 1) { + throw std::invalid_argument( + "MappingChannelWeights: the mappings cannot have permutations" + ); + } + if (mapping.particle_count() != mappings.at(0).particle_count() || + mapping.mirror_beams() != _mirror_beams) { + throw std::invalid_argument( + "MappingChannelWeights: the mappings must take the same momenta" + ); + } + } + for (auto& mask : flavor_masks) { + if (mask.size() != mappings.size()) { + throw std::invalid_argument( + "MappingChannelWeights: one flavor mask entry per channel expected" + ); + } + } + if (flavor_masks.empty()) { + throw std::invalid_argument("MappingChannelWeights needs a flavor"); + } + for (auto channel : diagram_channels) { + if (channel < 0 || std::cmp_greater_equal(channel, mappings.size())) { + throw std::invalid_argument( + "MappingChannelWeights: diagram channel out of range" + ); + } + } + if (!diagram_channels.empty()) { + if (flavor_diagram_masks.size() != flavor_masks.size()) { + throw std::invalid_argument( + "MappingChannelWeights: one diagram mask per flavor expected" + ); + } + for (auto& mask : flavor_diagram_masks) { + if (mask.size() != diagram_channels.size()) { + throw std::invalid_argument( + "MappingChannelWeights: one diagram mask entry per diagram " + "expected" + ); + } + } + } +} + +NamedVector MappingChannelWeights::build_function_impl( + FunctionBuilder& fb, const NamedVector& args +) const { + ValueVec conditions; + if (_mirror_beams) { + conditions.push_back(args.at("mirror_index")); + } + ValueVec densities; + for (auto& mapping : _mappings) { + auto inverse = mapping.build_inverse( + fb, {args.at("momenta"), args.at("x1"), args.at("x2")}, conditions + ); + // the discrete choices are sampled uniformly, each with probability 1/2 + Value density = fb.mul( + inverse.at("det"), + std::pow(0.5, static_cast(mapping.discrete_dim())) + ); + // the event is out of reach of the mapping if a random number is outside + // the unit hypercube + if (mapping.random_dim() > 0) { + density = fb.mul(density, fb.cut_all(inverse.at("random"), 0., 1.)); + } + densities.push_back(density); + } + Value flavor = args.at("flavor"); + Value channel_weights = fb.density_channel_weights( + fb.stack(densities), + fb.gather_vector(flavor, _flavor_masks), + args.at("channel_index") + ); + NamedVector ret{{"channel_weights", channel_weights}}; + if (!_diagram_channels.empty()) { + Value probs = fb.mul( + fb.select(channel_weights, _diagram_channels), + fb.gather_vector(flavor, _flavor_diagram_masks) + ); + Value random = + fb + .unstack(fb.random( + fb.batch_size({args.at("momenta")}), static_cast(1) + )) + .at(0); + auto [diagram, det] = fb.sample_discrete_probs(random, probs); + ret.push_back("diagram_index", diagram); + } + return ret; +} diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index 12188a6e12..0e55b2c1aa 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -17,12 +17,15 @@ std::size_t final_channel_count( const std::vector& second_chan_weight_remap, std::size_t second_remapped_chan_count, const std::optional& prop_chan_weights, - const std::optional& subchan_weights + const std::optional& subchan_weights, + const std::optional& mapping_chan_weights ) { if (second_chan_weight_remap.size() > 0) { return second_remapped_chan_count; } else if (subchan_weights) { return subchan_weights->channel_count(); + } else if (mapping_chan_weights && first_chan_weight_remap.size() == 0) { + return mapping_chan_weights->channel_count(); } else if (first_chan_weight_remap.size() > 0 || prop_chan_weights) { // with the denominators sde_strategy, prop_chan_weights already // produces first_remapped_chan_count weights without needing a remap @@ -76,7 +79,8 @@ Integrand::Integrand( const std::vector& flavor_per_subproc_remap, std::size_t compressed_channel_weight_count, const std::optional& custom_bias, - const std::optional& pdf_grid2 + const std::optional& pdf_grid2, + const std::optional& mapping_chan_weights ) : FunctionGenerator( "Integrand", @@ -104,7 +108,8 @@ Integrand::Integrand( second_chan_weight_remap, second_remapped_chan_count, prop_chan_weights, - subchan_weights + subchan_weights, + mapping_chan_weights ); if (madnis_channel_count > 1 && madnis_channel_count * 8 > compressed_channel_weight_count * 12) { @@ -196,6 +201,7 @@ Integrand::Integrand( _running_coupling(running_coupling), _energy_scale(energy_scale), _prop_chan_weights(prop_chan_weights), + _mapping_chan_weights(mapping_chan_weights), _subchan_weights(subchan_weights), _chan_weight_net(chan_weight_net), _first_chan_weight_remap(first_chan_weight_remap), @@ -332,9 +338,10 @@ Integrand::Integrand( second_chan_weight_remap, second_remapped_chan_count, prop_chan_weights, - subchan_weights + subchan_weights, + mapping_chan_weights ); - if (channel_count > 1 && !prop_chan_weights && + if (channel_count > 1 && !prop_chan_weights && !mapping_chan_weights && me_output_index(me_first, MatrixElement::diagram_amp2_out) == -1) { throw std::invalid_argument( "diagram-based channel weights need the diagram_amp2 output of the " @@ -348,6 +355,19 @@ Integrand::Integrand( "mirrored flavors need a PhaseSpaceMapping built with mirror_beams = true" ); } + if (mapping_chan_weights) { + if (prop_chan_weights) { + throw std::invalid_argument( + "propagator and mapping channel weights cannot be combined" + ); + } + if (mapping_chan_weights->mirror_beams() != mapping.mirror_beams()) { + throw std::invalid_argument( + "the channel weight mappings must have the beam orientation of the " + "integrand's mapping" + ); + } + } _channel_part_ret_types = compute_channel_part_ret_types(); } @@ -411,6 +431,9 @@ NamedVector Integrand::compute_channel_part_ret_types() const { ret.push_back("x1_acc", acc_float); ret.push_back("x2_acc", acc_float); ret.push_back("flavor_id", acc_int); + if (_mapping_chan_weights && _mapping.mirror_beams()) { + ret.push_back("mirror_index_acc", acc_int); + } ret.push_back("weight_after_cuts", acc_float); if (_madnis_training && !std::holds_alternative(_discrete_flavor)) { ret.push_back("extra_weight_after_cuts", acc_float); @@ -729,6 +752,9 @@ NamedVector Integrand::build_channel_part( out.push_back("x1_acc", x_acc.at(0)); out.push_back("x2_acc", x_acc.at(1)); out.push_back("flavor_id", flavor_id); + if (_mapping_chan_weights && _mapping.mirror_beams()) { + out.push_back("mirror_index_acc", fb.batch_gather(indices_acc, mirror_index)); + } out.push_back("weight_after_cuts", weight_after_cuts); if (_madnis_training && extra_weight_after_cuts) { out.push_back("extra_weight_after_cuts", extra_weight_after_cuts); @@ -790,12 +816,38 @@ NamedVector Integrand::build_common_part( _second_chan_weight_remap, _second_remapped_chan_count, _prop_chan_weights, - _subchan_weights + _subchan_weights, + _mapping_chan_weights ); bool use_compressed_channel_weights = channel_count > 1 && channel_count * 8 > _compressed_channel_weight_count * 12; Value chan_weights_acc; - if (channel_count > 1 && _prop_chan_weights) { + Value mapping_diagram_acc; + if (_mapping_chan_weights && + (channel_count > 1 || _mapping_chan_weights->has_diagrams())) { + ValueVec cw_args{ + momenta_acc, + x1_acc, + x2_acc, + flavor_id, + fb.batch_gather(indices_acc, args.at("chan_index")) + }; + if (_mapping_chan_weights->mirror_beams()) { + cw_args.push_back(args.at("mirror_index_acc")); + } + auto cw_result = _mapping_chan_weights->build_function(fb, cw_args); + chan_weights_acc = cw_result.at(0); + if (_mapping_chan_weights->has_diagrams()) { + mapping_diagram_acc = cw_result.at(1); + } + if (_first_chan_weight_remap.size() > 0) { + chan_weights_acc = fb.collect_channel_weights( + chan_weights_acc, + _first_chan_weight_remap.at(0), + _first_remapped_chan_count + ); + } + } else if (channel_count > 1 && _prop_chan_weights) { chan_weights_acc = _prop_chan_weights->build_function(fb, {momenta_acc}).at(0); if (_first_chan_weight_remap.size() > 0) { chan_weights_acc = fb.collect_channel_weights( @@ -840,7 +892,7 @@ NamedVector Integrand::build_common_part( if (_diff_xs.size() == 1) { dxs_vec = _diff_xs.at(0).build_function(fb, xs_args).values(); ps_flavor_id = flavor_id; - if (channel_count > 1 && !_prop_chan_weights) { + if (channel_count > 1 && !_prop_chan_weights && !_mapping_chan_weights) { chan_weights_acc = dxs_vec.at(amp2_index); if (_first_chan_weight_remap.size() > 0) { chan_weights_acc = fb.collect_channel_weights( @@ -869,7 +921,7 @@ NamedVector Integrand::build_common_part( split_out.push_back(out); split_out.push_back(indices); } - if (channel_count > 1 && !_prop_chan_weights) { + if (channel_count > 1 && !_prop_chan_weights && !_mapping_chan_weights) { Value split_cw = outputs.at(amp2_index); if (_first_chan_weight_remap.size() > 0) { split_cw = fb.collect_channel_weights( @@ -891,7 +943,7 @@ NamedVector Integrand::build_common_part( } ++i; } - if (channel_count > 1 && !_prop_chan_weights) { + if (channel_count > 1 && !_prop_chan_weights && !_mapping_chan_weights) { chan_weights_acc = fb.batch_merge_by_index(split_channel_weights); } } @@ -1060,6 +1112,13 @@ NamedVector Integrand::build_common_part( std::pair{"helicity_index", MatrixElement::helicity_index_out}, std::pair{"diagram_index", MatrixElement::diagram_index_out}}) { int index = me_output_index(me_first, me_output); + if (index == -1 && me_output == MatrixElement::diagram_index_out && + mapping_diagram_acc) { + outputs.push_back( + name, scatter_or_drop(zeros_int, mapping_diagram_acc) + ); + continue; + } outputs.push_back( name, index == -1 diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index c93b6d3a13..580e4e9c4e 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1833,6 +1833,37 @@ PYBIND11_MODULE(_madspace_py, m) { pydoc::doc("SubchannelWeights::channel_count") ); + py::classh( + m, "MappingChannelWeights", pydoc::doc("MappingChannelWeights") + ) + .def( + py::init< + const std::vector&, + const nested_vector2&, + const std::vector&, + const nested_vector2&>(), + py::arg("mappings"), + py::arg("flavor_masks"), + py::arg("diagram_channels") = std::vector{}, + py::arg("flavor_diagram_masks") = nested_vector2{}, + pydoc::doc("MappingChannelWeights::MappingChannelWeights") + ) + .def( + "channel_count", + &MappingChannelWeights::channel_count, + pydoc::doc("MappingChannelWeights::channel_count") + ) + .def( + "mirror_beams", + &MappingChannelWeights::mirror_beams, + pydoc::doc("MappingChannelWeights::mirror_beams") + ) + .def( + "has_diagrams", + &MappingChannelWeights::has_diagrams, + pydoc::doc("MappingChannelWeights::has_diagrams") + ); + py::classh( m, "MomentumPreprocessing", pydoc::doc("MomentumPreprocessing") ) @@ -2375,7 +2406,8 @@ PYBIND11_MODULE(_madspace_py, m) { const std::vector&, std::size_t, const std::optional&, - const std::optional&>(), + const std::optional&, + const std::optional&>(), py::arg("mapping"), py::arg("diff_xs"), py::arg("adaptive_map") = std::monostate{}, @@ -2406,6 +2438,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("compressed_channel_weight_count") = 50, py::arg("custom_bias") = std::nullopt, py::arg("pdf_grid2") = std::nullopt, + py::arg("mapping_chan_weights") = std::nullopt, pydoc::doc("Integrand::Integrand") ) .def( @@ -2455,6 +2488,11 @@ PYBIND11_MODULE(_madspace_py, m) { &Integrand::energy_scale, pydoc::doc("Integrand::energy_scale") ) + .def( + "mapping_chan_weights", + &Integrand::mapping_chan_weights, + pydoc::doc("Integrand::mapping_chan_weights") + ) .def( "prop_chan_weights", &Integrand::prop_chan_weights, diff --git a/madspace/tests/test_mapping_channel_weights.py b/madspace/tests/test_mapping_channel_weights.py new file mode 100644 index 0000000000..e233ed1b7f --- /dev/null +++ b/madspace/tests/test_mapping_channel_weights.py @@ -0,0 +1,144 @@ +"""MappingChannelWeights: channel weights from the densities of the channel mappings. + +- The weights are the masked densities of the inverse mappings, normalised per + event, and the sampled diagram follows them. +- A color-ordered chain with a resonance among its children integrates to the + same as the plain chain, and can be inverted at the events of the plain one. +""" + +import pytest +import torch +from madspace_testing import default_context, generator, rand, randint + +import madspace as ms + +E_CM = 13000.0 +M_Z, W_Z = 91.1876, 2.4952 +PIDS = [2, -2, 21, 11, -11] + + +def cuts(): + O = ms.Observable + return ms.Cuts( + [ + ms.CutItem(O(PIDS, O.obs_pt, [O.jet_pids]), min=20.0), + ms.CutItem(O(PIDS, O.obs_pt, [[11, -11]]), min=10.0), + ] + ) + + +def chain(order, resonant): + """q q~ > g e- e+ as a chain, the lepton pair from a Z if resonant.""" + if resonant: + propagators = [ms.Propagator(), ms.Propagator(mass=M_Z, width=W_Z, pdg_id=23)] + vertices = [["i0", "p0", "o0"], ["i1", "p0", "p1"], ["p1", "o1", "o2"]] + else: + propagators = [ms.Propagator(), ms.Propagator()] + vertices = [["i0", "p0", "o0"], ["p0", "p1", "o1"], ["i1", "p1", "o2"]] + diagram = ms.Diagram([0.0, 0.0], [0.0] * 3, propagators, vertices) + return ms.PhaseSpaceMapping( + ms.Topology(diagram), + E_CM, + t_channel_mode=ms.PhaseSpaceMapping.color_ordered, + cuts=cuts(), + permutations=[list(range(5))], + color_ordered=ms.ColorOrderedMapping.Options(order), + ) + + +RESONANT = chain([0, 2, 3, 1], True) +PLAIN = chain([0, 2, 3, 4, 1], False) + + +def sample(mapping, n, seed): + g = generator(seed) + inputs = [rand((n, mapping.random_dim()), generator=g)] + if mapping.discrete_dim(): + inputs.append( + randint(0, 2, (n, mapping.discrete_dim()), generator=g, dtype=torch.int32) + ) + p, x1, x2, det = mapping.map_forward(inputs) + return p, x1, x2, det * 2 ** mapping.discrete_dim() + + +def density(mapping, p, x1, x2): + return mapping.map_inverse([p, x1, x2])[-1] / 2 ** mapping.discrete_dim() + + +def test_resonant_chain_integrates_like_the_plain_one(): + def bw(p): + q = p[:, 3] + p[:, 4] + m2 = q[:, 0] ** 2 - torch.sum(q[:, 1:] ** 2, dim=1) + return 1e6 / ((m2 - M_Z**2) ** 2 + (M_Z * W_Z) ** 2) + + n = 400_000 + results = [] + for mapping in (RESONANT, PLAIN): + p, _, _, w = sample(mapping, n, 1) + values = torch.where(w > 0, w * torch.nan_to_num(bw(p)), 0.0) + results.append((values.mean(), values.std() / n**0.5)) + (res, res_err), (plain, plain_err) = results + assert res_err < 0.02 * res + assert abs(res - plain) < 3 * (res_err**2 + plain_err**2) ** 0.5 + + +def test_resonant_chain_inverts_the_plain_events(): + p, x1, x2, w = sample(PLAIN, 100_000, 2) + ok = w > 0 + out = RESONANT.map_inverse([p[ok], x1[ok], x2[ok]]) + r = out[0] + assert torch.all((r >= 0.0) & (r <= 1.0)) + p_back, *_ = RESONANT.map_forward([r]) + assert torch.max(torch.abs(p_back - p[ok])) < 1e-8 * E_CM + + +@pytest.mark.parametrize("flavor", [0, 1]) +def test_weights_are_the_normalised_masked_densities(flavor): + masks = [[1.0, 1.0], [0.0, 1.0]] + # diagrams 0 and 1 of flavor 0 on channels 0 and 1, diagram 2 of flavor 1 + weights = ms.MappingChannelWeights( + [RESONANT, PLAIN], masks, [0, 1, 1], [[1.0, 1.0, 0.0], [0.0, 0.0, 1.0]] + ) + runtime = ms.FunctionRuntime(weights.function(), default_context()) + p, x1, x2, w = sample(PLAIN, 40_000, 3) + ok = w > 0 + p, x1, x2 = p[ok], x1[ok], x2[ok] + n = p.shape[0] + flavors = torch.full((n,), flavor, dtype=torch.int32) + channels = torch.full((n,), 1, dtype=torch.int32) + alpha, diagram = ( + torch.from_dlpack(t) for t in runtime.call([p, x1, x2, flavors, channels]) + ) + expected = torch.stack( + [density(RESONANT, p, x1, x2), density(PLAIN, p, x1, x2)], dim=1 + ) * torch.tensor(masks[flavor]) + expected /= expected.sum(dim=1, keepdim=True) + torch.testing.assert_close(alpha, expected, rtol=1e-12, atol=1e-14) + counts = torch.bincount(diagram.long(), minlength=3).double() + if flavor == 1: + assert counts.tolist() == [0.0, 0.0, n] + else: + assert counts[2] == 0 + mean = alpha[:, 0].sum() + assert abs(counts[0] - mean) < 5 * mean**0.5 + + +def test_channels_without_the_flavor_get_no_weight(): + weights = ms.MappingChannelWeights([PLAIN, PLAIN], [[0.0, 1.0]]) + runtime = ms.FunctionRuntime(weights.function(), default_context()) + p, x1, x2, w = sample(PLAIN, 1000, 4) + ok = w > 0 + n = int(ok.sum()) + alpha = torch.from_dlpack( + runtime.call( + [ + p[ok], + x1[ok], + x2[ok], + torch.zeros(n, dtype=torch.int32), + torch.ones(n, dtype=torch.int32), + ] + )[0] + ) + torch.testing.assert_close(alpha[:, 1], torch.ones(n, dtype=alpha.dtype)) + torch.testing.assert_close(alpha[:, 0], torch.zeros(n, dtype=alpha.dtype)) From 3333a6663e4ba0dc5b21d882221e6d230867b0f2 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 26 Oct 2026 18:43:36 +0100 Subject: [PATCH 50/88] pyamplicol: UMAMI phase-space channels with colour-singlet partners and resonances The resolved physics lists the unstable s-channel propagators a process can put on shell (extensions.s_channel_resonances), from a species-level sweep over the final-state leg subsets within the process's coupling orders. Every LC colour order is integrated over several phase-space channels: its colour singlets after the antiquark ends of the lines, in every order and distribution as in AmpliCol, the non-cyclic orders of three or more lines, and every configuration of mutually nested or disjoint resonances, the non-resonant one included. Channels are shared only by contributions with the same kind of particle in every leg. A UMAMI channel collects the selectors linked by shared phase-space channels (phasespace_channels, with its selectors' multichannels indexing them); schema_version 2. A coloured resonance that can go on shell, or one decaying to coloured and colourless particles, is an error for now. --- pyamplicol/docs/user/umami-api.md | 16 +- pyamplicol/src/pyamplicol/artifacts/umami.py | 282 +++++++++++++++++- .../generation/s_channel_resonances.py | 105 +++++++ .../src/pyamplicol/generation/service.py | 7 + .../unit/test_umami_phasespace_partners.py | 101 +++++++ pyamplicol/umami_example/test_integrate.py | 10 +- 6 files changed, 499 insertions(+), 22 deletions(-) create mode 100644 pyamplicol/src/pyamplicol/generation/s_channel_resonances.py create mode 100644 pyamplicol/tests/unit/test_umami_phasespace_partners.py diff --git a/pyamplicol/docs/user/umami-api.md b/pyamplicol/docs/user/umami-api.md index bd9297ac3e..898c79326e 100644 --- a/pyamplicol/docs/user/umami-api.md +++ b/pyamplicol/docs/user/umami-api.md @@ -209,6 +209,18 @@ A provider's metadata contains: optional αs parameter. - `channels[].processes`: the channel/flavour selectors consumed by UMAMI, with physical PDGs and integration weights referenced by `matrix_elements`. +- `channels[].phasespace_channels`: the colour-ordered phase-space channels + integrating the selectors of a channel, each a cyclic `phasespace_order` of + the legs and the s-channel `resonances` (`pdg`, decay `legs`) it maps out. + A selector lists its phase-space channels in `multichannels`; they differ in + where its colour singlets attach to the colour lines (after the antiquark + ends, in every order and distribution, as in AmpliCol), in the non-cyclic + order of three or more lines, and in which resonances of the process they + follow, with a non-resonant one always included. A UMAMI channel collects + the selectors linked by shared phase-space channels. Weight the channels of a + selector by a partition of unity, for example their normalised mapping + densities, rather than adding them. A coloured resonance that can go on shell + is rejected for now. - `pdg_ids`, `color_orders`, `color_flows` and `helicities`: shared physical configuration tables. Indices are zero-based and colour-flow references can be absent for contracted colour. @@ -219,8 +231,8 @@ A provider's metadata contains: Physical incoming PDGs remain oriented for PDF evaluation. Aliases name the same computation; they are not automatically extra integrable contributions. -Each exported contribution uses one phase-space ordering, rather than several -alternative maps whose naive sum would double count it. +An exported contribution integrated over several phase-space channels must be +weighted between them, since their naive sum would double count it. ## Automatic symmetry grouping diff --git a/pyamplicol/src/pyamplicol/artifacts/umami.py b/pyamplicol/src/pyamplicol/artifacts/umami.py index b609e38fbc..6501e701f3 100644 --- a/pyamplicol/src/pyamplicol/artifacts/umami.py +++ b/pyamplicol/src/pyamplicol/artifacts/umami.py @@ -17,6 +17,7 @@ from collections.abc import Mapping, Sequence from contextlib import suppress from dataclasses import dataclass +from itertools import combinations_with_replacement, permutations from typing import TYPE_CHECKING, Any from .api_bundle import ApiBundlePayload @@ -347,6 +348,162 @@ def _relabelled_runtime(runtime, relabel): } +# Phase-space partners of one colour order beyond this many are dropped: any +# subset gives correct channel weights. Every resonance configuration keeps one. +_MAX_PARTNERS = 64 +# More resonance candidates than this in one process are not enumerated. +_MAX_RESONANCES = 12 + + +def _roles(pdgs, incoming, particles): + """The colour representation of each leg in the all-outgoing convention.""" + roles = [] + for slot, pdg in enumerate(pdgs): + representation = int(particles[int(pdg)]["color"]) + if slot < incoming and abs(representation) == 3: + representation = -representation + roles.append(representation) + return roles + + +def _line_blocks(word_slots, roles): + """The complete open colour lines of an LC word in its order, or None for a + trace or a word without fundamental lines.""" + blocks, current = [], None + for slot in word_slots: + if roles[slot] == 3: + current = [slot] + elif current is not None: + current.append(slot) + if roles[slot] == -3: + blocks.append(current) + current = None + if not blocks or current is not None or sum(map(len, blocks)) != len(word_slots): + return None + return blocks + + +def _configurations(resonances): + """The maximal sets of resonances that can be mapped together, pairwise + nested or disjoint, and the empty set for the non-resonant region.""" + if len(resonances) > _MAX_RESONANCES: + raise ValueError( + f"{len(resonances)} candidate s-channel resonances are too many to " + "enumerate their phase-space configurations" + ) + + def compatible(a, b): + return a < b or b < a or not a & b + + families = [ + bits + for bits in range(1, 1 << len(resonances)) + if all( + compatible(resonances[i][1], resonances[j][1]) + for i in range(len(resonances)) + for j in range(i + 1, len(resonances)) + if bits >> i & 1 and bits >> j & 1 + ) + ] + maximal = [ + bits + for bits in families + if not any(other != bits and other & bits == bits for other in families) + ] + return [()] + [ + tuple(r for i, r in enumerate(resonances) if bits >> i & 1) for bits in maximal + ] + + +def _resonance_legs(resonance, configuration): + """The legs of a resonance with its nested resonances kept contiguous.""" + inner = [ + other + for other in configuration + if other[1] < resonance[1] + and not any(other[1] < r[1] < resonance[1] for r in configuration) + ] + covered = set().union(*(other[1] for other in inner)) + blocks = [_resonance_legs(other, configuration) for other in inner] + blocks += [[slot] for slot in resonance[1] - covered] + return [slot for block in sorted(blocks, key=min) for slot in block] + + +def _placements(word_slots, blocks, units): + """Phase-space orders of one colour order: as in AmpliCol, the colour + singlet units follow the antiquark ends of the lines, in every order and + every distribution over the lines, and with three or more lines every + non-cyclic order of the lines is a partner too (the same colour tensor).""" + if blocks is None: + return [list(word_slots) + [slot for unit in units for slot in unit]] + line_orders = [blocks] + ( + [[blocks[0], *rest] for rest in permutations(blocks[1:])][1:] + if len(blocks) > 2 + else [] + ) + orders = [] + for lines in line_orders: + for unit_order in permutations(units): + for cuts in combinations_with_replacement( + range(len(units) + 1), len(lines) - 1 + ): + bounds = (0, *cuts, len(units)) + order = [] + for index, line in enumerate(lines): + order += line + for unit in unit_order[bounds[index] : bounds[index + 1]]: + order += unit + orders.append(order) + return orders + + +def _phasespace_partners(word_slots, singlet_slots, roles, resonances, process): + """The phase-space orders and resonance configurations integrating one + colour order together, the first one its own: ``[(order, configuration)]`` + with ``configuration`` a tuple of ``(pdg, legs)``.""" + blocks = _line_blocks(word_slots, roles) if word_slots else None + applicable = [] + for pdg, legs, color in resonances: + if color != 1: + raise ValueError( + f"{process}: the coloured s-channel resonance {pdg} can go on shell, " + "which [amplicol] phase-space mappings do not follow yet; generate " + "the process with on-shell decays or with the madmatrix backend" + ) + coloured = {slot for slot in legs if roles[slot] != 1} + if not coloured: + applicable.append((pdg, legs)) + elif coloured != legs: + raise ValueError( + f"{process}: the s-channel resonance {pdg} decays to coloured and " + "colourless particles, which [amplicol] phase-space mappings do not " + "follow yet" + ) + elif blocks is not None and any(set(block) == legs for block in blocks): + # a colour singlet decaying to a complete colour line of this order + applicable.append((pdg, legs)) + configurations = _configurations(applicable) + per_configuration = max(1, _MAX_PARTNERS // len(configurations)) + partners = [] + for configuration in configurations: + top = [ + r + for r in configuration + if not any(r[1] < other[1] for other in configuration) + ] + units = [ + _resonance_legs(r, configuration) + for r in top + if all(roles[slot] == 1 for slot in r[1]) + ] + covered = {slot for unit in units for slot in unit} + units += [[slot] for slot in singlet_slots if slot not in covered] + units.sort(key=min) + orders = _placements(word_slots, blocks, units) + partners += [(order, configuration) for order in orders[:per_configuration]] + return partners + + class _InternTable(list): """A list with a hash index of its entries, so interning is O(1).""" @@ -442,6 +599,40 @@ def _add_terms(entry, terms): existing["factor"] += term["factor"] +def _group_phasespace_channels(data, selectors, phasespace_channels): + """One UMAMI channel per set of contributions linked by shared phase-space + channels: its contributions are integrated together, each over the + phase-space channels listed in its ``multichannels``.""" + parent = {key: key for key in phasespace_channels} + + def root(key): + while parent[key] != key: + parent[key] = parent[parent[key]] + key = parent[key] + return key + + for _, keys in selectors: + for key in keys[1:]: + parent[root(key)] = root(keys[0]) + groups: dict[tuple, dict] = {} + for entry, keys in selectors: + group = groups.setdefault(root(keys[0]), {"lookup": {}, "processes": []}) + entry["multichannels"] = [ + group["lookup"].setdefault(key, len(group["lookup"])) for key in keys + ] + group["processes"].append(entry) + data["channels"] = [] + for group in groups.values(): + channels = [phasespace_channels[key] for key in group["lookup"]] + data["channels"].append( + { + "phasespace_order": channels[0]["phasespace_order"], + "phasespace_channels": channels, + "processes": group["processes"], + } + ) + + def _merge_flavour_orbits(data): """Merge different flavours with a common proven computation. @@ -450,9 +641,9 @@ def _merge_flavour_orbits(data): A merged flavour is evaluated by the kept selector, so the complete runtime evaluation (permutation, factor and colour tags included) must agree. """ - candidates = {} incoming = data["provider"]["incoming_count"] for channel in data["channels"]: + candidates = {} kept = [] for entry in channel["processes"]: runtime = entry["runtime"] @@ -464,6 +655,7 @@ def _merge_flavour_orbits(data): tuple(runtime["permutation"]), runtime["factor"], tuple(data["pdg_ids"][next(iter(pdgs))][:incoming]), + tuple(entry["multichannels"]), ) previous = next( ( @@ -481,9 +673,6 @@ def _merge_flavour_orbits(data): previous["members"].extend(entry["members"]) channel["processes"] = kept data["channels"] = [c for c in data["channels"] if c["processes"]] - for index, channel in enumerate(data["channels"]): - for entry in channel["processes"]: - entry["multichannels"] = [index] def _export_helicities(data, runtime_helicities): @@ -634,7 +823,7 @@ def _build_provider( assumptions.append(FLAVOR_BLIND_ASSUMPTION) data: dict[str, Any] = { "kind": "pyamplicol-umami-metadata", - "schema_version": 1, + "schema_version": 2, "provider": { "id": provider_id, "library": "libumami.so" if single else f"libumami_{provider_id}.so", @@ -658,8 +847,11 @@ def _build_provider( } structural_representatives: dict[str, dict] = {} orbit_entries: dict[tuple, tuple[dict, list[int]]] = {} - channel_lookup = {} channel_patterns: dict[tuple, list[int]] = {} + # (order, resonances) of every phase-space channel, and the exported + # contributions with the channels that integrate them + phasespace_channels: dict[tuple, dict] = {} + selectors: list[tuple[dict, list[tuple]]] = [] physical_count = 0 runtime_helicities: list[list[list[int]]] = [] runtime_indices: dict[str, int] = {} @@ -726,6 +918,22 @@ def _build_provider( ) replay = _safe_replay(process) sectors_by_word = {_word(sector): sector for sector in process.color_sectors} + roles = _roles(pdgs, incoming, particles) + resonances = [ + ( + int(r["pdg"]), + frozenset(labels[int(label)] for label in r["labels"]), + int(r["color"]), + ) + for r in _records( + physics.get("extensions", {}).get("s_channel_resonances", ()) + ) + ] + if relation is not None: + resonances = [ + (pdg, frozenset(relation[slot] for slot in legs), color) + for pdg, legs, color in resonances + ] colors = sorted( _records(physics.get("color_components")), key=lambda color: tuple(color.get("word", ())), @@ -860,17 +1068,61 @@ def _build_provider( "members": [member], "multichannels": [], } - # A single complete map per channel pattern. Maps are shared only if - # their fixed external masses agree (already a provider invariant). - if pattern not in channel_lookup: - channel_lookup[pattern] = len(data["channels"]) - data["channels"].append( - {"phasespace_order": channel_order, "processes": []} + # The phase-space channels of the contribution, in the relabelled + # legs. Maps are shared only if their fixed external masses agree + # (already a provider invariant) and they follow the same resonances. + partners = ( + [(order, ())] + if flow is None + else _phasespace_partners( + [slot for slot in order if roles[slot] != 1], + [slot for slot in order if roles[slot] == 1], + roles, + resonances, + source_id, + ) + ) + # Contributions share a phase-space channel only with the same kind + # of particle in every leg, so that the cuts select the same legs + kinds = [None] * count + for slot, pdg in enumerate(pdgs): + particle = particles[int(pdg)] + kinds[relabel[slot]] = ( + abs(int(particle["color"])), + abs(float(particle.get("charge", 0.0))), + particle.get("spin"), + ) + keys = [] + for partner_order, configuration in partners: + _, partner_rotated = _channel_pattern( + [relabel[slot] for slot in partner_order], + incoming, + data["provider"]["masses"], + ) + partner_resonances = sorted( + (sorted(relabel[slot] for slot in legs), pdg) + for pdg, legs in configuration + ) + key = ( + tuple(partner_rotated), + tuple((tuple(legs), abs(pdg)) for legs, pdg in partner_resonances), + tuple(kinds), + ) + if key not in keys: + keys.append(key) + phasespace_channels.setdefault( + key, + { + "phasespace_order": partner_rotated, + "resonances": [ + {"pdg": pdg, "legs": legs} + for legs, pdg in partner_resonances + ], + }, ) - channel_index = channel_lookup[pattern] - entry["multichannels"] = [channel_index] - data["channels"][channel_index]["processes"].append(entry) + selectors.append((entry, keys)) orbit_entries[orbit] = (entry, aligned_order, runtime, relabel) + _group_phasespace_channels(data, selectors, phasespace_channels) _merge_flavour_orbits(data) _export_helicities(data, runtime_helicities) alpha_names = {p["alpha_s_parameter"] for p in data["runtime_processes"]} diff --git a/pyamplicol/src/pyamplicol/generation/s_channel_resonances.py b/pyamplicol/src/pyamplicol/generation/s_channel_resonances.py new file mode 100644 index 0000000000..4131c5c1d9 --- /dev/null +++ b/pyamplicol/src/pyamplicol/generation/s_channel_resonances.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: 0BSD +"""The unstable s-channel propagators a process can put on shell. + +A phase-space mapping has to follow these resonances, or it undersamples their +peaks. This is a species-level reachability sweep over the subsets of external +legs, like the DAG prepass: it ignores helicity and colour, keeps the coupling +orders, and may overestimate. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +from ..models.base import CouplingOrders, Model +from ..processes.ir import CanonicalProcessIR +from .dag_reachability import ( + _combine_coupling_order_tuples, + _coupling_orders_within_limits, + _masks_by_size, + _normalize_coupling_order_limits, + _ordered_splits, +) + + +def s_channel_resonances( + process: CanonicalProcessIR, + model: Model, + max_coupling_orders: Mapping[str, int] | None = None, +) -> list[dict[str, object]]: + """``[{"pdg", "labels", "color"}]`` of every unstable particle that a tree + current of two or more final-state legs can carry, with a mass above their + threshold and a complementary current to close the amplitude with. The + pdg carries the summed charge of the outgoing legs it decays to.""" + + if not any(p.width > 0.0 and p.mass > 0.0 for p in model.particles.values()): + return [] + limits = _normalize_coupling_order_limits(max_coupling_orders) + legs = process.legs + every = final = 0 + species: dict[int, dict[int, set[CouplingOrders]]] = {} + for leg in legs: + bit = 1 << (leg.label - 1) + every |= bit + if leg.is_final: + final |= bit + if leg.outgoing_pdg is not None: + species.setdefault(bit, {}).setdefault(int(leg.outgoing_pdg), set()).add(()) + for mask in _masks_by_size(every): + if mask & (mask - 1) == 0 or mask == every: + continue + current = species.setdefault(mask, {}) + for left_mask, right_mask in _ordered_splits(mask): + left, right = species.get(left_mask), species.get(right_mask) + if not left or not right: + continue + for left_pdg, left_orders in left.items(): + for right_pdg, right_orders in right.items(): + for vertex in model.vertices_accepting( + left_pdg, right_pdg, color_accuracy=process.color_accuracy + ): + if model.skip_duplicate_vertex_orientation(vertex): + continue + vertex_orders = model.vertex_coupling_orders(vertex) + for a in left_orders: + for b in right_orders: + orders = _combine_coupling_order_tuples( + a, b, vertex_orders + ) + if _coupling_orders_within_limits(orders, limits): + current.setdefault(vertex.particles[2], set()).add( + orders + ) + found: dict[tuple[int, tuple[int, ...]], dict[str, object]] = {} + for mask, particles in species.items(): + if mask & ~final or mask.bit_count() < 2: + continue + decay_legs = [leg for leg in legs if mask >> (leg.label - 1) & 1] + threshold = sum(float(model.mass(int(leg.pdg))) for leg in decay_legs) + rest = species.get(every & ~mask, {}) + for pdg, orders in particles.items(): + particle = model.particle(pdg) + if particle.width <= 0.0 or particle.mass <= threshold: + continue + closing = rest.get(model.anti_particle(pdg), ()) + if not any( + _coupling_orders_within_limits( + _combine_coupling_order_tuples(a, b), limits + ) + for a in orders + for b in closing + ): + continue + charge = sum(model.charge(int(leg.pdg)) for leg in decay_legs) + if abs(model.charge(pdg) - charge) > 1e-9: + pdg = model.anti_particle(pdg) + labels = tuple(leg.label for leg in decay_legs) + found.setdefault( + (pdg, labels), + { + "pdg": pdg, + "labels": list(labels), + "color": int(model.color_rep(pdg)), + }, + ) + return [found[key] for key in sorted(found)] diff --git a/pyamplicol/src/pyamplicol/generation/service.py b/pyamplicol/src/pyamplicol/generation/service.py index 3a0df17f71..65ce7ae73a 100644 --- a/pyamplicol/src/pyamplicol/generation/service.py +++ b/pyamplicol/src/pyamplicol/generation/service.py @@ -228,6 +228,7 @@ build_runtime_expression_schema, build_runtime_schema_layout, ) +from .s_channel_resonances import s_channel_resonances from .stage_compiler import ( build_and_write_generic_stage_evaluator_artifacts, write_model_parameter_evaluator_artifact, @@ -6147,6 +6148,9 @@ def report_selector(details: Mapping[str, object]) -> None: ), color_plan=prepared.complete_color_plan, ) + physics["extensions"]["s_channel_resonances"] = s_channel_resonances( + expanded.process_ir, model, self._coupling_order_limits or None + ) profile["physics-metadata"] = time.perf_counter() - physics_started run = self._run_config assert run is not None @@ -6744,6 +6748,9 @@ def _construct_eager_plan_v3_artifact( runtime_model_parameter_records_from_eager_input(lowering_input) ), ) + physics["extensions"]["s_channel_resonances"] = s_channel_resonances( + ir, model, self._coupling_order_limits or None + ) return EagerPlanV3ProcessArtifact( process_id=process_name, expression=ir.process, diff --git a/pyamplicol/tests/unit/test_umami_phasespace_partners.py b/pyamplicol/tests/unit/test_umami_phasespace_partners.py new file mode 100644 index 0000000000..4fb27d05cf --- /dev/null +++ b/pyamplicol/tests/unit/test_umami_phasespace_partners.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: 0BSD +"""The phase-space channels that integrate one leading-colour order together: +where its colour singlets attach to the colour lines, which non-cyclic order +its lines take, and which s-channel resonances it maps out.""" + +from __future__ import annotations + +import pytest + +from pyamplicol.artifacts.umami import ( + _configurations, + _line_blocks, + _phasespace_partners, +) + +# all-outgoing colour representations +Q, QBAR, G, S = 3, -3, 8, 1 + + +def test_line_blocks_follow_the_word(): + roles = [QBAR, Q, G, Q, QBAR, G] + assert _line_blocks([1, 2, 0, 3, 5, 4], roles) == [[1, 2, 0], [3, 5, 4]] + # a gluon trace has no open lines + assert _line_blocks([0, 1, 2], [G, G, G]) is None + + +def test_singlets_follow_the_antiquark_in_every_order(): + # q q~ > e+ e- g, the word q g q~ crossed: slots 1 (q~ in) ... 0 (q in) + roles = [QBAR, Q, S, S, G] + partners = _phasespace_partners([1, 4, 0], [2, 3], roles, [], "q q~ > e+ e- g") + assert [order for order, _ in partners] == [[1, 4, 0, 2, 3], [1, 4, 0, 3, 2]] + assert all(configuration == () for _, configuration in partners) + + +def test_singlets_are_distributed_over_the_lines(): + # two lines and one singlet: after the first or the second antiquark + roles = [Q, QBAR, Q, QBAR, S] + partners = _phasespace_partners([0, 1, 2, 3], [4], roles, [], "4q + W") + assert [order for order, _ in partners] == [[0, 1, 2, 3, 4], [0, 1, 4, 2, 3]] + + +def test_three_lines_take_both_cyclic_orders(): + roles = [Q, QBAR] * 3 + partners = _phasespace_partners(list(range(6)), [], roles, [], "6q") + assert [order for order, _ in partners] == [ + [0, 1, 2, 3, 4, 5], + [0, 1, 4, 5, 2, 3], + ] + + +def test_a_resonance_keeps_its_decay_products_together(): + roles = [QBAR, Q, S, S, G] + z = (23, frozenset({2, 3}), 1) + partners = _phasespace_partners([1, 4, 0], [2, 3], roles, [z], "q q~ > e+ e- g") + # the non-resonant placements and the Z, with its legs as one unit + assert partners == [ + ([1, 4, 0, 2, 3], ()), + ([1, 4, 0, 3, 2], ()), + ([1, 4, 0, 2, 3], ((23, frozenset({2, 3})),)), + ] + + +def test_a_colour_singlet_decaying_to_a_complete_line(): + roles = [Q, QBAR, Q, QBAR, G] + z = (23, frozenset({2, 3}), 1) + # the Z decays to the line 2 3 of this order, but not to legs of two lines + assert len(_phasespace_partners([0, 1, 2, 3], [], roles[:4], [z], "4q")) == 2 + assert len(_phasespace_partners([0, 3, 2, 1], [], roles[:4], [z], "4q")) == 1 + + +def test_overlapping_resonances_are_alternatives(): + # e+ e- e+ ve: a Z from either positron, with the W from the other + z1, z2 = (23, frozenset({2, 4})), (23, frozenset({3, 4})) + w1, w2 = (24, frozenset({2, 5})), (24, frozenset({3, 5})) + configurations = _configurations([z1, z2, w1, w2]) + assert configurations[0] == () + assert sorted(map(set, configurations[1:]), key=str) == sorted( + [{z1, w2}, {z2, w1}], key=str + ) + + +def test_nested_resonances_stay_contiguous(): + roles = [QBAR, Q, S, S, S, S] + h = (25, frozenset({2, 3, 4, 5}), 1) + z1, z2 = (23, frozenset({2, 4}), 1), (23, frozenset({3, 5}), 1) + partners = _phasespace_partners([1, 0], [2, 3, 4, 5], roles, [h, z1, z2], "4l") + resonant = [order for order, configuration in partners if len(configuration) == 3] + assert resonant == [[1, 0, 2, 4, 3, 5]] + + +@pytest.mark.parametrize( + "resonance, message", + [ + ((6, frozenset({2, 3, 4}), 3), "coloured s-channel resonance 6"), + ((25, frozenset({2, 4}), 1), "coloured and colourless"), + ], +) +def test_unsupported_resonances_raise(resonance, message): + roles = [QBAR, Q, Q, S, S] + with pytest.raises(ValueError, match=message): + _phasespace_partners([1, 2, 0], [3, 4], roles, [resonance], "process") diff --git a/pyamplicol/umami_example/test_integrate.py b/pyamplicol/umami_example/test_integrate.py index 31b8705f4b..776b75ac25 100644 --- a/pyamplicol/umami_example/test_integrate.py +++ b/pyamplicol/umami_example/test_integrate.py @@ -70,12 +70,12 @@ def load_channels(path, provider=None): channels = [] for ci, ch in enumerate(data["channels"]): inits, finals, procs = set(), [], [] + if len(ch.get("phasespace_channels", [ch])) != 1: + raise ValueError( + "this example integrates each channel with a single map; weighting " + "several phase-space channels is not implemented here" + ) for proc in ch["processes"]: - partners = proc.get("multichannels", [ci]) - if partners != [ci]: - raise ValueError( - "this example requires one nonduplicated map per contribution" - ) cfgs = [] for me in proc["matrix_elements"]: pdg = list(data["pdg_ids"][me["pdg_ids"]]) From 4be41de6894a68be3b908be488b87a0185cd7f4f Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 26 Oct 2026 18:43:42 +0100 Subject: [PATCH 51/88] mg7: [amplicol] multi-channel phase space with colour-singlet partners and resonances Every UMAMI channel becomes one subprocess with a channel per phase-space channel: a colour-ordered chain whose children include the resonances as Breit-Wigner decays, weighted by MappingChannelWeights so that cwnet and MadNIS apply. The flavours of a selector are split by the charges of their resonances, and the LHE writer takes the resonances of an event from the phase-space channel drawn in proportion to its weight. The representative processes of the [amplicol] definitions cover every particle of each final-state leg, so that the default histograms of p p > e+ e- j j are of jets and not of gluons, which most subprocesses lack. New acceptance tests: p p > e+ e- j (Z) and p p > e+ ve j (W+). --- madgraph/iolibs/export_amplicol.py | 65 +++-- .../template_files/mg7/subprocess_amplicol.py | 234 +++++++++++++++--- .../mg7/subprocess_madmatrix.py | 3 + tests/acceptance_tests/test_mg7_amplicol.py | 44 +++- .../interface/test_amplicol_syntax.py | 71 +++++- 5 files changed, 353 insertions(+), 64 deletions(-) diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index ade077d8f6..52b3e89132 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -17,11 +17,13 @@ pyamplicol generates the colour-ordered matrix elements of a process definition in SubProcesses/A_: a leading-colour artifact (lc), a full-colour one (full), and the sources of the UMAMI library returning their ratio (reweight). -Its makefile builds the libraries into lib/. Every colour ordering (UMAMI -channel) of every provider becomes one entry of subprocesses.json, integrated -on its own with a colour-ordered phase-space mapping. Subprocesses related by a -flavour relabelling that is a symmetry of the allowed vertices are generated -once (pyamplicol's strict flavour-blind grouping) and evaluated through their +Its makefile builds the libraries into lib/. Every UMAMI channel of every +provider becomes one entry of subprocesses.json: colour orderings integrated +over a common set of colour-ordered phase-space channels, which differ in where +the colour singlets attach and in the s-channel resonances they follow, and are +weighted by their mapping densities. Subprocesses related by a flavour +relabelling that is a symmetry of the allowed vertices are generated once +(pyamplicol's strict flavour-blind grouping) and evaluated through their representative. """ @@ -223,17 +225,23 @@ def model_job(model, cache_path): def representative_processes(proc_defs, model): """Processes standing for the [amplicol] definitions where the output needs - MG processes (run card defaults, proc_characteristics): every initial state, - with the first particle of each final-state leg.""" + MG processes (run card defaults, proc_characteristics): every initial state + with the first particle of each final-state leg, and every particle of each + final-state leg once, so that the default histograms see the whole group + (a jet, not a gluon).""" import itertools import madgraph.core.base_objects as base_objects processes = [] for procdef in proc_defs: legs = procdef.get('legs') initial = [leg_pdgs(leg, model) for leg in legs if not leg.get('state')] - final = [leg_pdgs(leg, model)[0] for leg in legs if leg.get('state')] - for ids in itertools.product(*initial): - pdgs = list(ids) + final + finals = [leg_pdgs(leg, model) for leg in legs if leg.get('state')] + first = [pdgs[0] for pdgs in finals] + states = [list(ids) + first for ids in itertools.product(*initial)] + states += [[pdgs[0] for pdgs in initial] + first[:i] + [pdg] + first[i + 1:] + for i, options in enumerate(finals) for pdg in options[1:]] + for pdgs in states: + ids = pdgs[:len(initial)] processes.append(base_objects.Process({ 'legs': base_objects.LegList([ base_objects.Leg({'id': pdg, 'number': number + 1, @@ -264,6 +272,16 @@ def mirrored(pdgs, incoming, initial_legs, exported): and swapped not in exported) +def resonance_pdg(model, pdg, decay_pdgs): + """The resonance of the given species carrying the charge of its outgoing + decay products: W+ or W-, say, for the rows of one selector.""" + particle = model.get_particle(abs(pdg)) + if particle.get('self_antipart'): + return abs(pdg) + charge = sum(model.get_particle(p).get_charge() for p in decay_pdgs) + return abs(pdg) if abs(particle.get_charge() - charge) < 1e-6 else -abs(pdg) + + def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power, initial_legs): """One subprocesses.json entry per channel of a UMAMI provider; initial_legs @@ -275,25 +293,30 @@ def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power, for pdgs in pdg_ids} entries = [] for channel_index, channel in enumerate(metadata['channels']): + phasespace = channel['phasespace_channels'] flavors = [] for flavor_index, process in enumerate(channel['processes']): - if process['multichannels'] != [channel_index]: - raise madgraph.MadGraph5Error( - 'AmpliCol contributions shared by several channels are not ' - 'supported yet') # the rows of a selector share its evaluation and incoming legs: one - # flavour per mirror flag, its options written in proportion to factor - by_mirror = {} + # flavour per mirror flag and resonance charges, its options written + # in proportion to factor + by_kind = {} for term in process['matrix_elements']: pdgs = pdg_ids[term['pdg_ids']] mirror = mirrored(pdgs, incoming, initial_legs, exported) - flavor = by_mirror.setdefault(mirror, { + resonances = [ + [resonance_pdg(model, r['pdg'], [pdgs[leg] for leg in r['legs']]) + for r in phasespace[index]['resonances']] + for index in process['multichannels']] + flavor = by_kind.setdefault((mirror, str(resonances)), { 'index': flavor_index, 'options': [], 'weights': [], - 'factor': 0, 'mirror': mirror}) + 'factor': 0, 'mirror': mirror, + 'color_index': process['color_flows'] or 0, + 'channels': process['multichannels'], + 'resonance_pdgs': resonances}) flavor['options'].append(pdgs) flavor['weights'].append(term['factor']) flavor['factor'] += term['factor'] - flavors.extend(by_mirror.values()) + flavors.extend(by_kind.values()) first = flavors[0]['options'][0] entries.append({ 'backend': 'amplicol', @@ -306,7 +329,9 @@ def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power, 'artifact': pjoin(rel_dir, 'lc'), 'ratio_artifact': pjoin(rel_dir, 'reweight'), 'umami_channel': channel_index, - 'phasespace_order': channel['phasespace_order'], + 'phasespace_channels': [ + {'order': ps['phasespace_order'], 'resonances': ps['resonances']} + for ps in phasespace], 'masses': provider['masses'], 'supports_alpha_s': provider['supports_alpha_s'], 'flavors': flavors, diff --git a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py index 965b26c5e0..890e2a29c8 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py +++ b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py @@ -1,10 +1,12 @@ """Subprocesses with AmpliCol matrix elements ('[amplicol]' processes). -Each subprocess is one colour ordering (UMAMI channel) of an AmpliCol provider, -integrated on its own at leading colour with a colour-ordered phase-space -mapping. A second pass re- or unweights its events to full colour with the -ratio of the full-colour to the leading-colour matrix element summed over all -colour orderings, at the helicity of the event. +Each subprocess is one UMAMI channel of an AmpliCol provider: colour orderings +integrated at leading colour over a common set of colour-ordered phase-space +channels. They differ in where the colour singlets attach to the colour lines +and in the s-channel resonances they map out, and are weighted by the densities +of their mappings. A second pass re- or unweights the events to full colour with +the ratio of the full-colour to the leading-colour matrix element summed over +all colour orderings, at the helicity of the event. """ from __future__ import annotations @@ -53,16 +55,146 @@ def chain_topology(incoming_masses: list[float], outgoing_masses: list[float]): ) +def channel_topology( + incoming_masses: list[float], + outgoing_masses: list[float], + order: list[int], + resonances: list[tuple[int, list[int]]], + process: MadgraphProcess, +) -> tuple[ms.Topology, list[int]]: + """The topology of a colour-ordered phase-space channel, a t-channel chain + through the final-state legs and the resonances (pdg, legs) that decay to + them, and its colour order over the children of that chain.""" + n_in = len(incoming_masses) + legs = [frozenset(r_legs) for _, r_legs in resonances] + top = [i for i, r in enumerate(legs) if not any(r < other for other in legs)] + children = [] + for slot in order: + if slot < n_in: + continue + child = next((("r", i) for i in top if slot in legs[i]), ("o", slot)) + if child not in children: + children.append(child) + elif children[-1] != child: + raise RuntimeError(f"the resonance legs of {order} are not contiguous") + chain_count = len(children) - 1 + + def ref(child): + kind, index = child + return f"o{index - n_in}" if kind == "o" else f"p{chain_count + index}" + + refs = [ref(child) for child in children] + if len(children) == 1: + vertices = [["i0", "i1", refs[0]]] + else: + vertices = [["i0", "p0", refs[0]]] + vertices += [[f"p{i - 1}", f"p{i}", refs[i]] for i in range(1, chain_count)] + vertices.append(["i1", f"p{chain_count - 1}", refs[-1]]) + position = {slot: i for i, slot in enumerate(order)} + for index, r_legs in enumerate(legs): + inner = [ + i for i, other in enumerate(legs) + if other < r_legs and not any(other < mid < r_legs for mid in legs) + ] + covered = set().union(*(legs[i] for i in inner)) + decay = [(min(position[s] for s in legs[i]), ref(("r", i))) for i in inner] + decay += [(position[s], ref(("o", s))) for s in r_legs - covered] + vertices.append([f"p{chain_count + index}", *(r for _, r in sorted(decay))]) + propagators = [ms.Propagator() for _ in range(chain_count)] + for pdg, _ in resonances: + propagators.append(ms.Propagator( + mass=process.get_mass(abs(pdg)), + width=process.get_width(abs(pdg)), + pdg_id=pdg, + )) + topology = ms.Topology( + ms.Diagram(incoming_masses, outgoing_masses, propagators, vertices) + ) + color_order = [] + for slot in order: + node = slot if slot < n_in else n_in + children.index( + next((("r", i) for i in top if slot in legs[i]), ("o", slot)) + ) + if node not in color_order: + color_order.append(node) + return topology, color_order + + +def channel_topologies(meta: dict, process: MadgraphProcess) -> list[tuple]: + """channel_topology of every phase-space channel of a subprocess.""" + incoming_masses = [process.get_mass(pid) for pid in clean_pids(meta["incoming"])] + outgoing_masses = [process.get_mass(pid) for pid in clean_pids(meta["outgoing"])] + return [ + channel_topology( + incoming_masses, + outgoing_masses, + channel["order"], + [(r["pdg"], r["legs"]) for r in channel["resonances"]], + process, + ) + for channel in meta["phasespace_channels"] + ] + + +def lhe_diagrams(meta: dict) -> tuple[list[tuple[int, tuple]], list[list[float]]]: + """The diagrams that the LHE writer takes the resonances of an event from, + one per phase-space channel and resonance charges, and per flavour a mask + of those it can be written with. None if there are no resonances.""" + if not any(channel["resonances"] for channel in meta["phasespace_channels"]): + return [], [] + diagrams = [] + for flavor in meta["flavors"]: + for channel, pdgs in zip(flavor["channels"], flavor["resonance_pdgs"]): + if (channel, tuple(pdgs)) not in diagrams: + diagrams.append((channel, tuple(pdgs))) + masks = [ + [ + float(diagram in zip(flavor["channels"], map(tuple, flavor["resonance_pdgs"]))) + for diagram in diagrams + ] + for flavor in meta["flavors"] + ] + return diagrams, masks + + def amplicol_subproc_args(meta: dict, process: MadgraphProcess) -> ms.SubprocArgs: - """LHE completion of an AmpliCol subprocess: its events have no diagram.""" + """LHE completion of an AmpliCol subprocess: its diagrams are the + phase-space channels, with the resonances they map out.""" + diagrams, _ = lhe_diagrams(meta) + if diagrams: + topologies = channel_topologies(meta, process) + particle_count = len(meta["incoming"]) + len(meta["outgoing"]) + colors = sorted({flavor["color_index"] for flavor in meta["flavors"]}) + by_channel = [ + [(d, pdgs) for d, (c, pdgs) in enumerate(diagrams) if c == channel] + for channel in range(len(topologies)) + ] + # the chain propagators come first: one less than the root's children + chain_counts = [len(order) - len(meta["incoming"]) - 1 for _, order in topologies] + layout = dict( + topologies=[topology for topology, _ in topologies], + permutations=[ + [list(range(particle_count))] * len(channel) for channel in by_channel + ], + diagram_indices=[[d for d, _ in channel] for channel in by_channel], + diagram_color_indices=[[colors] * len(channel) for channel in by_channel], + diagram_propagator_pdgs=[ + [[0] * count + list(pdgs) for _, pdgs in channel] + for channel, count in zip(by_channel, chain_counts) + ], + ) + else: + layout = dict( + topologies=[chain_topology( + [process.get_mass(pid) for pid in clean_pids(meta["incoming"])], + [process.get_mass(pid) for pid in clean_pids(meta["outgoing"])], + )], + permutations=[[]], + diagram_indices=[[]], + diagram_color_indices=[[]], + ) return ms.SubprocArgs( - topologies=[chain_topology( - [process.get_mass(pid) for pid in clean_pids(meta["incoming"])], - [process.get_mass(pid) for pid in clean_pids(meta["outgoing"])], - )], - permutations=[[]], - diagram_indices=[[]], - diagram_color_indices=[[]], + **layout, color_flows=meta["color_flows"], pdg_color_types={ int(key): value for key, value in meta["pdg_color_types"].items() @@ -136,31 +268,59 @@ def build_multi_channel_data(self): return None def build_multichannel_phasespace(self) -> PhaseSpace: - mapping = ms.PhaseSpaceMapping( - self.incoming_masses + self.outgoing_masses, - self.process.e_cm, - leptonic=self.process.leptonic, - invariant_power=self.process.run_card["phasespace"]["invariant_power"], - mode=ms.PhaseSpaceMapping.color_ordered, - cuts=self.cuts, - color_ordered=ms.ColorOrderedMapping.Options(self.meta["phasespace_order"]), - beam_rapidity=self.process.beam_rapidity, - mirror_beams=self.mirror_beams, - ) - prefix = f"subproc{self.subproc_id}.amplicol" - discrete_sym, discrete_flavor = self.build_discrete( - 1, len(self.meta["flavors"]), prefix - ) - channel = Channel( - phasespace_mapping=mapping, - adaptive_mapping=self.build_vegas(mapping, prefix), - discrete_sym=discrete_sym, - discrete_flavor=discrete_flavor, - channel_weight_indices=[0], - name="C", - active_flavors=[], + """One channel per phase-space channel of the UMAMI channel, weighted + by the densities of their mappings.""" + flavors = self.meta["flavors"] + particle_count = len(self.incoming_masses) + len(self.outgoing_masses) + channels = [] + for index, (topology, color_order) in enumerate( + channel_topologies(self.meta, self.process) + ): + mapping = ms.PhaseSpaceMapping( + topology, + self.process.e_cm, + leptonic=self.process.leptonic, + invariant_power=self.process.run_card["phasespace"]["invariant_power"], + t_channel_mode=ms.PhaseSpaceMapping.color_ordered, + cuts=self.cuts, + permutations=[list(range(particle_count))], + color_ordered=ms.ColorOrderedMapping.Options(color_order), + beam_rapidity=self.process.beam_rapidity, + mirror_beams=self.mirror_beams, + ) + prefix = f"subproc{self.subproc_id}.amplicol{index}" + discrete_sym, discrete_flavor = self.build_discrete( + 1, len(flavors), prefix + ) + channels.append(Channel( + phasespace_mapping=mapping, + adaptive_mapping=self.build_vegas(mapping, prefix), + discrete_sym=discrete_sym, + discrete_flavor=discrete_flavor, + channel_weight_indices=[index], + name=f"C{index}", + active_flavors=[[ + f for f, flavor in enumerate(flavors) if index in flavor["channels"] + ]], + )) + diagrams, diagram_masks = lhe_diagrams(self.meta) + weights = None + if len(channels) > 1 or diagrams: + weights = ms.MappingChannelWeights( + [channel.phasespace_mapping for channel in channels], + [ + [float(index in flavor["channels"]) for index in range(len(channels))] + for flavor in flavors + ], + [channel for channel, _ in diagrams], + diagram_masks, + ) + return PhaseSpace( + mode="multichannel", + channels=channels, + symfact=[None] * len(channels), + mapping_chan_weights=weights, ) - return PhaseSpace(mode="multichannel", channels=[channel], symfact=[None]) def build_flat_phasespace(self) -> PhaseSpace: return self.build_multichannel_phasespace() diff --git a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py index 7d17692695..078c691372 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py +++ b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py @@ -38,6 +38,7 @@ class PhaseSpace: prop_chan_weights: ms.PropagatorChannelWeights | None = None subchan_weights: ms.SubchannelWeights | None = None cwnet: ms.ChannelWeightNetwork | None = None + mapping_chan_weights: ms.MappingChannelWeights | None = None class MultiChannelData(NamedTuple): @@ -942,6 +943,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: cwnet=self.build_cwnet(len(phasespace.symfact)), prop_chan_weights=phasespace.prop_chan_weights, subchan_weights=phasespace.subchan_weights, + mapping_chan_weights=phasespace.mapping_chan_weights, ) def build_vegas(self, mapping: ms.PhaseSpaceMapping, prefix: str) -> ms.VegasMapping: @@ -1172,6 +1174,7 @@ def build_integrands( madnis_args["compressed_channel_weight_count"], custom_bias, None if self.process.leptonic else self.process.pdf_grid2, + phasespace.mapping_chan_weights, )) #print(integrands[1].function()) #for i in integrands: print(i.function()) diff --git a/tests/acceptance_tests/test_mg7_amplicol.py b/tests/acceptance_tests/test_mg7_amplicol.py index e5f5559cd1..3d4a60ee8a 100644 --- a/tests/acceptance_tests/test_mg7_amplicol.py +++ b/tests/acceptance_tests/test_mg7_amplicol.py @@ -45,8 +45,11 @@ _PDF = 'NNPDF23_lo_as_0130_qed' -def _edit_run_card(path, **amplicol): +def _edit_run_card(path, output_format=None, **amplicol): card = open(path).read() + if output_format: + card = re.sub(r'(?m)^output_format = ".*"', 'output_format = "%s"' % + output_format, card) card = card.replace('fixed_ren_scale = false', 'fixed_ren_scale = true') card = card.replace('fixed_fact_scale = false', 'fixed_fact_scale = true') card = re.sub(r'(?m)^(pdf[12]) = ".*"$', r'\1 = "%s"' % _PDF, card) @@ -84,6 +87,16 @@ def _color_balanced(lhe_path): return True +def _resonances(lhe_path): + """The (pdg, mass) of the status-2 particles of every event.""" + with gzip.open(lhe_path, 'rt') as stream: + text = stream.read() + return [[(int(particle[0]), float(particle[10])) + for particle in (line.split() for line in event.strip().splitlines()[1:]) + if len(particle) > 10 and particle[1] == '2'] + for event in re.findall(r'\n(.*?)', text, re.S)] + + class MG7AmplicolTest(unittest.TestCase): def setUp(self): @@ -160,6 +173,35 @@ def test_gg_ttxg(self): amplicol = self.run_mg7('amp', ['g g > t t~ g [amplicol]'])[:2] self.assert_compatible(amplicol, madmatrix) + def assert_resonances(self, run_dir, pdg, mass, width): + """Some events write the resonance, within bw_cutoff widths of its + mass, and no other one.""" + events = glob.glob(pjoin(run_dir, 'Events', '*', 'events.lhe.gz')) + self.assertTrue(events) + self.assertTrue(_color_balanced(events[-1])) + resonances = [r for event in _resonances(events[-1]) for r in event] + self.assertGreater(len(resonances), 0.2 * _EVENTS) + self.assertEqual({r[0] for r in resonances}, {pdg}) + self.assertLess(max(abs(m - mass) for _, m in resonances), 15 * width) + + def test_pp_eej(self): + """Colour singlets attached to the quark line in either order and a Z + resonance: three phase-space channels per colour order, weighted by + their mapping densities, and the Z in the event record.""" + madmatrix = self.run_mg7('mm', ['p p > e+ e- j'])[:2] + amplicol = self.run_mg7('amp', ['p p > e+ e- j [amplicol]'], + output_format='lhe') + self.assert_compatible(amplicol[:2], madmatrix) + self.assert_resonances(amplicol[2], 23, 91.188, 2.441404) + + def test_pp_wpj(self): + """A charged resonance: W+, never W-, from the charge of its decay.""" + madmatrix = self.run_mg7('mm', ['p p > e+ ve j'])[:2] + amplicol = self.run_mg7('amp', ['p p > e+ ve j [amplicol]'], + output_format='lhe') + self.assert_compatible(amplicol[:2], madmatrix) + self.assert_resonances(amplicol[2], 24, 80.419, 2.0476) + def test_mixed(self): """[amplicol] and madmatrix processes in one output.""" jj = self.run_mg7('jj', ['p p > j j'])[:2] diff --git a/tests/unit_tests/interface/test_amplicol_syntax.py b/tests/unit_tests/interface/test_amplicol_syntax.py index c08d9f8660..9c5196b618 100644 --- a/tests/unit_tests/interface/test_amplicol_syntax.py +++ b/tests/unit_tests/interface/test_amplicol_syntax.py @@ -124,7 +124,8 @@ def test_multiparticles(self): self.assertEqual(expression, 'mp0 mp0 > mp0 mp0') self.assertEqual(sorted(multiparticles['mp0']), sorted(['g', 'u', 'd', 's', 'c', 'u~', 'd~', 's~', 'c~'])) - self.assertEqual(len(processes), 81) + # every initial state, and every other particle of each final-state leg + self.assertEqual(len(processes), 81 + 2 * 8) def test_coupling_orders(self): """The orders MG selects by default, as the largest non-QCD orders.""" @@ -137,6 +138,23 @@ def test_coupling_orders(self): self.assertEqual(export_amplicol.coupling_orders( self.cmd._curr_amplicol_defs[0], self.cmd._curr_model), orders) + def test_default_histograms_see_every_final_state_particle(self): + """The representative processes cover every particle of a final-state + leg: the default histograms of p p > e+ e- j j are of jets, which + every subprocess has, not of gluons, which most of them lack.""" + import madgraph.iolibs.export_amplicol as export_amplicol + import madgraph.various.banner as banner_mod + self.cmd.exec_cmd('generate p p > e+ e- j j [amplicol]', printcmd=False, + precmd=True) + processes = export_amplicol.representative_processes( + self.cmd._curr_amplicol_defs, self.cmd._curr_model) + run_card = banner_mod.RunCardMG7() + run_card.create_default_for_process( + {'ninitial': 2, 'nexternal': 6}, '', [processes]) + histograms = list(run_card.dynamic_sections['histograms']) + self.assertIn('jet_2-pt', histograms) + self.assertFalse([name for name in histograms if name.startswith('g_')]) + def test_mirrored_beams(self): """pyamplicol keeps one process per unordered initial state; the other beam order is integrated as the mirror where the definition allows it.""" @@ -164,11 +182,13 @@ def test_selector_rows_are_folded_into_weighted_options(self): 'supports_alpha_s': True}, 'channels': [{ 'phasespace_order': [0, 2, 3, 4, 1], + 'phasespace_channels': [ + {'phasespace_order': [0, 2, 3, 4, 1], 'resonances': []}], 'processes': [ - {'multichannels': [0], 'matrix_elements': [ + {'multichannels': [0], 'color_flows': 0, 'matrix_elements': [ {'pdg_ids': 0, 'factor': 1}, {'pdg_ids': 2, 'factor': 2}, {'pdg_ids': 1, 'factor': 1}]}, - {'multichannels': [0], 'matrix_elements': [ + {'multichannels': [0], 'color_flows': 1, 'matrix_elements': [ {'pdg_ids': 3, 'factor': 1}]}, ], }], @@ -181,11 +201,50 @@ def test_selector_rows_are_folded_into_weighted_options(self): self.assertEqual(len(entries), 1) flavors = entries[0]['flavors'] # d d~ > s s~ g has its mirror exported, the other rows are mirrored + common = {'channels': [0], 'resonance_pdgs': [[]]} self.assertEqual(flavors, [ {'index': 0, 'options': [[1, -1, 2, -2, 21], [1, -1, 4, -4, 21]], - 'weights': [1, 2], 'factor': 3, 'mirror': True}, + 'weights': [1, 2], 'factor': 3, 'mirror': True, 'color_index': 0, + **common}, {'index': 0, 'options': [[1, -1, 3, -3, 21]], 'weights': [1], - 'factor': 1, 'mirror': False}, + 'factor': 1, 'mirror': False, 'color_index': 0, **common}, {'index': 1, 'options': [[-1, 1, 3, -3, 21]], 'weights': [1], - 'factor': 1, 'mirror': False}, + 'factor': 1, 'mirror': False, 'color_index': 1, **common}, ]) + + def test_resonance_charges_split_the_selector_rows(self): + """A selector integrated over phase-space channels with resonances + writes them with the charge of the decay products of each row: the + charge-conjugated rows of a W selector become flavours of their own.""" + import madgraph.iolibs.export_amplicol as export_amplicol + metadata = { + 'pdg_ids': [[2, -1, -11, 12, 21], [4, -3, -11, 12, 21], + [1, -2, 11, -12, 21]], + 'provider': {'incoming_count': 2, 'masses': [0.] * 5, + 'supports_alpha_s': True}, + 'channels': [{ + 'phasespace_order': [0, 4, 1, 2, 3], + 'phasespace_channels': [ + {'phasespace_order': [0, 4, 1, 2, 3], 'resonances': []}, + {'phasespace_order': [0, 4, 1, 2, 3], + 'resonances': [{'pdg': 24, 'legs': [2, 3]}]}], + 'processes': [ + {'multichannels': [0, 1], 'color_flows': 0, + 'matrix_elements': [{'pdg_ids': index, 'factor': 1} + for index in range(3)]}, + ], + }], + 'color_flows': [], 'helicities': [], 'runtime_processes': [], + } + partons = {21, 1, 2, 3, 4, -1, -2, -3, -4} + entries = export_amplicol.subprocess_entries( + metadata, self.cmd._curr_model, 'SubProcesses/A1', 'me.so', 'ratio.so', + 3, [partons, partons]) + flavors = entries[0]['flavors'] + self.assertEqual([f['options'] for f in flavors], [ + [[2, -1, -11, 12, 21], [4, -3, -11, 12, 21]], [[1, -2, 11, -12, 21]]]) + self.assertEqual([f['resonance_pdgs'] for f in flavors], + [[[], [24]], [[], [-24]]]) + self.assertEqual(entries[0]['phasespace_channels'][1], + {'order': [0, 4, 1, 2, 3], + 'resonances': [{'pdg': 24, 'legs': [2, 3]}]}) From f0d718381941b906a75587847aa71624dcd17139 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 27 Oct 2026 10:27:01 +0100 Subject: [PATCH 52/88] pyamplicol: config registry test counts process.flavor_blind_strict The strict flavour-blind option added an 85th configuration field. --- pyamplicol/tests/unit/test_config_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyamplicol/tests/unit/test_config_models.py b/pyamplicol/tests/unit/test_config_models.py index 6be37cefc1..f0e9f29c8d 100644 --- a/pyamplicol/tests/unit/test_config_models.py +++ b/pyamplicol/tests/unit/test_config_models.py @@ -40,7 +40,7 @@ def test_schema_v1_registry_contains_every_contract_leaf() -> None: - assert len(FIELD_REGISTRY) == 84 + assert len(FIELD_REGISTRY) == 85 assert "generation.umami_grouping" not in FIELD_REGISTRY assert "evaluator.jit.direct_translation" not in FIELD_REGISTRY assert FIELD_REGISTRY["action"].required From dfe3b1f19146973e89ebdf582552de777683a207 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 27 Oct 2026 10:36:02 +0100 Subject: [PATCH 53/88] madspace: limit pretty MadNIS training box to the last 10 subprocesses mg7: remove unused legacy MadNIS status display code --- madgraph/iolibs/template_files/mg7/common.py | 9 -- .../template_files/mg7/madgraph_process.py | 84 +------------------ .../madspace/driver/madnis_training.hpp | 1 + madspace/src/driver/madnis_training.cpp | 42 +++++++--- 4 files changed, 34 insertions(+), 102 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/common.py b/madgraph/iolibs/template_files/mg7/common.py index 253f4d06d9..597a793bd8 100644 --- a/madgraph/iolibs/template_files/mg7/common.py +++ b/madgraph/iolibs/template_files/mg7/common.py @@ -53,15 +53,6 @@ def get_start_time(): return time.time(), time.process_time() -def format_time(t: int, centi: bool = False): - hours, t = divmod(t, 3600) - minutes, seconds = divmod(t, 60) - if centi: - return f"{int(hours):02}:{int(minutes):02}:{seconds:02.2f}" - else: - return f"{int(hours):02}:{int(minutes):02}:{seconds:02.0f}" - - def resolve_verbosity(verbosity: str) -> str: """Resolve the run_card "auto" verbosity to "pretty"/"log" depending on whether stdout is attached to a terminal; other values pass through diff --git a/madgraph/iolibs/template_files/mg7/madgraph_process.py b/madgraph/iolibs/template_files/mg7/madgraph_process.py index ca44e13cbc..af28b667fb 100644 --- a/madgraph/iolibs/template_files/mg7/madgraph_process.py +++ b/madgraph/iolibs/template_files/mg7/madgraph_process.py @@ -9,7 +9,7 @@ import tempfile import time -from madgraph.iolibs.template_files.mg7.common import backend_of, format_time, fptype_of, get_start_time, interference_helicities_summed, lhe_helicities, lhe_weight_info, logger, me_parameters, ms, resolve_auto_backend, resolve_seed, resolve_verbosity +from madgraph.iolibs.template_files.mg7.common import backend_of, fptype_of, get_start_time, interference_helicities_summed, lhe_helicities, lhe_weight_info, logger, me_parameters, ms, resolve_auto_backend, resolve_seed, resolve_verbosity from madgraph.various import misc from madgraph.various.banner import RunCardMG7 from madgraph.models.check_param_card import ParamCard @@ -1589,88 +1589,6 @@ def train_madnis(self) -> None: self.event_generator = self.build_event_generator(madnis_phasespaces) self.madnis_trained = True - def update_madnis_status_single( - self, batch: int, batch_target: int, loss: float, lr: float, channel_count: int - ) -> None: - now = time.time() - if batch + 1 < batch_target: - if now - self.last_update_time < 0.1: - return - self.last_update_time = now - progress_bar = ms.format_progress((batch + 1) / batch_target, 52) - time_diff = now - self.madnis_wall_time - time_str = f"{format_time(time_diff)}" - else: - progress_bar = "" - wall_diff = now - self.madnis_wall_time - cpu_diff = time.process_time() - self.madnis_cpu_time - time_str = ( - f"{format_time(wall_diff, centi=True)} wall, " - f"{format_time(cpu_diff, centi=True)} cpu" - ) - batch_str = f"{batch + 1} / {batch_target}" - self.madnis_box.set_column(1, [ - f"{batch_str:<15} {progress_bar}", - f"{loss:>.4f}", - f"{channel_count}", - time_str - ]) - self.madnis_box.print_update() - - def update_madnis_status_multi( - self, - subproc_id: int, - batch: int, - batch_target: int, - loss: float, - lr: float, - channel_count: int - ) -> None: - now = time.time() - subproc_count = len(self.subprocesses) - if batch + 1 < batch_target: - if now - self.last_update_time < 0.1: - return - self.last_update_time = now - progress_bar = ms.format_progress((batch + 1) / batch_target, 34) - progress_bar_all = ms.format_progress( - (subproc_id * batch_target + batch + 1) / (subproc_count * batch_target), - 52 - ) - time_str = f"{format_time(now - self.madnis_wall_time)}" - subproc_str = f"{subproc_id} / {subproc_count}" - elif subproc_id < subproc_count - 1: - progress_bar = "" - progress_bar_all = ms.format_progress( - ((subproc_id + 1) * batch_target + 1) / (subproc_count * batch_target), - 52 - ) - time_str = f"{format_time(now - self.madnis_wall_time)}" - subproc_str = f"{subproc_id} / {subproc_count}" - else: - progress_bar = "" - progress_bar_all = "" - wall_diff = now - self.madnis_wall_time - cpu_diff = time.process_time() - self.madnis_cpu_time - time_str = ( - f"{format_time(wall_diff, centi=True)} wall, " - f"{format_time(cpu_diff, centi=True)} cpu" - ) - subproc_str = f"{subproc_count} / {subproc_count}" - batch_str = f"{batch + 1} / {batch_target}" - self.madnis_upper_box.set_column(1, [ - f"{subproc_str:<15} {progress_bar_all}", - time_str, - ]) - self.madnis_lower_box.set_row(subproc_id + 1, [ - f"{subproc_id}", - f"{loss:>.4f}", - f"{channel_count}", - f"{batch_str:<15} {progress_bar}", - ]) - self.madnis_upper_box.print_update() - self.madnis_lower_box.print_update() - @property def gridpack_run_mode(self) -> str: """How this run ends: a gridpack run may skip the events (see the diff --git a/madspace/include/madspace/driver/madnis_training.hpp b/madspace/include/madspace/driver/madnis_training.hpp index 8d274c0ab6..5ba75131ec 100644 --- a/madspace/include/madspace/driver/madnis_training.hpp +++ b/madspace/include/madspace/driver/madnis_training.hpp @@ -352,6 +352,7 @@ class MultiMadnisTraining { std::chrono::time_point _last_print_time; PrettyBox _pretty_box_upper; PrettyBox _pretty_box_lower; + nested_vector2 _pretty_subproc_rows; std::shared_ptr _status_file; nlohmann::json _trainings_status; }; diff --git a/madspace/src/driver/madnis_training.cpp b/madspace/src/driver/madnis_training.cpp index c577aa149f..614c310a82 100644 --- a/madspace/src/driver/madnis_training.cpp +++ b/madspace/src/driver/madnis_training.cpp @@ -5,6 +5,13 @@ using namespace madspace; +namespace { + +// maximum number of subprocess rows shown in the pretty training box +constexpr std::size_t max_pretty_subproc_rows = 10; + +} // namespace + MadnisTraining::MadnisTraining( ContextPtr generator_context, ContextPtr optimizer_context, @@ -506,7 +513,12 @@ void MadnisTraining::start_single_job( // buffer content depend on thread scheduling bool store_buffer = _batch_index >= _config.buffer_skip_batches; _generator_context->thread_pool().submit( - [this, channel_index, global_channel_index, batch_size, job_id, store_buffer, + [this, + channel_index, + global_channel_index, + batch_size, + job_id, + store_buffer, &job]() { auto& channel = _channels.at(channel_index); if (_seed) { @@ -1022,8 +1034,10 @@ void MultiMadnisTraining::print_progress_init() { } if (_training_args.size() > 1) { - _pretty_box_lower = - PrettyBox("Subprocesses", _training_args.size() + 1, {12, 12, 12, 0}); + std::size_t row_count = + std::min(_training_args.size(), max_pretty_subproc_rows); + _pretty_subproc_rows.assign(_training_args.size(), std::vector(4)); + _pretty_box_lower = PrettyBox("Subprocesses", row_count + 1, {12, 12, 12, 0}); _pretty_box_lower.set_row(0, {"Subprocess", "Loss", "Channels", "Batch"}); _pretty_box_upper = PrettyBox("MadNIS training", 2, {18, 0}, _pretty_box_lower.line_count()); @@ -1147,13 +1161,21 @@ void MultiMadnisTraining::print_progress_update( _pretty_box_upper.set_column( 1, {std::format("{:<15} {}", subproc_str, progress_bar_all), time_str} ); - _pretty_box_lower.set_row( - subproc_index + 1, - {std::format("{}", subproc_index), - std::format("{:>.4f}", loss), - std::format("{}", chan_count), - std::format("{:<15} {}", batch_str, progress_bar)} - ); + _pretty_subproc_rows.at(subproc_index) = { + std::format("{}", subproc_index), + std::format("{:>.4f}", loss), + std::format("{}", chan_count), + std::format("{:<15} {}", batch_str, progress_bar) + }; + // scroll so that the last rows up to the current subprocess are shown + std::size_t row_count = std::min(subproc_count, max_pretty_subproc_rows); + std::size_t first_row = + subproc_index + 1 > row_count ? subproc_index + 1 - row_count : 0; + for (std::size_t row = 0; row < row_count; ++row) { + _pretty_box_lower.set_row( + row + 1, _pretty_subproc_rows.at(first_row + row) + ); + } _pretty_box_upper.print_update(); _pretty_box_lower.print_update(); } From 668fce10cb57488ddfa89784d9277bac9d7d1152 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 27 Oct 2026 17:17:37 +0100 Subject: [PATCH 54/88] pyamplicol: umami writer test builders provide read_staged --- pyamplicol/tests/unit/test_umami_metadata.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyamplicol/tests/unit/test_umami_metadata.py b/pyamplicol/tests/unit/test_umami_metadata.py index 5c03863708..507257026c 100644 --- a/pyamplicol/tests/unit/test_umami_metadata.py +++ b/pyamplicol/tests/unit/test_umami_metadata.py @@ -741,6 +741,7 @@ def test_writer_append_recovers_old_compact_otf_sdk(tmp_path, monkeypatch): written = {} builder = SimpleNamespace( staged_path=lambda path: tmp_path / path, + read_staged=lambda path: (tmp_path / path).read_bytes(), payload_records=lambda: [], add_bytes=lambda path, content, **kwargs: written.setdefault(path, content), ) @@ -965,6 +966,7 @@ def test_writer_exports_only_ungenerated_group_members(tmp_path): written = {} builder = SimpleNamespace( staged_path=lambda path: tmp_path / path, + read_staged=lambda path: (tmp_path / path).read_bytes(), payload_records=lambda: [], add_bytes=lambda path, content, **kwargs: written.setdefault(path, content), ) From 4033b1d6d8cb18243344b0a84544fcee4c400c56 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Tue, 27 Oct 2026 17:25:30 +0100 Subject: [PATCH 55/88] mg7: move the launcher modules from template_files/mg7 to madgraph/mg7 The launcher modules are imported, not copied into outputs, so they are no longer kept with the templates. The files copied into outputs and gridpacks (api.cpp/h, run_card.toml, custom_functions, gridpack scripts) stay in template_files/mg7, where launch, npy_to_lhe and init_custom_func are kept as redirects for the bin/ scripts of existing outputs. --- madgraph/interface/madgraph_interface.py | 6 +- madgraph/interface/tutorials/lo.py | 6 +- madgraph/iolibs/export_amplicol.py | 4 +- madgraph/iolibs/export_cpp.py | 6 +- .../template_files/mg7/init_custom_func.py | 266 +----------------- madgraph/iolibs/template_files/mg7/launch.py | 78 +---- .../iolibs/template_files/mg7/npy_to_lhe.py | 253 +---------------- madgraph/mg7/__init__.py | 0 .../template_files => }/mg7/bootstrap.py | 0 .../{iolibs/template_files => }/mg7/cli.py | 12 +- .../{iolibs/template_files => }/mg7/common.py | 2 +- .../{iolibs/template_files => }/mg7/driver.py | 8 +- .../template_files => }/mg7/hwu_output.py | 0 madgraph/mg7/init_custom_func.py | 266 ++++++++++++++++++ madgraph/mg7/launch.py | 77 +++++ .../mg7/madgraph_process.py | 27 +- .../template_files => }/mg7/model_widths.py | 2 +- madgraph/mg7/npy_to_lhe.py | 252 +++++++++++++++++ .../template_files => }/mg7/observables.py | 2 +- .../template_files => }/mg7/options.py | 0 .../{iolibs/template_files => }/mg7/plots.py | 0 .../template_files => }/mg7/postprocess.py | 6 +- .../template_files => }/mg7/run_interface.py | 2 +- .../mg7/subprocess_amplicol.py | 8 +- .../mg7/subprocess_madmatrix.py | 6 +- .../mg7/systematics_summary.py | 0 madgraph/various/banner.py | 2 +- tests/acceptance_tests/test_auto_width.py | 4 +- tests/acceptance_tests/test_cmd_madevent.py | 2 +- tests/acceptance_tests/test_mg7_amplicol.py | 2 +- .../test_mg7_custom_functions.py | 4 +- .../interface/test_auto_width_engine.py | 6 +- .../interface/test_auto_width_nlo.py | 12 +- .../test_mg7_interference_helicity.py | 2 +- tests/unit_tests/interface/test_mg7_launch.py | 40 +-- .../interface/test_mg7_systematics_summary.py | 4 +- tests/unit_tests/interface/test_tutorials.py | 2 +- .../iolibs/test_mg7_discrete_flavor.py | 4 +- .../unit_tests/various/test_mg7_hwu_output.py | 2 +- tests/unit_tests/various/test_mg7_plots.py | 4 +- 40 files changed, 695 insertions(+), 684 deletions(-) create mode 100644 madgraph/mg7/__init__.py rename madgraph/{iolibs/template_files => }/mg7/bootstrap.py (100%) rename madgraph/{iolibs/template_files => }/mg7/cli.py (98%) rename madgraph/{iolibs/template_files => }/mg7/common.py (99%) rename madgraph/{iolibs/template_files => }/mg7/driver.py (97%) rename madgraph/{iolibs/template_files => }/mg7/hwu_output.py (100%) create mode 100644 madgraph/mg7/init_custom_func.py create mode 100644 madgraph/mg7/launch.py rename madgraph/{iolibs/template_files => }/mg7/madgraph_process.py (98%) rename madgraph/{iolibs/template_files => }/mg7/model_widths.py (98%) create mode 100644 madgraph/mg7/npy_to_lhe.py rename madgraph/{iolibs/template_files => }/mg7/observables.py (99%) rename madgraph/{iolibs/template_files => }/mg7/options.py (100%) rename madgraph/{iolibs/template_files => }/mg7/plots.py (100%) rename madgraph/{iolibs/template_files => }/mg7/postprocess.py (98%) rename madgraph/{iolibs/template_files => }/mg7/run_interface.py (99%) rename madgraph/{iolibs/template_files => }/mg7/subprocess_amplicol.py (98%) rename madgraph/{iolibs/template_files => }/mg7/subprocess_madmatrix.py (99%) rename madgraph/{iolibs/template_files => }/mg7/systematics_summary.py (100%) diff --git a/madgraph/interface/madgraph_interface.py b/madgraph/interface/madgraph_interface.py index 8dba34a2c3..8e2ed6b426 100755 --- a/madgraph/interface/madgraph_interface.py +++ b/madgraph/interface/madgraph_interface.py @@ -9805,11 +9805,11 @@ def do_launch(self, line): # what needs madspace). The bootstrap has to be told whether it may # take over the terminal: in a subprocess it could read that off # sys.argv/sys.stdin, but in process those describe MG5, not the run. - from madgraph.iolibs.template_files.mg7 import bootstrap as mg7_bootstrap + from madgraph.mg7 import bootstrap as mg7_bootstrap mg7_bootstrap.ensure_madspace( interactive=bool(self.use_rawinput) and not options['force'], jobs=self.options['nb_core']) - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch MG7 = mg7_launch.MG7Cmd(me_dir=me_dir, options=self.options) if options['interactive']: @@ -13766,7 +13766,7 @@ def integrate_widths_madspace(self, opts, decay_info): after others succeeded, the caller still has the decay_info it would hand to MadEvent. """ - from madgraph.iolibs.template_files.mg7 import bootstrap as mg7_bootstrap + from madgraph.mg7 import bootstrap as mg7_bootstrap # the survey of a decay uses no PDF: only the build needs the environment self.setup_mg7_environment(lhapdf=False) mg7_bootstrap.ensure_madspace(interactive=False, diff --git a/madgraph/interface/tutorials/lo.py b/madgraph/interface/tutorials/lo.py index 2514029a7d..300e5d1251 100644 --- a/madgraph/interface/tutorials/lo.py +++ b/madgraph/interface/tutorials/lo.py @@ -45,11 +45,11 @@ def _mg7_bootstrap(): Everything the tutorial wants to know about madspace goes through it, so the tutorial cannot disagree with the launcher about where madspace is or - whether it is there (see iolibs/template_files/mg7/bootstrap.py). + whether it is there (see madgraph/mg7/bootstrap.py). """ try: - from madgraph.iolibs.template_files.mg7 import bootstrap + from madgraph.mg7 import bootstrap except Exception: return None return bootstrap @@ -233,7 +233,7 @@ def _systematics_rows(info): if not info: return illustrative try: - from madgraph.iolibs.template_files.mg7 import systematics_summary + from madgraph.mg7 import systematics_summary except Exception: return illustrative diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index 52b3e89132..dc422c8242 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -73,7 +73,7 @@ def check_pyamplicol(interactive=None, jobs=None): """Install pyamplicol in pyamplicol/install if needed and put it first on sys.path, which the worker inherits. Never imported here, see amplicol_worker.""" - from madgraph.iolibs.template_files.mg7 import bootstrap + from madgraph.mg7 import bootstrap try: bootstrap.ensure_pyamplicol(interactive=interactive, jobs=jobs) except RuntimeError as error: @@ -83,7 +83,7 @@ def check_pyamplicol(interactive=None, jobs=None): def write_rusticol_config(subproc_dir): """SubProcesses/rusticol-config, the RUSTICOL_CONFIG of the makefiles.""" - from madgraph.iolibs.template_files.mg7 import bootstrap + from madgraph.mg7 import bootstrap path = pjoin(subproc_dir, 'rusticol-config') with open(path, 'w') as stream: stream.write(RUSTICOL_CONFIG % { diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index 55c122f543..739c7c99f5 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3223,7 +3223,7 @@ def mg7_launcher_source(interpreter, mg5_dir): " os.execv(_INTERPRETER, [_INTERPRETER,\n" " os.path.abspath(__file__)] + sys.argv[1:])\n" + "sys.path.append(%r)\n" % mg5_dir - + "from madgraph.iolibs.template_files.mg7.launch import main\n" + + "from madgraph.mg7.launch import main\n" "if __name__ == '__main__':\n" " os.chdir(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))\n" " try:\n" @@ -3319,7 +3319,7 @@ def copy_template(self, model): "#! /usr/bin/env python3\n" "import sys, os\n" f"sys.path.append('{MG5DIR}')\n" - "from madgraph.iolibs.template_files.mg7.npy_to_lhe import main\n" + "from madgraph.mg7.npy_to_lhe import main\n" "if __name__ == '__main__':\n" " main(me_dir=os.path.dirname(os.path.dirname(os.path.realpath(__file__))))\n" ) @@ -3331,7 +3331,7 @@ def copy_template(self, model): "#! /usr/bin/env python3\n" "import sys, os\n" f"sys.path.append('{MG5DIR}')\n" - "from madgraph.iolibs.template_files.mg7.init_custom_func import main\n" + "from madgraph.mg7.init_custom_func import main\n" "if __name__ == '__main__':\n" " sys.exit(main(me_dir=os.path.dirname(os.path.dirname(os.path.realpath(__file__)))))\n" ) diff --git a/madgraph/iolibs/template_files/mg7/init_custom_func.py b/madgraph/iolibs/template_files/mg7/init_custom_func.py index 90311d00a2..c8a5064f88 100644 --- a/madgraph/iolibs/template_files/mg7/init_custom_func.py +++ b/madgraph/iolibs/template_files/mg7/init_custom_func.py @@ -1,265 +1,3 @@ -"""Create a skeleton for a custom cut, scale, PDF or bias function. +"""Redirect for the bin/ scripts of outputs made before the move to madgraph/mg7.""" -The function is written for a single phase-space point in C, C++, Fortran, -CUDA or HIP. A generated UMAMI wrapper calls it on batches of events, so that -madspace can load it like a matrix element. The files go into -``custom_/`` in the process directory, and ``make`` there builds -``lib/custom_.so``. The run card option -``[custom_functions] `` takes the path of that library. - -Usage:: - - bin/init_custom_func [-f {cut,scale,pdf,bias}] [-l {c,cpp,fortran,cuda,hip}] - [-y] [--force] [--build | --no-build] - -Options that are not given are asked interactively, unless ``-y`` is passed. -""" - -import argparse -import logging -import os -import shutil -import subprocess -import sys - -logger = logging.getLogger("madevent") - -TEMPLATE_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "custom_functions") -UMAMI_HEADER = os.path.join( - os.path.dirname(os.path.dirname(os.path.realpath(__file__))), "madmatrix", "umami.h" -) - -FUNCTIONS = { - "cut": (["UMAMI_OUT_CUT"], "a cut, true keeps the event and false rejects it"), - "scale": ( - ["UMAMI_OUT_REN_SCALE", "UMAMI_OUT_FACT_SCALE_1", "UMAMI_OUT_FACT_SCALE_2"], - "the renormalization and factorization scales", - ), - "pdf": (["UMAMI_OUT_PDF"], "a parton density x f(x, Q)"), - "bias": (["UMAMI_OUT_BIAS"], "an event bias for the unweighting"), -} -# example parameter of parameters.h: (name, default, run card value) -EXAMPLE_PARAMETERS = { - "cut": ("pt_min", "20.", "30.0"), - "scale": ("ht_fraction", "0.5", "0.25"), - "pdf": ("gluon_norm", "3.", "2.5"), - "bias": ("pt0", "100.", "200.0"), -} -# language: (body template extension, description) -LANGUAGES = { - "c": ("c", "C"), - "cpp": ("cpp", "C++"), - "fortran": ("F90", "Fortran"), - "cuda": ("cu", "CUDA, for runs on NVIDIA GPUs"), - "hip": ("cu", "HIP, for runs on AMD GPUs"), - "python-numpy": ("py", "Python with numpy, for CPU runs, no compilation"), - "python-torch": ("py", "Python with PyTorch, for CPU or GPU runs, no compilation"), -} - -_NONINTERACTIVE = (not sys.stdin.isatty()) or any(a in ("-y", "--yes") for a in sys.argv[1:]) - - -def _ask(question, default, choices): - """Ask through MG5's cmd.ask when available, else input().""" - if _NONINTERACTIVE: - return default - try: - from madgraph.interface.extended_cmd import Cmd - - answer = Cmd().ask(question, default, choices=list(choices), timeout=0) - return default if answer is None else str(answer) - except Exception: - pass - try: - answer = input(f"{question} [{default}]: ").strip() - except EOFError: - return default - return answer or default - - -def ask_choice(title, options, default): - """Numbered menu over ``options`` (name -> description), returns a name.""" - names = list(options) - lines = [title] - width = max(len(name) for name in names) - for i, name in enumerate(names, 1): - lines.append(f" {i}. {name:{width}s} {options[name]}") - print("\n".join(lines)) - choices = names + [str(i) for i in range(1, len(names) + 1)] - while True: - answer = _ask("Enter a number or a name", default, choices).strip().lower() - if answer in names: - return answer - if answer.isdigit() and 1 <= int(answer) <= len(names): - return names[int(answer) - 1] - print(f" Please enter one of {', '.join(names)} or 1-{len(names)}.") - - -def ask_yes_no(question, default=True): - default_s = "y" if default else "n" - while True: - answer = _ask(question, default_s, ["y", "n"]).strip().lower() - if answer in ("y", "yes"): - return True - if answer in ("n", "no"): - return False - print(" Please enter 'y' or 'n'.") - - -def build_parser(): - parser = argparse.ArgumentParser( - prog="bin/init_custom_func", - description=( - "Create a skeleton for a custom cut, scale, PDF or bias function in " - "custom_/. Compiled languages are built into " - "lib/custom_.so, Python functions need no build." - ), - ) - parser.add_argument("-f", "--function", choices=list(FUNCTIONS)) - parser.add_argument("-l", "--language", choices=list(LANGUAGES)) - parser.add_argument( - "-y", "--yes", action="store_true", - help="do not ask, use the defaults (cut, cpp) for options that are not given", - ) - parser.add_argument( - "--force", action="store_true", help="overwrite an existing custom_ directory" - ) - build = parser.add_mutually_exclusive_group() - build.add_argument("--build", dest="build", action="store_true", default=None, - help="run make after creating the files (default)") - build.add_argument("--no-build", dest="build", action="store_false") - run_card = parser.add_mutually_exclusive_group() - run_card.add_argument("--set-run-card", dest="set_run_card", action="store_true", - default=None, help="enable the function in Cards/run_card.toml (default)") - run_card.add_argument("--no-set-run-card", dest="set_run_card", action="store_false") - return parser - - -def is_python(language): - return language.startswith("python-") - - -def function_path(function, language): - """The run card path of the function, relative to the process directory.""" - if is_python(language): - return os.path.join(f"custom_{function}", f"custom_{function}.py") - return os.path.join("lib", f"custom_{function}.so") - - -def write_files(out_dir, function, language): - """Write the function body and, for compiled languages, the wrapper, - umami.h and the Makefile.""" - extension = LANGUAGES[language][0] - body_file = f"custom_{function}.{extension}" - if is_python(language): - framework = language.split("-", 1)[1] - os.makedirs(out_dir, exist_ok=True) - shutil.copy( - os.path.join(TEMPLATE_DIR, f"{function}_{framework}.py"), - os.path.join(out_dir, body_file), - ) - return body_file - values = { - "kind": function, - "function_name": f"custom_{function}", - "output_keys": ", ".join(FUNCTIONS[function][0]), - "is_pdf": int(function == "pdf"), - "is_scale": int(function == "scale"), - "return_type": "bool" if function == "cut" else "double", - "body_file": body_file, - "language": language, - "example": EXAMPLE_PARAMETERS[function][0], - "example_default": EXAMPLE_PARAMETERS[function][1], - "example_value": EXAMPLE_PARAMETERS[function][2], - } - os.makedirs(out_dir, exist_ok=True) - shutil.copy(os.path.join(TEMPLATE_DIR, f"{function}.{extension}"), os.path.join(out_dir, body_file)) - shutil.copy(UMAMI_HEADER, os.path.join(out_dir, "umami.h")) - wrapper = "wrapper_gpu.cu" if language in ("cuda", "hip") else "wrapper_cpu.c" - for template in (wrapper, "Makefile", "parameters.h"): - with open(os.path.join(TEMPLATE_DIR, template)) as f: - text = f.read() % values - with open(os.path.join(out_dir, template), "w") as f: - f.write(text) - return body_file - - -def build(out_dir): - """Run make in ``out_dir``, logging to compile.log. Returns True on success.""" - log_path = os.path.join(out_dir, "compile.log") - with open(log_path, "w") as log: - result = subprocess.run(["make"], cwd=out_dir, stdout=log, stderr=subprocess.STDOUT) - if result.returncode != 0: - with open(log_path) as log: - print(log.read()) - print(f"Building failed, see {log_path}") - return False - return True - - -def set_run_card(me_dir, function, lib_path): - from madgraph.various.banner import RunCardMG7 - - card_path = os.path.join(me_dir, "Cards", "run_card.toml") - run_card = RunCardMG7(card_path) - run_card.set(f"custom_functions.{function}", lib_path) - run_card.write(card_path) - - -def main(me_dir=None, argv=None): - global _NONINTERACTIVE - args = build_parser().parse_args(argv) - if args.yes: - _NONINTERACTIVE = True - me_dir = me_dir or os.getcwd() - - function = args.function or ask_choice( - "Which custom function do you want to create?", - {name: desc for name, (_, desc) in FUNCTIONS.items()}, - "cut", - ) - language = args.language or ask_choice( - "Which language do you want to write it in?", - {name: desc for name, (_, desc) in LANGUAGES.items()}, - "cpp", - ) - - out_dir = os.path.join(me_dir, f"custom_{function}") - if os.path.exists(out_dir): - if not (args.force or (not _NONINTERACTIVE and ask_yes_no( - f"{out_dir} already exists. Overwrite it?", False))): - print(f"{out_dir} already exists, use --force to overwrite it.") - return 1 - shutil.rmtree(out_dir) - - body_file = write_files(out_dir, function, language) - lib_path = function_path(function, language) - print(f"Created {os.path.relpath(out_dir, me_dir)}/. Write your function in " - f"{os.path.relpath(os.path.join(out_dir, body_file), me_dir)}.") - - do_build = args.build - if is_python(language): - if do_build: - print("Python functions need no build.") - do_build = False - elif do_build is None: - do_build = ask_yes_no("Build the default function now?", True) - if do_build and not build(out_dir): - return 1 - - do_set = args.set_run_card - if do_set is None: - do_set = ask_yes_no(f"Set custom_functions.{function} = \"{lib_path}\" in the run card?", True) - if do_set: - set_run_card(me_dir, function, lib_path) - print(f"Enabled it in Cards/run_card.toml: custom_functions.{function} = \"{lib_path}\".") - else: - print(f"To enable it, set custom_functions.{function} = \"{lib_path}\" in Cards/run_card.toml.") - if not do_build and not is_python(language): - print(f"Build it with `make -C {os.path.relpath(out_dir, me_dir)}`. " - "bin/generate_events also rebuilds it before a run.") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) +from madgraph.mg7.init_custom_func import main # noqa: F401 diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index 84bf57465c..c4a6105e79 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -1,77 +1,3 @@ -"""Entry point of the mg7 launcher, re-exporting its modules.""" +"""Redirect for the bin/ scripts of outputs made before the move to madgraph/mg7.""" -from madgraph.iolibs.template_files.mg7.common import * # noqa: F401,F403 -from madgraph.iolibs.template_files.mg7.common import ms, logger # noqa: F401 -from madgraph.iolibs.template_files.mg7.observables import ( # noqa: F401 - CutItem, - HistItem, - _MERGED_PID_REPRESENTATIVE, - clean_pids, - _SIGNED_OBSERVABLES, - _flavor_options, - observable_pids, - build_observable, - decay_products, - pid_is_qcd, -) -from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import ( # noqa: F401 - Channel, - PhaseSpace, - MultiChannelData, - build_topologies, - build_multi_channel_data, - MadgraphSubprocess, -) -from madgraph.iolibs.template_files.mg7.madgraph_process import ( # noqa: F401 - MadgraphProcess, - PCH_MIN_SUBPROCESSES, - _single_flavor_metas, -) -from madgraph.iolibs.template_files.mg7.options import ( # noqa: F401 - _ROOTED_OPTIONS, - load_mg5_options, - lhapdf_paths, -) -from madgraph.iolibs.template_files.mg7.model_widths import ( # noqa: F401 - _model_hash_warned, - _model_cache, - read_stored_model, - _proc_characteristic, - load_process_model, - compute_auto_widths, -) -from madgraph.iolibs.template_files.mg7.postprocess import ( # noqa: F401 - _find_event_file, - _find_or_convert_event_file, - _report_failure, - _off, - _has_effective_handler, - _TOOL_LOGGING_READY, - _setup_logging, - selected_tools, - run_selected_tools, - _add_time_of_flight, - _lhapdf_config_path, - _run_systematics, - run_lhe_postprocessing, -) -from madgraph.iolibs.template_files.mg7.driver import ( # noqa: F401 - compute_decay_widths, - _decay_flavor_options, - _release_channel_generators, - run_single, - detect_run_scan, - detect_param_scan, - run_scan, - run_generation, - _lhe_needed_by, - force_lhe_output_if_needed, -) -from madgraph.iolibs.template_files.mg7.cli import ( # noqa: F401 - MG7Cmd, - build_selector_cmd, - ask_edit_cards, - prune_unselected_tool_cards, - _raise_open_file_limit, - main, -) +from madgraph.mg7.launch import main # noqa: F401 diff --git a/madgraph/iolibs/template_files/mg7/npy_to_lhe.py b/madgraph/iolibs/template_files/mg7/npy_to_lhe.py index df9c62c2f5..8f194255cc 100644 --- a/madgraph/iolibs/template_files/mg7/npy_to_lhe.py +++ b/madgraph/iolibs/template_files/mg7/npy_to_lhe.py @@ -1,252 +1,3 @@ -"""Convert an mg7 numpy event file (``events.npy``) into an LHE file. +"""Redirect for the bin/ scripts of outputs made before the move to madgraph/mg7.""" -The mg7 output writes its events as npy by default. Anything that needs an LHE -file -- a shower, MadSpin, reweighting, an analysis -- normally switches the run -to the LHE format before the generation (``force_lhe_output_if_needed``). This -module covers the case where that request comes *after* the run: it rebuilds -the LHE file from what the npy run left behind. - -* ``output_format = "lhe_npy"``: the file already holds complete LHE events - (pdg ids, status, mothers, colour, helicity); they are only reformatted. -* ``output_format = "compact_npy"``: the file holds the momenta plus the - subprocess/diagram/colour/flavour/helicity indices the generation picked. - Those are completed exactly as the direct LHE output does it, with the - :class:`madspace.LHECompleter` the run saved next to the events - (``lhe_completer.json``). The completion draws random numbers (colour flow - among equivalent ones, resonance assignment), so the result is statistically - but not byte-identical to an ``output_format = "lhe"`` run with the same seed. - -In both cases the LHE header and ```` block (with the systematics -````) are taken verbatim from the ``header.lhe`` file the run writes -next to its npy events; the ``rwgt_`` columns become the per-event -```` block. The ```` reweighting-input block (``[systematics] -write_inputs``) is not reconstructed. - -Usage:: - - bin/npy_to_lhe [RUN_NAME | RUN_DIR | path/to/events.npy] [-o OUTPUT] [--seed N] -""" - -import argparse -import logging -import os -import re -import sys - -logger = logging.getLogger("madevent") - -COMPLETER_FILE = "lhe_completer.json" -HEADER_FILE = "header.lhe" -EVENT_FILE = "events.npy" - -# events converted per LHEFileWriter.write_string call -_CHUNK = 10000 - - -def _madspace(): - from madgraph.iolibs.template_files.mg7.bootstrap import ensure_madspace - ensure_madspace(interactive=False) - import madspace - return madspace - - -def _read_header(run_path): - """The header.lhe text without its closing tag, ready to be followed by - the events.""" - with open(os.path.join(run_path, HEADER_FILE)) as f: - text = f.read() - end = text.rfind("") - return text if end < 0 else text[:end] - - -def _seed(header, seed): - if seed is not None: - return int(seed) - match = re.search(r"\s*(-?\d+)\s*", header) - return int(match.group(1)) if match else 0 - - -# ---------------------------------------------------------------------------- -# conversion -# ---------------------------------------------------------------------------- -def resolve_event_file(path, me_dir=None) -> str: - """``path`` may be an npy file, a run directory or a run name (looked up - under ``/Events``).""" - candidates = [path, os.path.join(path, EVENT_FILE)] - if me_dir is not None: - candidates.append(os.path.join(me_dir, "Events", path, EVENT_FILE)) - for candidate in candidates: - if os.path.isfile(candidate): - return os.path.abspath(candidate) - raise FileNotFoundError("no %s found for '%s'" % (EVENT_FILE, path)) - - -def can_convert(run_path) -> bool: - """True when ``run_path`` holds npy events this module can turn into LHE.""" - return (os.path.isfile(os.path.join(run_path, EVENT_FILE)) - and os.path.isfile(os.path.join(run_path, HEADER_FILE))) - - -def _particle_count(names): - count = 0 - while "part%d_energy" % (count + 1) in names: - count += 1 - return count - - -def _compact_events(ms, events, completer, columns, rng): - names = events.dtype.names - nparticles = _particle_count(names) - momenta = [tuple(events["part%d_%s" % (i + 1, c)] - for c in ("energy", "px", "py", "pz")) - for i in range(nparticles)] - rwgt = [events[c] for c in columns] - for n in range(len(events)): - particles = [] - for energy, px, py, pz in momenta: - if energy[n] == 0.: - break - particles.append(ms.LHEParticle(px=px[n], py=py[n], pz=pz[n], - energy=energy[n])) - event = ms.LHEEvent(weight=events["weight"][n], - scale=events["ren_scale"][n], - alpha_qcd=events["alpha_qcd"][n], - particles=particles) - completer.complete_event_data( - event, int(events["subprocess_index"][n]), - int(events["diagram_index"][n]), int(events["color_index"][n]), - int(events["flavor_index"][n]), int(events["helicity_index"][n]), - rng) - yield event, [float(w[n]) for w in rwgt] - - -_LHE_PARTICLE_FIELDS = ("pdg_id", "status_code", "mother1", "mother2", "color", - "anti_color", "px", "py", "pz", "energy", "mass", - "lifetime", "spin") - - -def _lhe_npy_events(ms, events, columns): - names = events.dtype.names - count = 0 - while "part%d_pdg_id" % (count + 1) in names: - count += 1 - fields = [[events["part%d_%s" % (i + 1, f)] for f in _LHE_PARTICLE_FIELDS] - for i in range(count)] - rwgt = [events[c] for c in columns] - for n in range(len(events)): - particles = [] - for values in fields: - if values[0][n] == 0: - break - particles.append(ms.LHEParticle( - *[v[n].item() for v in values])) - event = ms.LHEEvent(int(events["process_id"][n]), - events["weight"][n], events["scale"][n], - events["alpha_qed"][n], events["alpha_qcd"][n], - particles) - yield event, [float(w[n]) for w in rwgt] - - -def convert(path, output=None, seed=None, compress=True, me_dir=None) -> str: - """Write the LHE file for the npy events at ``path`` (npy file, run - directory or run name) and return its path. By default it is written as - ``events.lhe.gz`` next to the npy file.""" - import numpy as np - from madgraph.various import misc - - ms = _madspace() - npy_path = resolve_event_file(path, me_dir) - run_path = os.path.dirname(npy_path) - # A standalone conversion starts in a fresh process, so it has not seen - # the gzip configuration installed by the event-generation launcher. - card_root = me_dir or os.path.dirname(os.path.dirname(run_path)) - card_path = os.path.join(card_root, "Cards", "run_card.toml") - if os.path.isfile(card_path): - from madgraph.various.banner import RunCardMG7 - misc.configure_gzip(RunCardMG7(card_path).gzip_configuration()) - if not os.path.isfile(os.path.join(run_path, HEADER_FILE)): - raise FileNotFoundError( - "%s is missing: the run predates the npy->LHE conversion support " - "and cannot be converted." % os.path.join(run_path, HEADER_FILE)) - header = _read_header(run_path) - - events = np.load(npy_path, mmap_mode="r") - names = events.dtype.names - columns = [n for n in names if re.fullmatch(r"rwgt_\d+", n)] - ids = [int(c.split("_", 1)[1]) for c in columns] - if "diagram_index" in names: - completer_path = os.path.join(run_path, COMPLETER_FILE) - if not os.path.isfile(completer_path): - raise FileNotFoundError("%s is missing" % completer_path) - completer = ms.LHECompleter.load(completer_path) - rng = ms.MixMaxRandom(_seed(header, seed)) - stream = lambda chunk: _compact_events(ms, chunk, completer, columns, rng) - elif "part1_pdg_id" in names: - stream = lambda chunk: _lhe_npy_events(ms, chunk, columns) - else: - raise ValueError("%s is not an mg7 event file" % npy_path) - - if output is None: - output = os.path.join(run_path, "events.lhe") - if output.endswith(".gz"): - output, compress = output[:-3], True - logger.info("converting %s (%d events) to LHE", npy_path, len(events)) - with open(output, "w") as out: - out.write(header) - for start in range(0, len(events), _CHUNK): - chunk = np.asarray(events[start:start + _CHUNK]) - text = [] - for event, weights in stream(chunk): - if ids: - event.rwgt_ids = ids - event.rwgt = weights - text.append(event.format()) - out.write("".join(text)) - out.write("\n") - if compress: - misc.gzip(output) - output += ".gz" - logger.info("LHE events written to %s", output) - return output - - -def main(argv=None, me_dir=None) -> None: - """Command-line entry point; run names are looked up under - ``/Events`` (default: the working directory).""" - parser = argparse.ArgumentParser( - description="Convert an mg7 events.npy file into an LHE file.") - parser.add_argument("path", nargs="?", default=None, - help="run name, run directory or events.npy file " - "(default: the most recent run with npy events " - "and no LHE file)") - parser.add_argument("-o", "--output", default=None, - help="output file (default: events.lhe.gz next to the " - "npy file)") - parser.add_argument("--seed", type=int, default=None, - help="seed of the LHE completion (default: the run seed)") - parser.add_argument("--no-gzip", action="store_false", dest="compress", - help="do not compress the output") - args = parser.parse_args(argv) - logging.basicConfig(level=logging.INFO, format="%(message)s") - - me_dir = me_dir or os.getcwd() - path = args.path - if path is None: - events_dir = os.path.join(me_dir, "Events") - runs = [os.path.join(events_dir, d) for d in - (os.listdir(events_dir) if os.path.isdir(events_dir) else [])] - runs = [d for d in runs if can_convert(d) and not any( - os.path.exists(os.path.join(d, n)) - for n in ("events.lhe", "events.lhe.gz"))] - if not runs: - sys.exit("no run with unconverted npy events in %s" % events_dir) - path = max(runs, key=lambda d: os.path.getmtime( - os.path.join(d, EVENT_FILE))) - try: - convert(path, args.output, args.seed, args.compress, me_dir) - except (FileNotFoundError, ValueError) as error: - sys.exit(str(error)) - - -if __name__ == "__main__": - main() +from madgraph.mg7.npy_to_lhe import main # noqa: F401 diff --git a/madgraph/mg7/__init__.py b/madgraph/mg7/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/madgraph/iolibs/template_files/mg7/bootstrap.py b/madgraph/mg7/bootstrap.py similarity index 100% rename from madgraph/iolibs/template_files/mg7/bootstrap.py rename to madgraph/mg7/bootstrap.py diff --git a/madgraph/iolibs/template_files/mg7/cli.py b/madgraph/mg7/cli.py similarity index 98% rename from madgraph/iolibs/template_files/mg7/cli.py rename to madgraph/mg7/cli.py index f1097793db..8ff98d6b5a 100644 --- a/madgraph/iolibs/template_files/mg7/cli.py +++ b/madgraph/mg7/cli.py @@ -6,14 +6,14 @@ import re import resource -from madgraph.iolibs.template_files.mg7.common import logger +from madgraph.mg7.common import logger from madgraph.interface.extended_cmd import Cmd from madgraph.various.banner import RunCardMG7 -from madgraph.iolibs.template_files.mg7 import npy_to_lhe -from madgraph.iolibs.template_files.mg7.driver import compute_decay_widths, force_lhe_output_if_needed, run_generation -from madgraph.iolibs.template_files.mg7.model_widths import compute_auto_widths, load_process_model -from madgraph.iolibs.template_files.mg7.options import lhapdf_paths, load_mg5_options -from madgraph.iolibs.template_files.mg7.postprocess import _off, _setup_logging +from madgraph.mg7 import npy_to_lhe +from madgraph.mg7.driver import compute_decay_widths, force_lhe_output_if_needed, run_generation +from madgraph.mg7.model_widths import compute_auto_widths, load_process_model +from madgraph.mg7.options import lhapdf_paths, load_mg5_options +from madgraph.mg7.postprocess import _off, _setup_logging class MG7Cmd(Cmd): """The mg7 run interface: the command object that drives a launch. diff --git a/madgraph/iolibs/template_files/mg7/common.py b/madgraph/mg7/common.py similarity index 99% rename from madgraph/iolibs/template_files/mg7/common.py rename to madgraph/mg7/common.py index 597a793bd8..3b59a4b445 100644 --- a/madgraph/iolibs/template_files/mg7/common.py +++ b/madgraph/mg7/common.py @@ -10,7 +10,7 @@ # Install madspace on first use, then import it. The bootstrap lives in its own # module (which must not import madspace) so that an in-process caller can run # it *before* importing this one -- see bootstrap.ensure_madspace. -from madgraph.iolibs.template_files.mg7.bootstrap import ( +from madgraph.mg7.bootstrap import ( ensure_madspace as _ensure_madspace, drop_install_path as _drop_install_path, MADSPACE_DIR as _MADSPACE_DIR, diff --git a/madgraph/iolibs/template_files/mg7/driver.py b/madgraph/mg7/driver.py similarity index 97% rename from madgraph/iolibs/template_files/mg7/driver.py rename to madgraph/mg7/driver.py index 91ccf5b29e..d23add81fc 100644 --- a/madgraph/iolibs/template_files/mg7/driver.py +++ b/madgraph/mg7/driver.py @@ -6,11 +6,11 @@ import re import shutil -from madgraph.iolibs.template_files.mg7.common import logger +from madgraph.mg7.common import logger from madgraph.various.banner import RunCardMG7 -from madgraph.iolibs.template_files.mg7.model_widths import compute_auto_widths -from madgraph.iolibs.template_files.mg7.postprocess import run_lhe_postprocessing, run_selected_tools, selected_tools -from madgraph.iolibs.template_files.mg7.madgraph_process import MadgraphProcess +from madgraph.mg7.model_widths import compute_auto_widths +from madgraph.mg7.postprocess import run_lhe_postprocessing, run_selected_tools, selected_tools +from madgraph.mg7.madgraph_process import MadgraphProcess def compute_decay_widths( me_dir: str = ".", precision: float = 0.01, nb_core: int | None = None, diff --git a/madgraph/iolibs/template_files/mg7/hwu_output.py b/madgraph/mg7/hwu_output.py similarity index 100% rename from madgraph/iolibs/template_files/mg7/hwu_output.py rename to madgraph/mg7/hwu_output.py diff --git a/madgraph/mg7/init_custom_func.py b/madgraph/mg7/init_custom_func.py new file mode 100644 index 0000000000..3133851e6d --- /dev/null +++ b/madgraph/mg7/init_custom_func.py @@ -0,0 +1,266 @@ +"""Create a skeleton for a custom cut, scale, PDF or bias function. + +The function is written for a single phase-space point in C, C++, Fortran, +CUDA or HIP. A generated UMAMI wrapper calls it on batches of events, so that +madspace can load it like a matrix element. The files go into +``custom_/`` in the process directory, and ``make`` there builds +``lib/custom_.so``. The run card option +``[custom_functions] `` takes the path of that library. + +Usage:: + + bin/init_custom_func [-f {cut,scale,pdf,bias}] [-l {c,cpp,fortran,cuda,hip}] + [-y] [--force] [--build | --no-build] + +Options that are not given are asked interactively, unless ``-y`` is passed. +""" + +import argparse +import logging +import os +import shutil +import subprocess +import sys + +logger = logging.getLogger("madevent") + +_TEMPLATE_FILES = os.path.join( + os.path.dirname(os.path.dirname(os.path.realpath(__file__))), "iolibs", "template_files" +) +TEMPLATE_DIR = os.path.join(_TEMPLATE_FILES, "mg7", "custom_functions") +UMAMI_HEADER = os.path.join(_TEMPLATE_FILES, "madmatrix", "umami.h") + +FUNCTIONS = { + "cut": (["UMAMI_OUT_CUT"], "a cut, true keeps the event and false rejects it"), + "scale": ( + ["UMAMI_OUT_REN_SCALE", "UMAMI_OUT_FACT_SCALE_1", "UMAMI_OUT_FACT_SCALE_2"], + "the renormalization and factorization scales", + ), + "pdf": (["UMAMI_OUT_PDF"], "a parton density x f(x, Q)"), + "bias": (["UMAMI_OUT_BIAS"], "an event bias for the unweighting"), +} +# example parameter of parameters.h: (name, default, run card value) +EXAMPLE_PARAMETERS = { + "cut": ("pt_min", "20.", "30.0"), + "scale": ("ht_fraction", "0.5", "0.25"), + "pdf": ("gluon_norm", "3.", "2.5"), + "bias": ("pt0", "100.", "200.0"), +} +# language: (body template extension, description) +LANGUAGES = { + "c": ("c", "C"), + "cpp": ("cpp", "C++"), + "fortran": ("F90", "Fortran"), + "cuda": ("cu", "CUDA, for runs on NVIDIA GPUs"), + "hip": ("cu", "HIP, for runs on AMD GPUs"), + "python-numpy": ("py", "Python with numpy, for CPU runs, no compilation"), + "python-torch": ("py", "Python with PyTorch, for CPU or GPU runs, no compilation"), +} + +_NONINTERACTIVE = (not sys.stdin.isatty()) or any(a in ("-y", "--yes") for a in sys.argv[1:]) + + +def _ask(question, default, choices): + """Ask through MG5's cmd.ask when available, else input().""" + if _NONINTERACTIVE: + return default + try: + from madgraph.interface.extended_cmd import Cmd + + answer = Cmd().ask(question, default, choices=list(choices), timeout=0) + return default if answer is None else str(answer) + except Exception: + pass + try: + answer = input(f"{question} [{default}]: ").strip() + except EOFError: + return default + return answer or default + + +def ask_choice(title, options, default): + """Numbered menu over ``options`` (name -> description), returns a name.""" + names = list(options) + lines = [title] + width = max(len(name) for name in names) + for i, name in enumerate(names, 1): + lines.append(f" {i}. {name:{width}s} {options[name]}") + print("\n".join(lines)) + choices = names + [str(i) for i in range(1, len(names) + 1)] + while True: + answer = _ask("Enter a number or a name", default, choices).strip().lower() + if answer in names: + return answer + if answer.isdigit() and 1 <= int(answer) <= len(names): + return names[int(answer) - 1] + print(f" Please enter one of {', '.join(names)} or 1-{len(names)}.") + + +def ask_yes_no(question, default=True): + default_s = "y" if default else "n" + while True: + answer = _ask(question, default_s, ["y", "n"]).strip().lower() + if answer in ("y", "yes"): + return True + if answer in ("n", "no"): + return False + print(" Please enter 'y' or 'n'.") + + +def build_parser(): + parser = argparse.ArgumentParser( + prog="bin/init_custom_func", + description=( + "Create a skeleton for a custom cut, scale, PDF or bias function in " + "custom_/. Compiled languages are built into " + "lib/custom_.so, Python functions need no build." + ), + ) + parser.add_argument("-f", "--function", choices=list(FUNCTIONS)) + parser.add_argument("-l", "--language", choices=list(LANGUAGES)) + parser.add_argument( + "-y", "--yes", action="store_true", + help="do not ask, use the defaults (cut, cpp) for options that are not given", + ) + parser.add_argument( + "--force", action="store_true", help="overwrite an existing custom_ directory" + ) + build = parser.add_mutually_exclusive_group() + build.add_argument("--build", dest="build", action="store_true", default=None, + help="run make after creating the files (default)") + build.add_argument("--no-build", dest="build", action="store_false") + run_card = parser.add_mutually_exclusive_group() + run_card.add_argument("--set-run-card", dest="set_run_card", action="store_true", + default=None, help="enable the function in Cards/run_card.toml (default)") + run_card.add_argument("--no-set-run-card", dest="set_run_card", action="store_false") + return parser + + +def is_python(language): + return language.startswith("python-") + + +def function_path(function, language): + """The run card path of the function, relative to the process directory.""" + if is_python(language): + return os.path.join(f"custom_{function}", f"custom_{function}.py") + return os.path.join("lib", f"custom_{function}.so") + + +def write_files(out_dir, function, language): + """Write the function body and, for compiled languages, the wrapper, + umami.h and the Makefile.""" + extension = LANGUAGES[language][0] + body_file = f"custom_{function}.{extension}" + if is_python(language): + framework = language.split("-", 1)[1] + os.makedirs(out_dir, exist_ok=True) + shutil.copy( + os.path.join(TEMPLATE_DIR, f"{function}_{framework}.py"), + os.path.join(out_dir, body_file), + ) + return body_file + values = { + "kind": function, + "function_name": f"custom_{function}", + "output_keys": ", ".join(FUNCTIONS[function][0]), + "is_pdf": int(function == "pdf"), + "is_scale": int(function == "scale"), + "return_type": "bool" if function == "cut" else "double", + "body_file": body_file, + "language": language, + "example": EXAMPLE_PARAMETERS[function][0], + "example_default": EXAMPLE_PARAMETERS[function][1], + "example_value": EXAMPLE_PARAMETERS[function][2], + } + os.makedirs(out_dir, exist_ok=True) + shutil.copy(os.path.join(TEMPLATE_DIR, f"{function}.{extension}"), os.path.join(out_dir, body_file)) + shutil.copy(UMAMI_HEADER, os.path.join(out_dir, "umami.h")) + wrapper = "wrapper_gpu.cu" if language in ("cuda", "hip") else "wrapper_cpu.c" + for template in (wrapper, "Makefile", "parameters.h"): + with open(os.path.join(TEMPLATE_DIR, template)) as f: + text = f.read() % values + with open(os.path.join(out_dir, template), "w") as f: + f.write(text) + return body_file + + +def build(out_dir): + """Run make in ``out_dir``, logging to compile.log. Returns True on success.""" + log_path = os.path.join(out_dir, "compile.log") + with open(log_path, "w") as log: + result = subprocess.run(["make"], cwd=out_dir, stdout=log, stderr=subprocess.STDOUT) + if result.returncode != 0: + with open(log_path) as log: + print(log.read()) + print(f"Building failed, see {log_path}") + return False + return True + + +def set_run_card(me_dir, function, lib_path): + from madgraph.various.banner import RunCardMG7 + + card_path = os.path.join(me_dir, "Cards", "run_card.toml") + run_card = RunCardMG7(card_path) + run_card.set(f"custom_functions.{function}", lib_path) + run_card.write(card_path) + + +def main(me_dir=None, argv=None): + global _NONINTERACTIVE + args = build_parser().parse_args(argv) + if args.yes: + _NONINTERACTIVE = True + me_dir = me_dir or os.getcwd() + + function = args.function or ask_choice( + "Which custom function do you want to create?", + {name: desc for name, (_, desc) in FUNCTIONS.items()}, + "cut", + ) + language = args.language or ask_choice( + "Which language do you want to write it in?", + {name: desc for name, (_, desc) in LANGUAGES.items()}, + "cpp", + ) + + out_dir = os.path.join(me_dir, f"custom_{function}") + if os.path.exists(out_dir): + if not (args.force or (not _NONINTERACTIVE and ask_yes_no( + f"{out_dir} already exists. Overwrite it?", False))): + print(f"{out_dir} already exists, use --force to overwrite it.") + return 1 + shutil.rmtree(out_dir) + + body_file = write_files(out_dir, function, language) + lib_path = function_path(function, language) + print(f"Created {os.path.relpath(out_dir, me_dir)}/. Write your function in " + f"{os.path.relpath(os.path.join(out_dir, body_file), me_dir)}.") + + do_build = args.build + if is_python(language): + if do_build: + print("Python functions need no build.") + do_build = False + elif do_build is None: + do_build = ask_yes_no("Build the default function now?", True) + if do_build and not build(out_dir): + return 1 + + do_set = args.set_run_card + if do_set is None: + do_set = ask_yes_no(f"Set custom_functions.{function} = \"{lib_path}\" in the run card?", True) + if do_set: + set_run_card(me_dir, function, lib_path) + print(f"Enabled it in Cards/run_card.toml: custom_functions.{function} = \"{lib_path}\".") + else: + print(f"To enable it, set custom_functions.{function} = \"{lib_path}\" in Cards/run_card.toml.") + if not do_build and not is_python(language): + print(f"Build it with `make -C {os.path.relpath(out_dir, me_dir)}`. " + "bin/generate_events also rebuilds it before a run.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/madgraph/mg7/launch.py b/madgraph/mg7/launch.py new file mode 100644 index 0000000000..a8f289c3de --- /dev/null +++ b/madgraph/mg7/launch.py @@ -0,0 +1,77 @@ +"""Entry point of the mg7 launcher, re-exporting its modules.""" + +from madgraph.mg7.common import * # noqa: F401,F403 +from madgraph.mg7.common import ms, logger # noqa: F401 +from madgraph.mg7.observables import ( # noqa: F401 + CutItem, + HistItem, + _MERGED_PID_REPRESENTATIVE, + clean_pids, + _SIGNED_OBSERVABLES, + _flavor_options, + observable_pids, + build_observable, + decay_products, + pid_is_qcd, +) +from madgraph.mg7.subprocess_madmatrix import ( # noqa: F401 + Channel, + PhaseSpace, + MultiChannelData, + build_topologies, + build_multi_channel_data, + MadgraphSubprocess, +) +from madgraph.mg7.madgraph_process import ( # noqa: F401 + MadgraphProcess, + PCH_MIN_SUBPROCESSES, + _single_flavor_metas, +) +from madgraph.mg7.options import ( # noqa: F401 + _ROOTED_OPTIONS, + load_mg5_options, + lhapdf_paths, +) +from madgraph.mg7.model_widths import ( # noqa: F401 + _model_hash_warned, + _model_cache, + read_stored_model, + _proc_characteristic, + load_process_model, + compute_auto_widths, +) +from madgraph.mg7.postprocess import ( # noqa: F401 + _find_event_file, + _find_or_convert_event_file, + _report_failure, + _off, + _has_effective_handler, + _TOOL_LOGGING_READY, + _setup_logging, + selected_tools, + run_selected_tools, + _add_time_of_flight, + _lhapdf_config_path, + _run_systematics, + run_lhe_postprocessing, +) +from madgraph.mg7.driver import ( # noqa: F401 + compute_decay_widths, + _decay_flavor_options, + _release_channel_generators, + run_single, + detect_run_scan, + detect_param_scan, + run_scan, + run_generation, + _lhe_needed_by, + force_lhe_output_if_needed, +) +from madgraph.mg7.cli import ( # noqa: F401 + MG7Cmd, + build_selector_cmd, + ask_edit_cards, + prune_unselected_tool_cards, + _raise_open_file_limit, + main, +) diff --git a/madgraph/iolibs/template_files/mg7/madgraph_process.py b/madgraph/mg7/madgraph_process.py similarity index 98% rename from madgraph/iolibs/template_files/mg7/madgraph_process.py rename to madgraph/mg7/madgraph_process.py index af28b667fb..499264b311 100644 --- a/madgraph/iolibs/template_files/mg7/madgraph_process.py +++ b/madgraph/mg7/madgraph_process.py @@ -9,23 +9,23 @@ import tempfile import time -from madgraph.iolibs.template_files.mg7.common import backend_of, fptype_of, get_start_time, interference_helicities_summed, lhe_helicities, lhe_weight_info, logger, me_parameters, ms, resolve_auto_backend, resolve_seed, resolve_verbosity +from madgraph.mg7.common import backend_of, fptype_of, get_start_time, interference_helicities_summed, lhe_helicities, lhe_weight_info, logger, me_parameters, ms, resolve_auto_backend, resolve_seed, resolve_verbosity from madgraph.various import misc from madgraph.various.banner import RunCardMG7 from madgraph.models.check_param_card import ParamCard from madgraph.models.model_reader import Alphas_Runner -from madgraph.iolibs.template_files.mg7 import hwu_output -from madgraph.iolibs.template_files.mg7 import npy_to_lhe -from madgraph.iolibs.template_files.mg7 import options -from madgraph.iolibs.template_files.mg7 import plots -from madgraph.iolibs.template_files.mg7 import systematics_summary -from madgraph.iolibs.template_files.mg7.bootstrap import INSTALL_DIR as _INSTALL_DIR -from madgraph.iolibs.template_files.mg7.bootstrap import MADSPACE_DIR as _MADSPACE_DIR +from madgraph.mg7 import hwu_output +from madgraph.mg7 import npy_to_lhe +from madgraph.mg7 import options +from madgraph.mg7 import plots +from madgraph.mg7 import systematics_summary +from madgraph.mg7.bootstrap import INSTALL_DIR as _INSTALL_DIR +from madgraph.mg7.bootstrap import MADSPACE_DIR as _MADSPACE_DIR from madgraph.iolibs.template_files.mg7.custom_function_loader import load_custom_function -from madgraph.iolibs.template_files.mg7.observables import CutItem, HistItem, build_observable, clean_pids, observable_pids -from madgraph.iolibs.template_files.mg7.options import lhapdf_paths -from madgraph.iolibs.template_files.mg7.subprocess_amplicol import AmplicolSubprocess, amplicol_subproc_args, is_amplicol -from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import MadgraphSubprocess, PhaseSpace, build_multi_channel_data +from madgraph.mg7.observables import CutItem, HistItem, build_observable, clean_pids, observable_pids +from madgraph.mg7.options import lhapdf_paths +from madgraph.mg7.subprocess_amplicol import AmplicolSubprocess, amplicol_subproc_args, is_amplicol +from madgraph.mg7.subprocess_madmatrix import MadgraphSubprocess, PhaseSpace, build_multi_channel_data # From this many subprocess directories to compile, the standard headers are # precompiled once ('make PCH=1', see madmatrix.mk): each file of a small process @@ -2035,7 +2035,8 @@ def save_gridpack(self) -> None: bin_path = os.path.join(gridpack_path, "bin") os.mkdir(bin_path) gen_events_file = os.path.join(bin_path, "generate_events") - template_dir = os.path.dirname(__file__) + template_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), + "iolibs", "template_files", "mg7") shutil.copy(os.path.join(template_dir, "gridpack.py"), gen_events_file) shutil.copy(os.path.join(template_dir, "custom_function_loader.py"), bin_path) os.chmod(gen_events_file, 0o755) diff --git a/madgraph/iolibs/template_files/mg7/model_widths.py b/madgraph/mg7/model_widths.py similarity index 98% rename from madgraph/iolibs/template_files/mg7/model_widths.py rename to madgraph/mg7/model_widths.py index fb3d01c7c8..bd37ee210d 100644 --- a/madgraph/iolibs/template_files/mg7/model_widths.py +++ b/madgraph/mg7/model_widths.py @@ -7,7 +7,7 @@ from madgraph.models.check_param_card import group_auto_widths from madgraph.various import misc -from madgraph.iolibs.template_files.mg7.bootstrap import MG_ROOT as _MG_ROOT +from madgraph.mg7.bootstrap import MG_ROOT as _MG_ROOT logger = logging.getLogger("madgraph7") diff --git a/madgraph/mg7/npy_to_lhe.py b/madgraph/mg7/npy_to_lhe.py new file mode 100644 index 0000000000..d7d810097d --- /dev/null +++ b/madgraph/mg7/npy_to_lhe.py @@ -0,0 +1,252 @@ +"""Convert an mg7 numpy event file (``events.npy``) into an LHE file. + +The mg7 output writes its events as npy by default. Anything that needs an LHE +file -- a shower, MadSpin, reweighting, an analysis -- normally switches the run +to the LHE format before the generation (``force_lhe_output_if_needed``). This +module covers the case where that request comes *after* the run: it rebuilds +the LHE file from what the npy run left behind. + +* ``output_format = "lhe_npy"``: the file already holds complete LHE events + (pdg ids, status, mothers, colour, helicity); they are only reformatted. +* ``output_format = "compact_npy"``: the file holds the momenta plus the + subprocess/diagram/colour/flavour/helicity indices the generation picked. + Those are completed exactly as the direct LHE output does it, with the + :class:`madspace.LHECompleter` the run saved next to the events + (``lhe_completer.json``). The completion draws random numbers (colour flow + among equivalent ones, resonance assignment), so the result is statistically + but not byte-identical to an ``output_format = "lhe"`` run with the same seed. + +In both cases the LHE header and ```` block (with the systematics +````) are taken verbatim from the ``header.lhe`` file the run writes +next to its npy events; the ``rwgt_`` columns become the per-event +```` block. The ```` reweighting-input block (``[systematics] +write_inputs``) is not reconstructed. + +Usage:: + + bin/npy_to_lhe [RUN_NAME | RUN_DIR | path/to/events.npy] [-o OUTPUT] [--seed N] +""" + +import argparse +import logging +import os +import re +import sys + +logger = logging.getLogger("madevent") + +COMPLETER_FILE = "lhe_completer.json" +HEADER_FILE = "header.lhe" +EVENT_FILE = "events.npy" + +# events converted per LHEFileWriter.write_string call +_CHUNK = 10000 + + +def _madspace(): + from madgraph.mg7.bootstrap import ensure_madspace + ensure_madspace(interactive=False) + import madspace + return madspace + + +def _read_header(run_path): + """The header.lhe text without its closing tag, ready to be followed by + the events.""" + with open(os.path.join(run_path, HEADER_FILE)) as f: + text = f.read() + end = text.rfind("") + return text if end < 0 else text[:end] + + +def _seed(header, seed): + if seed is not None: + return int(seed) + match = re.search(r"\s*(-?\d+)\s*", header) + return int(match.group(1)) if match else 0 + + +# ---------------------------------------------------------------------------- +# conversion +# ---------------------------------------------------------------------------- +def resolve_event_file(path, me_dir=None) -> str: + """``path`` may be an npy file, a run directory or a run name (looked up + under ``/Events``).""" + candidates = [path, os.path.join(path, EVENT_FILE)] + if me_dir is not None: + candidates.append(os.path.join(me_dir, "Events", path, EVENT_FILE)) + for candidate in candidates: + if os.path.isfile(candidate): + return os.path.abspath(candidate) + raise FileNotFoundError("no %s found for '%s'" % (EVENT_FILE, path)) + + +def can_convert(run_path) -> bool: + """True when ``run_path`` holds npy events this module can turn into LHE.""" + return (os.path.isfile(os.path.join(run_path, EVENT_FILE)) + and os.path.isfile(os.path.join(run_path, HEADER_FILE))) + + +def _particle_count(names): + count = 0 + while "part%d_energy" % (count + 1) in names: + count += 1 + return count + + +def _compact_events(ms, events, completer, columns, rng): + names = events.dtype.names + nparticles = _particle_count(names) + momenta = [tuple(events["part%d_%s" % (i + 1, c)] + for c in ("energy", "px", "py", "pz")) + for i in range(nparticles)] + rwgt = [events[c] for c in columns] + for n in range(len(events)): + particles = [] + for energy, px, py, pz in momenta: + if energy[n] == 0.: + break + particles.append(ms.LHEParticle(px=px[n], py=py[n], pz=pz[n], + energy=energy[n])) + event = ms.LHEEvent(weight=events["weight"][n], + scale=events["ren_scale"][n], + alpha_qcd=events["alpha_qcd"][n], + particles=particles) + completer.complete_event_data( + event, int(events["subprocess_index"][n]), + int(events["diagram_index"][n]), int(events["color_index"][n]), + int(events["flavor_index"][n]), int(events["helicity_index"][n]), + rng) + yield event, [float(w[n]) for w in rwgt] + + +_LHE_PARTICLE_FIELDS = ("pdg_id", "status_code", "mother1", "mother2", "color", + "anti_color", "px", "py", "pz", "energy", "mass", + "lifetime", "spin") + + +def _lhe_npy_events(ms, events, columns): + names = events.dtype.names + count = 0 + while "part%d_pdg_id" % (count + 1) in names: + count += 1 + fields = [[events["part%d_%s" % (i + 1, f)] for f in _LHE_PARTICLE_FIELDS] + for i in range(count)] + rwgt = [events[c] for c in columns] + for n in range(len(events)): + particles = [] + for values in fields: + if values[0][n] == 0: + break + particles.append(ms.LHEParticle( + *[v[n].item() for v in values])) + event = ms.LHEEvent(int(events["process_id"][n]), + events["weight"][n], events["scale"][n], + events["alpha_qed"][n], events["alpha_qcd"][n], + particles) + yield event, [float(w[n]) for w in rwgt] + + +def convert(path, output=None, seed=None, compress=True, me_dir=None) -> str: + """Write the LHE file for the npy events at ``path`` (npy file, run + directory or run name) and return its path. By default it is written as + ``events.lhe.gz`` next to the npy file.""" + import numpy as np + from madgraph.various import misc + + ms = _madspace() + npy_path = resolve_event_file(path, me_dir) + run_path = os.path.dirname(npy_path) + # A standalone conversion starts in a fresh process, so it has not seen + # the gzip configuration installed by the event-generation launcher. + card_root = me_dir or os.path.dirname(os.path.dirname(run_path)) + card_path = os.path.join(card_root, "Cards", "run_card.toml") + if os.path.isfile(card_path): + from madgraph.various.banner import RunCardMG7 + misc.configure_gzip(RunCardMG7(card_path).gzip_configuration()) + if not os.path.isfile(os.path.join(run_path, HEADER_FILE)): + raise FileNotFoundError( + "%s is missing: the run predates the npy->LHE conversion support " + "and cannot be converted." % os.path.join(run_path, HEADER_FILE)) + header = _read_header(run_path) + + events = np.load(npy_path, mmap_mode="r") + names = events.dtype.names + columns = [n for n in names if re.fullmatch(r"rwgt_\d+", n)] + ids = [int(c.split("_", 1)[1]) for c in columns] + if "diagram_index" in names: + completer_path = os.path.join(run_path, COMPLETER_FILE) + if not os.path.isfile(completer_path): + raise FileNotFoundError("%s is missing" % completer_path) + completer = ms.LHECompleter.load(completer_path) + rng = ms.MixMaxRandom(_seed(header, seed)) + stream = lambda chunk: _compact_events(ms, chunk, completer, columns, rng) + elif "part1_pdg_id" in names: + stream = lambda chunk: _lhe_npy_events(ms, chunk, columns) + else: + raise ValueError("%s is not an mg7 event file" % npy_path) + + if output is None: + output = os.path.join(run_path, "events.lhe") + if output.endswith(".gz"): + output, compress = output[:-3], True + logger.info("converting %s (%d events) to LHE", npy_path, len(events)) + with open(output, "w") as out: + out.write(header) + for start in range(0, len(events), _CHUNK): + chunk = np.asarray(events[start:start + _CHUNK]) + text = [] + for event, weights in stream(chunk): + if ids: + event.rwgt_ids = ids + event.rwgt = weights + text.append(event.format()) + out.write("".join(text)) + out.write("\n") + if compress: + misc.gzip(output) + output += ".gz" + logger.info("LHE events written to %s", output) + return output + + +def main(argv=None, me_dir=None) -> None: + """Command-line entry point; run names are looked up under + ``/Events`` (default: the working directory).""" + parser = argparse.ArgumentParser( + description="Convert an mg7 events.npy file into an LHE file.") + parser.add_argument("path", nargs="?", default=None, + help="run name, run directory or events.npy file " + "(default: the most recent run with npy events " + "and no LHE file)") + parser.add_argument("-o", "--output", default=None, + help="output file (default: events.lhe.gz next to the " + "npy file)") + parser.add_argument("--seed", type=int, default=None, + help="seed of the LHE completion (default: the run seed)") + parser.add_argument("--no-gzip", action="store_false", dest="compress", + help="do not compress the output") + args = parser.parse_args(argv) + logging.basicConfig(level=logging.INFO, format="%(message)s") + + me_dir = me_dir or os.getcwd() + path = args.path + if path is None: + events_dir = os.path.join(me_dir, "Events") + runs = [os.path.join(events_dir, d) for d in + (os.listdir(events_dir) if os.path.isdir(events_dir) else [])] + runs = [d for d in runs if can_convert(d) and not any( + os.path.exists(os.path.join(d, n)) + for n in ("events.lhe", "events.lhe.gz"))] + if not runs: + sys.exit("no run with unconverted npy events in %s" % events_dir) + path = max(runs, key=lambda d: os.path.getmtime( + os.path.join(d, EVENT_FILE))) + try: + convert(path, args.output, args.seed, args.compress, me_dir) + except (FileNotFoundError, ValueError) as error: + sys.exit(str(error)) + + +if __name__ == "__main__": + main() diff --git a/madgraph/iolibs/template_files/mg7/observables.py b/madgraph/mg7/observables.py similarity index 99% rename from madgraph/iolibs/template_files/mg7/observables.py rename to madgraph/mg7/observables.py index 27523e7d73..5c55cc57be 100644 --- a/madgraph/iolibs/template_files/mg7/observables.py +++ b/madgraph/mg7/observables.py @@ -2,7 +2,7 @@ from dataclasses import dataclass -from madgraph.iolibs.template_files.mg7.common import ms +from madgraph.mg7.common import ms @dataclass class CutItem: diff --git a/madgraph/iolibs/template_files/mg7/options.py b/madgraph/mg7/options.py similarity index 100% rename from madgraph/iolibs/template_files/mg7/options.py rename to madgraph/mg7/options.py diff --git a/madgraph/iolibs/template_files/mg7/plots.py b/madgraph/mg7/plots.py similarity index 100% rename from madgraph/iolibs/template_files/mg7/plots.py rename to madgraph/mg7/plots.py diff --git a/madgraph/iolibs/template_files/mg7/postprocess.py b/madgraph/mg7/postprocess.py similarity index 98% rename from madgraph/iolibs/template_files/mg7/postprocess.py rename to madgraph/mg7/postprocess.py index 2099f6112e..55c891bf2a 100644 --- a/madgraph/iolibs/template_files/mg7/postprocess.py +++ b/madgraph/mg7/postprocess.py @@ -5,8 +5,8 @@ import sys import time -from madgraph.iolibs.template_files.mg7 import npy_to_lhe -from madgraph.iolibs.template_files.mg7.options import lhapdf_paths, load_mg5_options +from madgraph.mg7 import npy_to_lhe +from madgraph.mg7.options import lhapdf_paths, load_mg5_options def _find_event_file(run_path): """Locate the LHE event file produced by generate_events (if any).""" @@ -162,7 +162,7 @@ def run_selected_tools(switch, process) -> None: log.info("Post-processing the generated events with: %s", ", ".join(tools)) try: - from madgraph.iolibs.template_files.mg7.run_interface import MG7RunCmd + from madgraph.mg7.run_interface import MG7RunCmd cmd = MG7RunCmd(os.getcwd(), load_mg5_options(), run_name, lhe_path) cmd.set_run_name(run_name, None, "parton") except Exception as error: diff --git a/madgraph/iolibs/template_files/mg7/run_interface.py b/madgraph/mg7/run_interface.py similarity index 99% rename from madgraph/iolibs/template_files/mg7/run_interface.py rename to madgraph/mg7/run_interface.py index 407ba0ff6a..6e541e800b 100644 --- a/madgraph/iolibs/template_files/mg7/run_interface.py +++ b/madgraph/mg7/run_interface.py @@ -261,7 +261,7 @@ def get_model(self): tools this class drives (MadSpin, reweight, the auto widths and the 'update dependent' pass of their card questions) work here too. """ - from madgraph.iolibs.template_files.mg7.model_widths import load_process_model + from madgraph.mg7.model_widths import load_process_model return load_process_model(self.me_dir) def do_set(self, line, log=True): diff --git a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py b/madgraph/mg7/subprocess_amplicol.py similarity index 98% rename from madgraph/iolibs/template_files/mg7/subprocess_amplicol.py rename to madgraph/mg7/subprocess_amplicol.py index d1bfa96372..cbc2277ad9 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py +++ b/madgraph/mg7/subprocess_amplicol.py @@ -13,16 +13,16 @@ from typing import TYPE_CHECKING -from madgraph.iolibs.template_files.mg7.common import lhe_helicities, ms -from madgraph.iolibs.template_files.mg7.observables import clean_pids -from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import ( +from madgraph.mg7.common import lhe_helicities, ms +from madgraph.mg7.observables import clean_pids +from madgraph.mg7.subprocess_madmatrix import ( Channel, MadgraphSubprocess, PhaseSpace, ) if TYPE_CHECKING: - from madgraph.iolibs.template_files.mg7.madgraph_process import MadgraphProcess + from madgraph.mg7.madgraph_process import MadgraphProcess def is_amplicol(meta: dict) -> bool: diff --git a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py b/madgraph/mg7/subprocess_madmatrix.py similarity index 99% rename from madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py rename to madgraph/mg7/subprocess_madmatrix.py index ef419b8783..6bbeebd6d9 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py +++ b/madgraph/mg7/subprocess_madmatrix.py @@ -7,12 +7,12 @@ from dataclasses import dataclass, field from typing import Literal, NamedTuple, TYPE_CHECKING -from madgraph.iolibs.template_files.mg7.common import logger, me_parameters, ms +from madgraph.mg7.common import logger, me_parameters, ms from madgraph.iolibs.template_files.mg7.custom_function_loader import load_surrogate -from madgraph.iolibs.template_files.mg7.observables import build_observable, clean_pids, decay_products, observable_pids, pid_is_qcd +from madgraph.mg7.observables import build_observable, clean_pids, decay_products, observable_pids, pid_is_qcd if TYPE_CHECKING: - from madgraph.iolibs.template_files.mg7.madgraph_process import MadgraphProcess + from madgraph.mg7.madgraph_process import MadgraphProcess @dataclass class Channel: diff --git a/madgraph/iolibs/template_files/mg7/systematics_summary.py b/madgraph/mg7/systematics_summary.py similarity index 100% rename from madgraph/iolibs/template_files/mg7/systematics_summary.py rename to madgraph/mg7/systematics_summary.py diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 0dcedf277d..8b79aa91b1 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -8201,7 +8201,7 @@ def _parameter_value(model, name): # MadGraph gives a leg whose flavours were merged one of these codes; the # mg7 runtime resolves them to the same representatives (clean_pids / - # _MERGED_PID_REPRESENTATIVE in template_files/mg7/launch.py), so the + # _MERGED_PID_REPRESENTATIVE in madgraph/mg7/launch.py), so the # observables have to be written in terms of those. merged_pid_representative = {81: 1, 82: 11, 83: 12} diff --git a/tests/acceptance_tests/test_auto_width.py b/tests/acceptance_tests/test_auto_width.py index 9dd218d631..e8ce737276 100644 --- a/tests/acceptance_tests/test_auto_width.py +++ b/tests/acceptance_tests/test_auto_width.py @@ -32,7 +32,7 @@ import madgraph.interface.master_interface as mgcmd import madgraph.models.check_param_card as check_param_card -from madgraph.iolibs.template_files.mg7 import bootstrap as mg7_bootstrap +from madgraph.mg7 import bootstrap as mg7_bootstrap pjoin = os.path.join @@ -214,7 +214,7 @@ def test_options_of_the_widths(self): with open(card, 'w') as stream: stream.write(text) - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch os.chdir(run_dir) mg7_launch.compute_auto_widths(card) os.chdir(self.cwd) diff --git a/tests/acceptance_tests/test_cmd_madevent.py b/tests/acceptance_tests/test_cmd_madevent.py index 3fef695638..fbd1b8cea2 100755 --- a/tests/acceptance_tests/test_cmd_madevent.py +++ b/tests/acceptance_tests/test_cmd_madevent.py @@ -1523,7 +1523,7 @@ def test_relaunch_switch_defaults_mg7(self): mg.exec_cmd('output mg7 %s' % out_dir) launcher = importlib.import_module( - 'madgraph.iolibs.template_files.mg7.launch') + 'madgraph.mg7.launch') tool_cards = ('pythia8_card.dat', 'madspin_card.dat', 'delphes_card.dat', 'reweight_card.dat', 'rivet_card.dat', 'madanalysis5_parton_card.dat', diff --git a/tests/acceptance_tests/test_mg7_amplicol.py b/tests/acceptance_tests/test_mg7_amplicol.py index 5204ce520f..598fe0f0aa 100644 --- a/tests/acceptance_tests/test_mg7_amplicol.py +++ b/tests/acceptance_tests/test_mg7_amplicol.py @@ -112,7 +112,7 @@ def _resonances(lhe_path): class MG7AmplicolTest(unittest.TestCase): def setUp(self): - from madgraph.iolibs.template_files.mg7 import bootstrap + from madgraph.mg7 import bootstrap # not imported: Symbolica would hold its lock against the pyamplicol worker missing = [name for name, installed in ( ('madspace', bootstrap.madspace_is_installed()), diff --git a/tests/acceptance_tests/test_mg7_custom_functions.py b/tests/acceptance_tests/test_mg7_custom_functions.py index 72fc3d0ed0..9fbcb8034d 100644 --- a/tests/acceptance_tests/test_mg7_custom_functions.py +++ b/tests/acceptance_tests/test_mg7_custom_functions.py @@ -208,7 +208,7 @@ def _import_madspace(): """The madspace of this MadGraph, as bin/generate_events imports it, or None. Done before anything imports another madspace from the path.""" try: - from madgraph.iolibs.template_files.mg7.bootstrap import ( + from madgraph.mg7.bootstrap import ( ensure_madspace, drop_install_path) ensure_madspace(interactive=False) import madspace @@ -508,7 +508,7 @@ def test_init_custom_func_languages_mg7(self): self.skipTest('madspace with custom functions unavailable') compilers = {'c': 'cc', 'cpp': 'c++', 'fortran': 'gfortran', 'cuda': 'nvcc', 'hip': 'hipcc'} - from madgraph.iolibs.template_files.mg7 import init_custom_func + from madgraph.mg7 import init_custom_func tested = [] for language, compiler in compilers.items(): diff --git a/tests/unit_tests/interface/test_auto_width_engine.py b/tests/unit_tests/interface/test_auto_width_engine.py index 561a63b019..3021a97bd8 100644 --- a/tests/unit_tests/interface/test_auto_width_engine.py +++ b/tests/unit_tests/interface/test_auto_width_engine.py @@ -36,7 +36,7 @@ import madgraph.core.base_objects as base_objects import madgraph.interface.madevent_interface as madevent_interface import madgraph.interface.master_interface as mgcmd -from madgraph.iolibs.template_files.mg7 import bootstrap as mg7_bootstrap +from madgraph.mg7 import bootstrap as mg7_bootstrap from madgraph.models.model_reader import Alphas_Runner pjoin = os.path.join @@ -445,7 +445,7 @@ def tearDown(self): shutil.rmtree(self.tmp, ignore_errors=True) def write_grid(self, mass, alpha_s_mz=0.118): - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch process = object.__new__(mg7_launch.MadgraphProcess) process.param_card = mock.Mock() process.param_card.get_value = lambda block, code: alpha_s_mz @@ -544,7 +544,7 @@ class DecayFlavorSplitTest(unittest.TestCase): } def setUp(self): - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch self.launch = mg7_launch def test_one_subprocess_per_flavour(self): diff --git a/tests/unit_tests/interface/test_auto_width_nlo.py b/tests/unit_tests/interface/test_auto_width_nlo.py index 0e42b18993..bc20c37be8 100644 --- a/tests/unit_tests/interface/test_auto_width_nlo.py +++ b/tests/unit_tests/interface/test_auto_width_nlo.py @@ -42,7 +42,7 @@ import madgraph.models.check_param_card as check_param_card import madgraph.models.import_ufo as import_ufo import madgraph.models.model_reader as model_reader -from madgraph.iolibs.template_files.mg7 import bootstrap as mg7_bootstrap +from madgraph.mg7 import bootstrap as mg7_bootstrap pjoin = os.path.join @@ -858,8 +858,8 @@ def tearDown(self): shutil.rmtree(self.tmp, ignore_errors=True) def test_one_compute_widths_per_set_of_options(self): - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch - from madgraph.iolibs.template_files.mg7 import model_widths + from madgraph.mg7 import launch as mg7_launch + from madgraph.mg7 import model_widths card = pjoin(self.tmp, 'param_card.dat') with open(card, 'w') as stream: stream.write('\n'.join([ @@ -885,7 +885,7 @@ def run(command, *args, **kwargs): 'compute_widths 24 --path=CARD ']) def test_no_auto_no_run(self): - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch card = pjoin(self.tmp, 'param_card.dat') with open(card, 'w') as stream: stream.write('DECAY 6 1.5\n') @@ -895,7 +895,7 @@ def test_no_auto_no_run(self): def compile_args(self, subprocesses, extra=()): """The arguments of the make that builds ``subprocesses`` directories.""" - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch process = object.__new__(mg7_launch.MadgraphProcess) process.run_card = {'run': {'device': ['cpu'], 'cpu_mode': 'simd_128', 'precision': 'color32', 'cpu_thread_pool_size': 2}} @@ -919,7 +919,7 @@ def test_the_headers_are_precompiled_from_four_subprocesses(self): def test_only_what_is_to_build_counts(self): """Six directories, but the libraries of three of them exist: below the threshold.""" - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch for i in range(3): open(pjoin(self.tmp, 'libmadmatrix_P%d_x_simd_128.so' % i), 'w').close() process = object.__new__(mg7_launch.MadgraphProcess) diff --git a/tests/unit_tests/interface/test_mg7_interference_helicity.py b/tests/unit_tests/interface/test_mg7_interference_helicity.py index 60c4cf2a71..2354a86bdd 100644 --- a/tests/unit_tests/interface/test_mg7_interference_helicity.py +++ b/tests/unit_tests/interface/test_mg7_interference_helicity.py @@ -46,7 +46,7 @@ class TestInterferenceHelicity(unittest.TestCase): def setUp(self): try: - import madgraph.iolibs.template_files.mg7.launch as launch + import madgraph.mg7.launch as launch except ImportError as error: # the mg7 runtime (madspace) is not there self.skipTest('mg7 launch not importable: %s' % error) self.launch = launch diff --git a/tests/unit_tests/interface/test_mg7_launch.py b/tests/unit_tests/interface/test_mg7_launch.py index d13ceca925..4501bed310 100644 --- a/tests/unit_tests/interface/test_mg7_launch.py +++ b/tests/unit_tests/interface/test_mg7_launch.py @@ -41,9 +41,9 @@ import madgraph.interface.extended_cmd as ext_cmd import madgraph.interface.master_interface as mgcmd -from madgraph.iolibs.template_files.mg7 import bootstrap as mg7_bootstrap +from madgraph.mg7 import bootstrap as mg7_bootstrap -LAUNCH_MODULE = 'madgraph.iolibs.template_files.mg7.launch' +LAUNCH_MODULE = 'madgraph.mg7.launch' _MG_ROOT = str(mg7_bootstrap.MG_ROOT) @@ -116,7 +116,7 @@ def MG7Cmd(me_dir='.', options=None): # `from pkg import launch` resolves through the parent package's # attribute when the real module has already been imported (which # MG7CmdTest does), so sys.modules alone is not enough to stub it. - import madgraph.iolibs.template_files.mg7 as mg7_pkg + import madgraph.mg7 as mg7_pkg self.mg7_pkg = mg7_pkg self._saved_attr = getattr(mg7_pkg, 'launch', None) mg7_pkg.launch = stub @@ -480,7 +480,7 @@ class MG7CmdTest(unittest.TestCase): """The real run interface (needs a compiled madspace to import).""" def setUp(self): - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch self.launch = mg7_launch self.tmpdir = tempfile.TemporaryDirectory() self.me_dir = make_mg7_dir(os.path.join(self.tmpdir.name, 'PROC')) @@ -633,7 +633,7 @@ def test_no_histograms_no_plot_directory(self): def test_missing_matplotlib_does_not_fail_the_run(self): """the numbers are in info.json: a run without the backend carries on""" - from madgraph.iolibs.template_files.mg7 import plots + from madgraph.mg7 import plots saved = plots._pyplot def no_matplotlib(): @@ -990,7 +990,7 @@ def test_setup_logging_defers_to_an_existing_handler(self): lg = logging.getLogger('madgraph') handler = logging.NullHandler() lg.addHandler(handler) - from madgraph.iolibs.template_files.mg7 import postprocess + from madgraph.mg7 import postprocess saved = postprocess._TOOL_LOGGING_READY postprocess._TOOL_LOGGING_READY = False try: @@ -1029,7 +1029,7 @@ def tearDown(self): self.logger.setLevel(self.saved) def test_info_is_dropped_during_setup(self): - from madgraph.iolibs.template_files.mg7.run_interface import _quiet_setup + from madgraph.mg7.run_interface import _quiet_setup with _quiet_setup(): self.logger.info('W E L C O M E to') @@ -1040,14 +1040,14 @@ def test_info_is_dropped_during_setup(self): def test_warnings_still_get_through(self): """Quiet is not silent: anything that needs saying still says it.""" - from madgraph.iolibs.template_files.mg7.run_interface import _quiet_setup + from madgraph.mg7.run_interface import _quiet_setup with _quiet_setup(): self.logger.warning('something worth knowing') self.assertIn('something worth knowing', self.captured.getvalue()) def test_the_level_is_restored_afterwards(self): - from madgraph.iolibs.template_files.mg7.run_interface import _quiet_setup + from madgraph.mg7.run_interface import _quiet_setup with _quiet_setup(): pass @@ -1055,7 +1055,7 @@ def test_the_level_is_restored_afterwards(self): self.assertIn('back to normal', self.captured.getvalue()) def test_it_restores_the_level_after_a_failure(self): - from madgraph.iolibs.template_files.mg7.run_interface import _quiet_setup + from madgraph.mg7.run_interface import _quiet_setup try: with _quiet_setup(): @@ -1078,7 +1078,7 @@ class TestPostProcessingIsSkippedWhenThereIsNothingToDo(unittest.TestCase): def tools_for(self, switch): """The tool list run_selected_tools builds, without running it.""" - from madgraph.iolibs.template_files.mg7 import launch + from madgraph.mg7 import launch return launch.selected_tools(switch) def test_a_not_available_switch_selects_no_tool(self): @@ -1112,7 +1112,7 @@ def test_run_selected_tools_returns_before_building_anything(self): """With no tool selected it must not construct MG7RunCmd -- that is what printed the MADEVENT banner and the configuration lines.""" - from madgraph.iolibs.template_files.mg7 import launch + from madgraph.mg7 import launch class _Process(object): run_path = '/nonexistent/run_01' @@ -1148,7 +1148,7 @@ def tearDown(self): def test_store_result_bookkeeping_writes_nothing(self): import madgraph.various.banner as banner_mod - from madgraph.iolibs.template_files.mg7.run_interface import MG7RunCmd + from madgraph.mg7.run_interface import MG7RunCmd # load_results_db only needs me_dir, process and the model name # (__del__ reads the other two); building a real MG7RunCmd needs a @@ -1190,7 +1190,7 @@ class MG7MeFrameTest(unittest.TestCase): def resolve(self, me_frame, is_decay, particle_count=4, outgoing=None): """``outgoing``, if given, lists the final states of several subprocesses of one output (g g > each of them)""" - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch process = object.__new__(mg7_launch.MadgraphProcess) process.run_card = {'run': {'me_frame': me_frame}} process.is_decay = is_decay @@ -1277,7 +1277,7 @@ class TestForceLHEOutput(unittest.TestCase): the LHE format only when something that reads LHE will run.""" def setUp(self): - from madgraph.iolibs.template_files.mg7 import launch + from madgraph.mg7 import launch from madgraph.various.banner import RunCardMG7 self.launch = launch self.tmpdir = tempfile.TemporaryDirectory() @@ -1365,7 +1365,7 @@ class TestNpyToLHE(unittest.TestCase): def setUp(self): import numpy as np - from madgraph.iolibs.template_files.mg7 import npy_to_lhe + from madgraph.mg7 import npy_to_lhe self.np = np self.npy_to_lhe = npy_to_lhe self.ms = npy_to_lhe._madspace() @@ -1463,14 +1463,14 @@ def test_defaults(self): def test_mean_ignores_under_and_overflow(self): """Both histogram kinds carry the under/overflow bins first and last; they have no position and must not shift the in-range bins.""" - from madgraph.iolibs.template_files.mg7 import launch + from madgraph.mg7 import launch mean = launch.MadgraphProcess._histogram_mean # bins [0,10) and [10,20) with weights 1 and 3, plus 100 in overflow self.assertAlmostEqual(mean(0., 20., [5., 1., 3., 100.]), 12.5) self.assertIsNone(mean(0., 20., [5., 0., 0., 100.])) def test_postprocessing_off_builds_no_event_histograms(self): - from madgraph.iolibs.template_files.mg7 import launch + from madgraph.mg7 import launch from madgraph.various.banner import RunCardMG7 process = launch.MadgraphProcess.__new__(launch.MadgraphProcess) process.run_card = RunCardMG7() @@ -1490,7 +1490,7 @@ class MG7ObservablePidsTest(unittest.TestCase): flavours, consistent over every flavour a subprocess groups together.""" def setUp(self): - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch self.launch = mg7_launch @staticmethod @@ -1536,7 +1536,7 @@ class MG7DecayProductsTest(unittest.TestCase): on-shell (decay-chain) propagator decays into, as MadEvent's check_decay.""" def setUp(self): - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch self.launch = mg7_launch @staticmethod diff --git a/tests/unit_tests/interface/test_mg7_systematics_summary.py b/tests/unit_tests/interface/test_mg7_systematics_summary.py index de9d6b39c9..7885c96917 100644 --- a/tests/unit_tests/interface/test_mg7_systematics_summary.py +++ b/tests/unit_tests/interface/test_mg7_systematics_summary.py @@ -26,7 +26,7 @@ import inspect import unittest -import madgraph.iolibs.template_files.mg7.systematics_summary as systematics_summary +import madgraph.mg7.systematics_summary as systematics_summary class _Summary(object): @@ -142,7 +142,7 @@ class TestOneImplementation(unittest.TestCase): """Neither caller may grow its own copy back.""" def sources(self): - import madgraph.iolibs.template_files.mg7.launch as launch + import madgraph.mg7.launch as launch return { 'the Systematics box': diff --git a/tests/unit_tests/interface/test_tutorials.py b/tests/unit_tests/interface/test_tutorials.py index 635601d439..8403ed3802 100644 --- a/tests/unit_tests/interface/test_tutorials.py +++ b/tests/unit_tests/interface/test_tutorials.py @@ -1924,7 +1924,7 @@ def test_the_percentages_come_from_the_shared_helper(self): """Not a third implementation: the run's own box and the scan summary read the same one.""" - from madgraph.iolibs.template_files.mg7 import systematics_summary + from madgraph.mg7 import systematics_summary up, down = systematics_summary.scale_percentages(self.INFO['systematics']) rows = self.module()._systematics_rows(self.INFO) diff --git a/tests/unit_tests/iolibs/test_mg7_discrete_flavor.py b/tests/unit_tests/iolibs/test_mg7_discrete_flavor.py index ec03a08472..4d2316c15a 100644 --- a/tests/unit_tests/iolibs/test_mg7_discrete_flavor.py +++ b/tests/unit_tests/iolibs/test_mg7_discrete_flavor.py @@ -22,14 +22,14 @@ import types import unittest -from madgraph.iolibs.template_files.mg7 import bootstrap as mg7_bootstrap +from madgraph.mg7 import bootstrap as mg7_bootstrap # Importing launch without a compiled madspace starts the madspace build, and # test_manager imports every unit test module while collecting: check first. launch = None if mg7_bootstrap.madspace_is_installed(): try: - import madgraph.iolibs.template_files.mg7.launch as launch + import madgraph.mg7.launch as launch except Exception as error: launch = None import_error = error diff --git a/tests/unit_tests/various/test_mg7_hwu_output.py b/tests/unit_tests/various/test_mg7_hwu_output.py index 8f8462a544..9e1c8c2162 100644 --- a/tests/unit_tests/various/test_mg7_hwu_output.py +++ b/tests/unit_tests/various/test_mg7_hwu_output.py @@ -29,7 +29,7 @@ import unittest import madgraph.various.histograms as histograms -from madgraph.iolibs.template_files.mg7 import hwu_output +from madgraph.mg7 import hwu_output def make_histogram(name='jet-pt', bin_count=2, weights=True): diff --git a/tests/unit_tests/various/test_mg7_plots.py b/tests/unit_tests/various/test_mg7_plots.py index 533dcf345d..13f157bc00 100644 --- a/tests/unit_tests/various/test_mg7_plots.py +++ b/tests/unit_tests/various/test_mg7_plots.py @@ -29,7 +29,7 @@ import unittest from madgraph import MG5DIR -from madgraph.iolibs.template_files.mg7 import plots +from madgraph.mg7 import plots def make_histogram(name='jet-pt', values=None, bin_count=4): @@ -156,7 +156,7 @@ def test_matplotlib_does_not_narrate_the_drawing(self): 'import json, logging, os, shutil, sys\n' 'logging.basicConfig(level=logging.DEBUG)\n' 'sys.path.insert(0, %r)\n' % MG5DIR + - 'from madgraph.iolibs.template_files.mg7 import plots\n' + 'from madgraph.mg7 import plots\n' 'out = %r\n' % os.path.join(self.path, 'sub') + 'plots.render([%r], out)\n' % make_histogram() ) From 36325bfd0c0e6c3fa5377052609e1aa23f916b92 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 13:54:26 +0100 Subject: [PATCH 56/88] madspace: freeze the maximum weight independently of the thread timing The freeze compared the written events, which are committed in a thread dependent order relative to the generate jobs. Count the expected unweighted events in update_max_weight instead, in the commit order of the jobs. Fixes the flaky test_reproducible in test_weighting_steps.py. --- .../madspace/driver/channel_generator.hpp | 8 ++++++-- madspace/src/driver/channel_generator.cpp | 19 +++++++++++-------- madspace/src/driver/event_generator.cpp | 4 ++-- 3 files changed, 19 insertions(+), 12 deletions(-) diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 8b1253c005..2f661edb14 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -305,8 +305,9 @@ class ChannelEventGenerator { std::size_t next_vegas_batch_size(); /// Discard the accumulated event and weight files. void clear_events(); - /// Update @ref max_weight from a batch of `weights`. - void update_max_weight(Tensor weights); + /// Update @ref max_weight from a batch of `weights`, which are unweighted + /// later if `unweight` is true. + void update_max_weight(Tensor weights, bool unweight); /// Append `unweighted_events` to @ref event_file, recording /// `job_max_weight`. void write_events(const TensorVec& unweighted_events, double job_max_weight); @@ -430,6 +431,9 @@ class ChannelEventGenerator { std::vector _step_states; std::vector _loaded_step_max_weights; double _first_step_count = 0.; + // expected number of unweighted events, counted in the commit order of the + // jobs to freeze the maximum weight independently of the thread timing + double _freeze_count = 0.; StepIntegral _final_integral; std::size_t _event_epoch = 0; diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index b34db776b1..c27bba0fae 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -742,6 +742,7 @@ void ChannelEventGenerator::clear_events() { _abs_cross_section.reset(); _large_weights.clear(); _first_step_count = 0.; + _freeze_count = 0.; _final_integral = {}; ++_event_epoch; for (auto& state : _step_states) { @@ -757,11 +758,9 @@ void ChannelEventGenerator::clear_events() { } } -void ChannelEventGenerator::update_max_weight(Tensor weights) { - // with steps, the first step freezes on its own accepted events - double unweighted_count = - _step_settings.empty() ? _status.count_unweighted : _first_step_count; - if (_max_weight_fixed || unweighted_count > _config.freeze_max_weight_after) { +void ChannelEventGenerator::update_max_weight(Tensor weights, bool unweight) { + // not the written events, they are committed in a thread dependent order + if (_max_weight_fixed || _freeze_count > _config.freeze_max_weight_after) { return; } @@ -781,6 +780,11 @@ void ChannelEventGenerator::update_max_weight(Tensor weights) { } std::sort(_large_weights.begin(), _large_weights.end(), std::greater{}); apply_truncation_budget(); + if (unweight && _max_weight > 0) { + for (std::size_t i = 0; i < w_view.size(); ++i) { + _freeze_count += std::min(1., std::abs(w_view[i]) / _max_weight); + } + } } void ChannelEventGenerator::set_target_count(std::size_t target_count) { @@ -791,10 +795,8 @@ void ChannelEventGenerator::set_target_count(std::size_t target_count) { // (e.g. generate()'s round-one seed) would otherwise never be corrected: a // channel that has met its smaller target is not revisited. Tighten the cap // now; unweight_all() re-unweights the whole file against it at the end. - double unweighted_count = - _step_settings.empty() ? _status.count_unweighted : _first_step_count; if (target_count < old_target && _max_weight != 0 && !_max_weight_fixed && - unweighted_count <= _config.freeze_max_weight_after) { + _freeze_count <= _config.freeze_max_weight_after) { apply_truncation_budget(); } } @@ -816,6 +818,7 @@ void ChannelEventGenerator::apply_truncation_budget() { if (_max_weight < w) { _status.count_unweighted *= _max_weight / w; _first_step_count *= _max_weight / w; + _freeze_count *= _max_weight / w; for (auto& state : _step_states) { state.count_accepted *= _max_weight / w; } diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index 95b6fa3a86..e81c56e090 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -75,7 +75,7 @@ void EventGenerator::commit_generate_job(GeneratorBatchJob& job) { } channel->integrate(job); update_integral_status(); - channel->update_max_weight(job.weights); + channel->update_max_weight(job.weights, job.unweight); --channel_job_count; --_context_job_counts.at(job.context_index); if (job.unweight) { @@ -338,7 +338,7 @@ void EventGenerator::survey(std::size_t survey_pass) { } channel->integrate(gen_job); update_integral(); - channel->update_max_weight(gen_job.weights); + channel->update_max_weight(gen_job.weights, gen_job.unweight); --channel_job_count; --_context_job_counts.at(gen_job.context_index); // Global order is ascending job id and so is the per-channel one, From 636ca9c5bc0f9a5fa5358ce2658c07ac5c2e1d31 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 19:53:53 +0100 Subject: [PATCH 57/88] ci: pyamplicol workflow imports bootstrap from madgraph.mg7 and triggers on umami.h --- .github/workflows/pyamplicol.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/pyamplicol.yml b/.github/workflows/pyamplicol.yml index b847a4b4e9..12e82d7d1b 100644 --- a/.github/workflows/pyamplicol.yml +++ b/.github/workflows/pyamplicol.yml @@ -73,7 +73,7 @@ jobs: if [[ -z "$base" || "$base" =~ ^0+$ ]] || ! git cat-file -e "$base^{commit}" 2>/dev/null; then changed=true elif git diff --name-only "$base" HEAD \ - | grep -qE '^(pyamplicol/|\.github/workflows/pyamplicol|\.github/actions/(pyamplicol|build_pyamplicol|install_pyamplicol)|VERSION$)'; then + | grep -qE '^(pyamplicol/|\.github/workflows/pyamplicol|\.github/actions/(pyamplicol|build_pyamplicol|install_pyamplicol)|madspace/include/madspace/umami\.h$|madgraph/mg7/bootstrap\.py$|VERSION$)'; then changed=true else changed=false @@ -173,7 +173,7 @@ jobs: python - <<'EOF' import sys sys.path.insert(0, ".") - from madgraph.iolibs.template_files.mg7 import bootstrap + from madgraph.mg7 import bootstrap bootstrap.ensure_pyamplicol(interactive=False) bootstrap.check_pyamplicol_license() assert bootstrap.check_pyamplicol_source_hash() From cff09fcbba0edf75ef2e626389eafb448c3a1123 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 19:53:58 +0100 Subject: [PATCH 58/88] ci: docs examples workflow triggers on madgraph/mg7 --- .github/workflows/docs_examples.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/docs_examples.yml b/.github/workflows/docs_examples.yml index 5138407ccf..5afbaa3343 100644 --- a/.github/workflows/docs_examples.yml +++ b/.github/workflows/docs_examples.yml @@ -13,6 +13,7 @@ on: - 'docs/source/madspace/examples/**' - 'docs/source/event_stream.rst' - 'madgraph/iolibs/template_files/mg7/**' + - 'madgraph/mg7/**' - 'madspace/**' - 'tests/acceptance_tests/test_madspace_examples.py' - '.github/workflows/docs_examples.yml' @@ -23,6 +24,7 @@ on: - 'docs/source/madspace/examples/**' - 'docs/source/event_stream.rst' - 'madgraph/iolibs/template_files/mg7/**' + - 'madgraph/mg7/**' - 'madspace/**' - 'tests/acceptance_tests/test_madspace_examples.py' - '.github/workflows/docs_examples.yml' From cc9d3c273c895ff5486848715f23df648d18fb65 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 19:54:28 +0100 Subject: [PATCH 59/88] ci: release smoke test fails when pyamplicol's license check fails --- .github/workflows/release.yml | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 25160316e2..ad98defd71 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -227,12 +227,16 @@ jobs: done test -d madspace/install/madspace - # Also checks that Symbolica accepts pyamplicol's library key. + # Also checks that Symbolica accepts pyamplicol's library key; only the + # wheel download is retried, so a rejected key fails the step. for attempt in $(seq 10); do - python3 pyamplicol/install.py --bin -y && break + if python3 pyamplicol/install.py --bin -y; then status=0; break; else status=$?; fi + # The wheel installed, so the license check failed: retrying cannot help. + test -d pyamplicol/install/pyamplicol && break echo "attempt $attempt: pyamplicol==${{ needs.check_version.outputs.version }} not installable yet, retrying..." sleep 15 done + test "$status" -eq 0 test -d pyamplicol/install/pyamplicol - name: Sign the tarball with Sigstore From 03b45479a00e2576ca9bec653f977842b1fb95ec Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 19:54:49 +0100 Subject: [PATCH 60/88] mg7: [amplicol] after another bracket is refused like before it --- madgraph/interface/madgraph_interface.py | 5 +++-- tests/unit_tests/interface/test_amplicol_syntax.py | 4 ++++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/madgraph/interface/madgraph_interface.py b/madgraph/interface/madgraph_interface.py index 8e2ed6b426..d22c449c88 100755 --- a/madgraph/interface/madgraph_interface.py +++ b/madgraph/interface/madgraph_interface.py @@ -190,8 +190,9 @@ def split_backend_line(line): """(line without the bracket, backend) for a process with a matrix-element backend bracket like '[amplicol]', (line, '') for any other line.""" - match = _BACKEND_BRACKET.search(line) - if not match or match.group('backend').lower() not in ME_BACKENDS: + match = next((m for m in _BACKEND_BRACKET.finditer(line) + if m.group('backend').lower() in ME_BACKENDS), None) + if not match: return line, '' rest = line[:match.start()] + ' ' + line[match.end():] if '[' in rest or ']' in rest: diff --git a/tests/unit_tests/interface/test_amplicol_syntax.py b/tests/unit_tests/interface/test_amplicol_syntax.py index 9c5196b618..84804ddf5d 100644 --- a/tests/unit_tests/interface/test_amplicol_syntax.py +++ b/tests/unit_tests/interface/test_amplicol_syntax.py @@ -48,6 +48,10 @@ def test_split(self): def test_other_brackets_are_refused(self): self.assertRaises(InvalidCmd, split, 'p p > j j [amplicol] [QCD]') + self.assertRaises(InvalidCmd, split, 'p p > j j [QCD] [amplicol]') + # the pre-check must not route it to aMC@NLO with 'amplicol' as order + process_type = MGCmd.MasterCmd.extract_process_type + self.assertRaises(InvalidCmd, process_type, 'p p > t t~ [QCD] [amplicol]') def test_process_type(self): process_type = MGCmd.MasterCmd.extract_process_type From ba95ed9b796ba171a81f2744175aa7c2e8e49063 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 19:56:23 +0100 Subject: [PATCH 61/88] mg7: always reap the pyamplicol worker and stop waiting for a busy Symbolica after an hour --- madgraph/iolibs/export_amplicol.py | 48 ++++++++--- .../various/test_export_amplicol.py | 83 +++++++++++++++++++ 2 files changed, 120 insertions(+), 11 deletions(-) create mode 100644 tests/unit_tests/various/test_export_amplicol.py diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index dc422c8242..9d914c429a 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -61,6 +61,9 @@ WORKER = pjoin(os.path.dirname(os.path.abspath(__file__)), 'amplicol_worker.py') # printed by Symbolica when another process holds its per-user thread SYMBOLICA_BUSY = 'Cannot start another restricted Symbolica thread' +# seconds between the attempts while Symbolica is busy, and before giving up +SYMBOLICA_RETRY = 10 +SYMBOLICA_MAX_WAIT = 3600 # rusticol-config of the pyamplicol in pyamplicol/install, for the makefiles @@ -96,18 +99,11 @@ def run_worker(job): """Run amplicol_worker on a job, waiting while Symbolica is busy elsewhere.""" env = dict(os.environ, SYMBOLICA_HIDE_BANNER='1', PYTHONPATH=os.pathsep.join(p for p in sys.path if p)) + start = time.monotonic() waiting = False while True: - worker = subprocess.Popen([sys.executable, WORKER, job], env=env, - stdout=subprocess.PIPE, stderr=subprocess.STDOUT, - text=True) - busy = False - for line in worker.stdout: - if SYMBOLICA_BUSY in line: - busy = True - elif not line.startswith(('\u250c', '\u2502', '\u2514', '\u251c')): - sys.stdout.write(line) - if worker.wait() == 0: + returncode, busy = run_worker_once(job, env) + if returncode == 0: return if not busy: raise madgraph.InvalidCmd('pyamplicol failed to generate the [amplicol] ' @@ -116,10 +112,40 @@ def run_worker(job): raise madgraph.InvalidCmd( 'Symbolica is imported in this MadGraph process, so the pyamplicol ' 'worker cannot start: do not import pyamplicol or symbolica here') + if time.monotonic() - start > SYMBOLICA_MAX_WAIT: + raise madgraph.InvalidCmd( + 'Symbolica has been used by another process for %d minutes: only one ' + 'pyamplicol process may run per user, stop the other one and output ' + 'again' % (SYMBOLICA_MAX_WAIT // 60)) if not waiting: logger.info('Symbolica is used by another process: waiting for it') waiting = True - time.sleep(10) + time.sleep(SYMBOLICA_RETRY) + + +def run_worker_once(job, env): + """The exit code of one worker run, and whether Symbolica was busy. The + worker never outlives this call, so it cannot keep holding Symbolica's lock.""" + worker = subprocess.Popen([sys.executable, WORKER, job], env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True) + busy = False + try: + for line in worker.stdout: + if SYMBOLICA_BUSY in line: + busy = True + elif not line.startswith(('\u250c', '\u2502', '\u2514', '\u251c')): + sys.stdout.write(line) + return worker.wait(), busy + finally: + if worker.poll() is None: + worker.terminate() + try: + worker.wait(timeout=10) + except subprocess.TimeoutExpired: + worker.kill() + worker.wait() + worker.stdout.close() def load_metadata(artifact): diff --git a/tests/unit_tests/various/test_export_amplicol.py b/tests/unit_tests/various/test_export_amplicol.py new file mode 100644 index 0000000000..79c1ace905 --- /dev/null +++ b/tests/unit_tests/various/test_export_amplicol.py @@ -0,0 +1,83 @@ +################################################################################ +# +# Copyright (c) 2026 The MadGraph5_aMC@NLO Development team and Contributors +# +# This file is a part of the MadGraph5_aMC@NLO project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph5_aMC@NLO license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""Test the pyamplicol worker handling of export_amplicol, with fake workers +instead of pyamplicol.""" + +import io +import os +import shutil +import subprocess +import sys +import tempfile +import unittest +from unittest import mock + +import madgraph +import madgraph.iolibs.export_amplicol as export_amplicol + + +class TestRunWorker(unittest.TestCase): + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmp) + + def fake_worker(self, code): + path = os.path.join(self.tmp, 'worker.py') + with open(path, 'w') as stream: + stream.write('import sys, time\n' + code) + patcher = mock.patch.object(export_amplicol, 'WORKER', path) + patcher.start() + self.addCleanup(patcher.stop) + + def test_success_prints_the_output(self): + self.fake_worker('print("generated")\n') + with mock.patch('sys.stdout', new=io.StringIO()) as out: + export_amplicol.run_worker('job.json') + self.assertIn('generated', out.getvalue()) + + def test_failure_raises(self): + self.fake_worker('sys.exit(1)\n') + with self.assertRaises(madgraph.InvalidCmd): + export_amplicol.run_worker('job.json') + + def test_busy_symbolica_is_retried_then_given_up(self): + self.fake_worker('print(%r); sys.exit(1)\n' % export_amplicol.SYMBOLICA_BUSY) + with mock.patch.object(export_amplicol, 'SYMBOLICA_RETRY', 0), \ + mock.patch.object(export_amplicol, 'SYMBOLICA_MAX_WAIT', 0.3), \ + mock.patch.object(export_amplicol, 'run_worker_once', + wraps=export_amplicol.run_worker_once) as once: + with self.assertRaisesRegex(madgraph.InvalidCmd, 'another process'): + export_amplicol.run_worker('job.json') + self.assertGreater(once.call_count, 1) + + def test_interrupt_kills_the_worker(self): + self.fake_worker('print("started", flush=True); time.sleep(60)\n') + workers = [] + popen = subprocess.Popen + + def record(*args, **kwargs): + workers.append(popen(*args, **kwargs)) + return workers[-1] + + with mock.patch('subprocess.Popen', side_effect=record), \ + mock.patch('sys.stdout.write', side_effect=KeyboardInterrupt): + with self.assertRaises(KeyboardInterrupt): + export_amplicol.run_worker('job.json') + self.assertIsNotNone(workers[0].returncode) + + +if __name__ == '__main__': + unittest.main() From 0bd0f6b72164e1cffd69c430d6722826a14f5770 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 19:56:53 +0100 Subject: [PATCH 62/88] mg7: launch test writes the run_card from the mg7 template --- tests/unit_tests/interface/test_mg7_launch.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit_tests/interface/test_mg7_launch.py b/tests/unit_tests/interface/test_mg7_launch.py index 4501bed310..67dd9951ed 100644 --- a/tests/unit_tests/interface/test_mg7_launch.py +++ b/tests/unit_tests/interface/test_mg7_launch.py @@ -1277,6 +1277,7 @@ class TestForceLHEOutput(unittest.TestCase): the LHE format only when something that reads LHE will run.""" def setUp(self): + import madgraph from madgraph.mg7 import launch from madgraph.various.banner import RunCardMG7 self.launch = launch @@ -1284,8 +1285,9 @@ def setUp(self): self.cwd = os.getcwd() os.chdir(self.tmpdir.name) os.mkdir('Cards') - template = os.path.join(os.path.dirname(launch.__file__), - 'run_card.toml') + template = os.path.join(madgraph.MG5DIR, 'madgraph', 'iolibs', + 'template_files', 'mg7', 'run_card.toml') + self.assertTrue(os.path.isfile(template), template) RunCardMG7().write(self.card, template=template) def tearDown(self): From cb4fa0fb6d2f30f5c703682b17cd73f02f714e10 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 19:57:13 +0100 Subject: [PATCH 63/88] mg7: refuse stale pyamplicol installs at output and before compiling [amplicol] libraries, accept system installs --- madgraph/iolibs/export_amplicol.py | 43 ++++++--- madgraph/mg7/bootstrap.py | 96 ++++++++++++++----- .../various/test_export_amplicol.py | 21 ++++ .../various/test_pyamplicol_install.py | 64 +++++++++++++ 4 files changed, 186 insertions(+), 38 deletions(-) diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index 9d914c429a..18c7185a74 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -43,6 +43,8 @@ pjoin = os.path.join MODEL_FILE = 'amplicol_model.pyamplicol-model' +# the pyamplicol that generated the output, checked before compiling +PYAMPLICOL_RECORD = 'amplicol_pyamplicol.json' MAKEFILE = """\ # Builds the AmpliCol matrix elements of this process into ../../lib. @@ -50,6 +52,7 @@ BUILD := $(CURDIR)/build all: +\t$(CURDIR)/../check-pyamplicol \t$(MAKE) -C lc/API/umami shared BUILD_DIR=$(BUILD)/lc RUSTICOL_CONFIG=$(RUSTICOL_CONFIG) \t$(MAKE) -C reweight shared BUILD_DIR=$(BUILD)/reweight RUSTICOL_CONFIG=$(RUSTICOL_CONFIG) %(copies)s @@ -72,27 +75,43 @@ PYTHONPATH=%(install)s${PYTHONPATH:+:$PYTHONPATH} exec %(python)s -m rusticol_config "$@" """ +# fails unless that pyamplicol is the one recorded in PYAMPLICOL_RECORD +CHECK_PYAMPLICOL = """\ +#!/bin/sh +PYTHONPATH=%(install)s${PYTHONPATH:+:$PYTHONPATH} exec %(python)s -c \\ + 'import sys; sys.path.insert(0, sys.argv[1]); from madgraph.mg7.bootstrap import check_pyamplicol_build; check_pyamplicol_build(sys.argv[2])' \\ + %(mg_root)s "$(dirname "$0")/%(record)s" +""" + def check_pyamplicol(interactive=None, jobs=None): """Install pyamplicol in pyamplicol/install if needed and put it first on - sys.path, which the worker inherits. Never imported here, see amplicol_worker.""" + sys.path, which the worker inherits. Never imported here, see amplicol_worker. + Returns its version and source hash.""" from madgraph.mg7 import bootstrap try: bootstrap.ensure_pyamplicol(interactive=interactive, jobs=jobs) + return bootstrap.check_pyamplicol_source_hash() except RuntimeError as error: raise madgraph.InvalidCmd('The [amplicol] processes need pyamplicol: %s' % error) - bootstrap.check_pyamplicol_source_hash() -def write_rusticol_config(subproc_dir): - """SubProcesses/rusticol-config, the RUSTICOL_CONFIG of the makefiles.""" +def write_rusticol_config(subproc_dir, pyamplicol): + """SubProcesses/rusticol-config, the RUSTICOL_CONFIG of the makefiles, and + check-pyamplicol with the pyamplicol it has to be.""" from madgraph.mg7 import bootstrap - path = pjoin(subproc_dir, 'rusticol-config') - with open(path, 'w') as stream: - stream.write(RUSTICOL_CONFIG % { - 'install': shlex.quote(str(bootstrap.PYAMPLICOL_INSTALL_DIR)), - 'python': shlex.quote(sys.executable)}) - os.chmod(path, 0o755) + values = {'install': shlex.quote(str(bootstrap.PYAMPLICOL_INSTALL_DIR)), + 'python': shlex.quote(sys.executable), + 'mg_root': shlex.quote(str(bootstrap.MG_ROOT)), + 'record': PYAMPLICOL_RECORD} + for name, template in (('rusticol-config', RUSTICOL_CONFIG), + ('check-pyamplicol', CHECK_PYAMPLICOL)): + path = pjoin(subproc_dir, name) + with open(path, 'w') as stream: + stream.write(template % values) + os.chmod(path, 0o755) + with open(pjoin(subproc_dir, PYAMPLICOL_RECORD), 'w') as stream: + json.dump(pyamplicol, stream, indent=2) def run_worker(job): @@ -389,9 +408,9 @@ def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None, """Write SubProcesses/A_ for every '[amplicol]' process definition and return the subprocesses.json entries. ``interactive`` and ``jobs`` are those of the pyamplicol installation, see bootstrap.ensure_pyamplicol.""" - check_pyamplicol(interactive, jobs) + pyamplicol = check_pyamplicol(interactive, jobs) subproc_dir = pjoin(dir_path, 'SubProcesses') - write_rusticol_config(subproc_dir) + write_rusticol_config(subproc_dir, pyamplicol) processes = [] for number, procdef in enumerate(proc_defs, start=1): name = 'A%d_%s' % (number, process_name(procdef, model, multiparticles or {})) diff --git a/madgraph/mg7/bootstrap.py b/madgraph/mg7/bootstrap.py index 74fae15d21..15e62463df 100644 --- a/madgraph/mg7/bootstrap.py +++ b/madgraph/mg7/bootstrap.py @@ -19,6 +19,7 @@ safe to do lazily. """ +import importlib.metadata import importlib.util import json import os @@ -44,9 +45,39 @@ def madspace_is_installed() -> bool: return (INSTALL_DIR / "madspace").is_dir() +def _pyamplicol_package(): + """The directory of the pyamplicol package that ``import pyamplicol`` would + load, or None. Found, not imported (Symbolica would take its per-user lock).""" + spec = importlib.util.find_spec("pyamplicol") + if spec is None or not spec.origin: + return None # e.g. the source directory as a namespace package + package = Path(spec.origin).parent + if (package.resolve().is_relative_to(PYAMPLICOL_DIR.resolve()) + and not package.resolve().is_relative_to(PYAMPLICOL_INSTALL_DIR.resolve())): + return None + return package + + def pyamplicol_is_installed() -> bool: - """True when a compiled pyamplicol is already available to import.""" - return (PYAMPLICOL_INSTALL_DIR / "pyamplicol").is_dir() + """True when a compiled pyamplicol is already available to import, in + pyamplicol/install or system-wide (``install pyamplicol --system``).""" + return (PYAMPLICOL_INSTALL_DIR / "pyamplicol").is_dir() or _pyamplicol_package() is not None + + +def installed_pyamplicol() -> dict: + """The version and source hash of the pyamplicol that would be imported + (None where unknown), read without importing it.""" + package = _pyamplicol_package() + if package is None: + return {"version": None, "source_hash": None} + source_hash = None + hash_file = package / "_source_hash.py" + if hash_file.is_file(): + match = re.search(r'SOURCE_HASH\s*=\s*"([0-9a-f]*)"', hash_file.read_text()) + source_hash = match.group(1) if match else None + dist = next(importlib.metadata.distributions( + name="pyamplicol", path=[str(package.parent)]), None) + return {"version": dist.version if dist else None, "source_hash": source_hash} def _run_installer(name, source_dir, interactive, jobs) -> None: @@ -114,34 +145,47 @@ def ensure_pyamplicol(interactive=None, jobs=None) -> None: sys.path.insert(0, str(PYAMPLICOL_INSTALL_DIR)) -def check_pyamplicol_source_hash() -> bool: - """Warn if the installed pyamplicol was built from other sources. +def check_pyamplicol_source_hash() -> dict: + """Fail unless the installed pyamplicol was built from the sources in + pyamplicol/, and return :func:`installed_pyamplicol`. - Returns whether the hashes match. Must run after :func:`ensure_pyamplicol`. + Must run after :func:`ensure_pyamplicol`. """ - # read, not imported: importing pyamplicol would take Symbolica's per-user - # lock in this process and block the generation worker - installed = None - spec = importlib.util.find_spec("pyamplicol") - if spec is not None and spec.origin: - hash_file = Path(spec.origin).parent / "_source_hash.py" - if hash_file.is_file(): - match = re.search(r'SOURCE_HASH\s*=\s*"([0-9a-f]*)"', hash_file.read_text()) - installed = match.group(1) if match else None - source = subprocess.run( + installed = installed_pyamplicol() + result = subprocess.run( [sys.executable, str(PYAMPLICOL_DIR / "source_hash.py")], - capture_output=True, text=True, check=True, - ).stdout.strip() - if installed == source: - return True - print() - print( - "\033[1m\033[31mWARNING\033[39m: pyamplicol source and installed package " - "are not compatible (source hash mismatch) — consider recompiling " - "pyamplicol (e.g. `install pyamplicol -y`)\033[0m" + capture_output=True, text=True, ) - print() - return False + if result.returncode != 0: + raise RuntimeError( + "cannot compute the hash of the pyamplicol sources:\n" + + result.stdout + result.stderr + ) + if installed["source_hash"] != result.stdout.strip(): + raise RuntimeError( + "the installed pyamplicol was not built from the sources in " + f"{PYAMPLICOL_DIR} (source hash mismatch): run `install pyamplicol` " + "to rebuild it" + ) + return installed + + +def check_pyamplicol_build(record_path) -> None: + """Exit with an error unless the pyamplicol that would be imported is the + one recorded in record_path by the output. Run by SubProcesses/check-pyamplicol + before the [amplicol] libraries are compiled.""" + with open(record_path) as stream: + recorded = json.load(stream) + installed = installed_pyamplicol() + if installed != {key: recorded.get(key) for key in installed}: + sys.exit( + f"The [amplicol] matrix elements of this output were generated with " + f"pyamplicol {recorded.get('version')} (source hash " + f"{recorded.get('source_hash')}), but the installed one is " + f"{installed['version']} (source hash {installed['source_hash']}): " + "output the process again, or reinstall the pyamplicol it was " + "generated with (`install pyamplicol`)." + ) def check_pyamplicol_license(pythonpath=None) -> None: diff --git a/tests/unit_tests/various/test_export_amplicol.py b/tests/unit_tests/various/test_export_amplicol.py index 79c1ace905..50710f9d4c 100644 --- a/tests/unit_tests/various/test_export_amplicol.py +++ b/tests/unit_tests/various/test_export_amplicol.py @@ -79,5 +79,26 @@ def record(*args, **kwargs): self.assertIsNotNone(workers[0].returncode) +class TestCheckPyamplicol(unittest.TestCase): + + def test_stale_install_is_an_error(self): + from madgraph.mg7 import bootstrap + with mock.patch.object(bootstrap, 'ensure_pyamplicol'), \ + mock.patch.object(bootstrap, 'check_pyamplicol_source_hash', + side_effect=RuntimeError('mismatch')): + with self.assertRaisesRegex(madgraph.InvalidCmd, 'mismatch'): + export_amplicol.check_pyamplicol() + + def test_check_script_rejects_another_pyamplicol(self): + tmp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, tmp) + export_amplicol.write_rusticol_config( + tmp, {'version': '0.0.0', 'source_hash': 'stale'}) + result = subprocess.run([os.path.join(tmp, 'check-pyamplicol')], + capture_output=True, text=True) + self.assertNotEqual(result.returncode, 0) + self.assertIn('generated with pyamplicol 0.0.0', result.stderr) + + if __name__ == '__main__': unittest.main() diff --git a/tests/unit_tests/various/test_pyamplicol_install.py b/tests/unit_tests/various/test_pyamplicol_install.py index a23d0a0db5..e25980077a 100644 --- a/tests/unit_tests/various/test_pyamplicol_install.py +++ b/tests/unit_tests/various/test_pyamplicol_install.py @@ -194,5 +194,69 @@ def test_generated_files_are_ignored(self): self.assertEqual(source_hash.compute_source_hash(root), before) +class TestInstalledPyamplicol(unittest.TestCase): + """bootstrap finds the installed pyamplicol without importing it.""" + + def setUp(self): + from madgraph.mg7 import bootstrap + self.bootstrap = bootstrap + self.tmp = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, self.tmp) + patcher = mock.patch.object(bootstrap, 'PYAMPLICOL_INSTALL_DIR', self.tmp / 'none') + patcher.start() + self.addCleanup(patcher.stop) + + def system_package(self, version='7.0.0', source_hash='abc123'): + site = self.tmp / 'site' + (site / 'pyamplicol').mkdir(parents=True) + (site / 'pyamplicol' / '__init__.py').write_text('raise ImportError\n') + (site / 'pyamplicol' / '_source_hash.py').write_text( + f'SOURCE_HASH = "{source_hash}"\n') + info = site / f'pyamplicol-{version}.dist-info' + info.mkdir() + (info / 'METADATA').write_text(f'Name: pyamplicol\nVersion: {version}\n') + patcher = mock.patch.object(sys, 'path', [str(site)] + sys.path) + patcher.start() + self.addCleanup(patcher.stop) + + def test_system_install_counts_as_installed(self): + self.system_package() + self.assertTrue(self.bootstrap.pyamplicol_is_installed()) + self.assertEqual(self.bootstrap.installed_pyamplicol(), + {'version': '7.0.0', 'source_hash': 'abc123'}) + + def test_source_directory_is_not_an_install(self): + with mock.patch.object(sys, 'path', [str(MG_ROOT)]): + self.assertFalse(self.bootstrap.pyamplicol_is_installed()) + + def hash_check(self, stdout, returncode=0): + result = subprocess.CompletedProcess([], returncode, stdout, 'boom') + with mock.patch('subprocess.run', return_value=result): + return self.bootstrap.check_pyamplicol_source_hash() + + def test_matching_source_hash(self): + self.system_package() + self.assertEqual(self.hash_check('abc123\n')['version'], '7.0.0') + + def test_source_hash_mismatch_is_an_error(self): + self.system_package() + with self.assertRaisesRegex(RuntimeError, 'install pyamplicol'): + self.hash_check('def456\n') + + def test_source_hash_failure_is_an_error(self): + self.system_package() + with self.assertRaisesRegex(RuntimeError, 'boom'): + self.hash_check('', returncode=1) + + def test_build_check(self): + record = self.tmp / 'record.json' + record.write_text(json.dumps({'version': '7.0.0', 'source_hash': 'abc123'})) + self.system_package() + self.bootstrap.check_pyamplicol_build(record) + record.write_text(json.dumps({'version': '6.0.0', 'source_hash': 'old'})) + with self.assertRaisesRegex(SystemExit, 'pyamplicol 6.0.0'): + self.bootstrap.check_pyamplicol_build(record) + + if __name__ == '__main__': unittest.main() From f27d8c70dcd66a9f2387346f82cdd31224291994 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 19:57:41 +0100 Subject: [PATCH 64/88] mg7: take the [amplicol] model restriction from the MG model, not the first '-' of its name --- madgraph/iolibs/export_amplicol.py | 17 +++++----- .../various/test_export_amplicol.py | 33 +++++++++++++++++++ 2 files changed, 42 insertions(+), 8 deletions(-) diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index 18c7185a74..17c4d8c8dd 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -255,14 +255,15 @@ def model_job(model, cache_path): """The UFO model of the output with its restriction, compiled for pyamplicol once per output.""" path = model.get('modelpath') - reference = os.path.basename(model.get('modelpath+restriction')) - restriction = 'default' - if '-' in reference: - suffix = reference.split('-', 1)[1] - if suffix == 'full': - restriction = 'none' - else: - restriction = pjoin(path, 'restrict_%s.dat' % suffix) + reference = model.get('name') + # the suffix after the UFO directory name, which may itself contain '-' + suffix = model.get('restrict_name') + if not suffix: + restriction = 'default' + elif suffix == 'full': + restriction = 'none' + else: + restriction = pjoin(path, 'restrict_%s.dat' % suffix) if not os.path.exists(cache_path): logger.info('Compiling the model %s for AmpliCol' % reference) return {'path': path, 'restriction': restriction, 'cache': cache_path} diff --git a/tests/unit_tests/various/test_export_amplicol.py b/tests/unit_tests/various/test_export_amplicol.py index 50710f9d4c..3b7cb71e57 100644 --- a/tests/unit_tests/various/test_export_amplicol.py +++ b/tests/unit_tests/various/test_export_amplicol.py @@ -100,5 +100,38 @@ def test_check_script_rejects_another_pyamplicol(self): self.assertIn('generated with pyamplicol 0.0.0', result.stderr) +class FakeModel(dict): + """The attributes of a restricted MG model that model_job reads.""" + + def __init__(self, path, name): + super().__init__(name=name, version_tag=path + '##0') + + def get(self, key): + import madgraph.core.base_objects as base_objects + if key in ('modelpath', 'restrict_name'): + return base_objects.Model.get(self, key) + return self[key] + + +class TestModelJob(unittest.TestCase): + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmp) + self.path = os.path.join(self.tmp, 'my-model_UFO') + os.mkdir(self.path) + + def restriction(self, name): + job = export_amplicol.model_job(FakeModel(self.path, name), + os.path.join(self.tmp, 'cache')) + return job['restriction'] + + def test_restrictions_of_a_model_name_with_dashes(self): + self.assertEqual(self.restriction('my-model_UFO'), 'default') + self.assertEqual(self.restriction('my-model_UFO-full'), 'none') + self.assertEqual(self.restriction('my-model_UFO-no_b-mass'), + os.path.join(self.path, 'restrict_no_b-mass.dat')) + + if __name__ == '__main__': unittest.main() From ea2f35f6f0eb1ea5df97e8315c7ae69c0cb07389 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 19:58:28 +0100 Subject: [PATCH 65/88] mg7: key the cached [amplicol] model by model, restriction and pyamplicol --- madgraph/iolibs/amplicol_worker.py | 20 ++++++- madgraph/iolibs/export_amplicol.py | 36 +++++++++++-- .../various/test_export_amplicol.py | 52 +++++++++++++++++++ 3 files changed, 101 insertions(+), 7 deletions(-) diff --git a/madgraph/iolibs/amplicol_worker.py b/madgraph/iolibs/amplicol_worker.py index 22daf01390..ee3237a7c9 100644 --- a/madgraph/iolibs/amplicol_worker.py +++ b/madgraph/iolibs/amplicol_worker.py @@ -26,14 +26,30 @@ def model_source(pyamplicol, model): - """The model compiled for pyamplicol, cached in the output.""" + """The model compiled for pyamplicol, cached in the output and compiled + again unless the key of the cache is model['key'].""" import os - if not os.path.exists(model['cache']): + import shutil + try: + with open(model['key_file']) as stream: + cached = json.load(stream).get('key') + except (OSError, ValueError, AttributeError): + cached = None + if cached != model['key'] or not os.path.exists(model['cache']): from pyamplicol.config import EvaluatorConfig, JITConfig + if os.path.exists(model['key_file']): + os.remove(model['key_file']) + if os.path.isdir(model['cache']): + shutil.rmtree(model['cache']) + elif os.path.exists(model['cache']): + os.remove(model['cache']) source = pyamplicol.ModelSource.from_path(model['path'], restriction=model['restriction']) source.compile(prepared_output=model['cache'], evaluator=EvaluatorConfig(jit=JITConfig(optimization_level=2))) + with open(model['key_file'], 'w') as stream: + json.dump({'key': model['key'], 'sources': model['sources']}, stream, + indent=2) return pyamplicol.ModelSource.from_path(model['cache']) diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index 17c4d8c8dd..fbbabff300 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -27,6 +27,7 @@ representative. """ +import hashlib import json import logging import os @@ -251,9 +252,10 @@ def coupling_orders(procdef, model): for order, weight in hierarchy.items() if weight > lowest} -def model_job(model, cache_path): +def model_job(model, cache_path, pyamplicol=None): """The UFO model of the output with its restriction, compiled for - pyamplicol once per output.""" + pyamplicol once per output and again when the model, its restriction or + pyamplicol change: the cache is keyed by a sidecar json file.""" path = model.get('modelpath') reference = model.get('name') # the suffix after the UFO directory name, which may itself contain '-' @@ -264,9 +266,32 @@ def model_job(model, cache_path): restriction = 'none' else: restriction = pjoin(path, 'restrict_%s.dat' % suffix) - if not os.path.exists(cache_path): + sources = {'model': os.path.realpath(path), 'restriction': restriction, + 'restriction_card': file_hash(pjoin(path, 'restrict_default.dat') + if restriction == 'default' else restriction), + 'pyamplicol': pyamplicol} + key = hashlib.sha256(json.dumps(sources, sort_keys=True).encode()).hexdigest() + key_file = os.path.splitext(cache_path)[0] + '.json' + if not (os.path.exists(cache_path) and cached_model_key(key_file) == key): logger.info('Compiling the model %s for AmpliCol' % reference) - return {'path': path, 'restriction': restriction, 'cache': cache_path} + return {'path': path, 'restriction': restriction, 'cache': cache_path, + 'key': key, 'key_file': key_file, 'sources': sources} + + +def file_hash(path): + if not os.path.isfile(path): + return None + with open(path, 'rb') as stream: + return hashlib.sha256(stream.read()).hexdigest() + + +def cached_model_key(key_file): + """The key of the cached model, see amplicol_worker.model_source.""" + try: + with open(key_file) as stream: + return json.load(stream).get('key') + except (OSError, ValueError, AttributeError): + return None def representative_processes(proc_defs, model): @@ -426,7 +451,8 @@ def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None, with tempfile.TemporaryDirectory(prefix='amplicol_job_') as scratch: job = pjoin(scratch, 'job.json') with open(job, 'w') as stream: - json.dump({'model': model_job(model, pjoin(subproc_dir, MODEL_FILE)), + json.dump({'model': model_job(model, pjoin(subproc_dir, MODEL_FILE), + pyamplicol), 'processes': processes, 'result': pjoin(scratch, 'providers.json')}, stream) run_worker(job) diff --git a/tests/unit_tests/various/test_export_amplicol.py b/tests/unit_tests/various/test_export_amplicol.py index 3b7cb71e57..5e5cf8c533 100644 --- a/tests/unit_tests/various/test_export_amplicol.py +++ b/tests/unit_tests/various/test_export_amplicol.py @@ -16,6 +16,7 @@ instead of pyamplicol.""" import io +import json import os import shutil import subprocess @@ -133,5 +134,56 @@ def test_restrictions_of_a_model_name_with_dashes(self): os.path.join(self.path, 'restrict_no_b-mass.dat')) +class TestModelCache(unittest.TestCase): + """The compiled model is reused only for the same model, restriction and + pyamplicol.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmp) + self.path = os.path.join(self.tmp, 'sm') + os.mkdir(self.path) + self.cache = os.path.join(self.tmp, export_amplicol.MODEL_FILE) + self.compiled = [] + source = mock.Mock() + source.compile.side_effect = lambda prepared_output, evaluator: ( + self.compiled.append(prepared_output), + open(prepared_output, 'w').close()) + self.pyamplicol = mock.Mock() + self.pyamplicol.ModelSource.from_path.return_value = source + config = mock.Mock() + patcher = mock.patch.dict(sys.modules, {'pyamplicol': self.pyamplicol, + 'pyamplicol.config': config}) + patcher.start() + self.addCleanup(patcher.stop) + + def compile(self, version): + from madgraph.iolibs import amplicol_worker + job = export_amplicol.model_job(FakeModel(self.path, 'sm'), self.cache, + {'version': version, 'source_hash': 'h'}) + amplicol_worker.model_source(self.pyamplicol, json.loads(json.dumps(job))) + return job + + def test_cache_is_reused_for_the_same_pyamplicol(self): + first = self.compile('1') + self.assertEqual(self.compile('1')['key'], first['key']) + self.assertEqual(len(self.compiled), 1) + + def test_cache_is_rebuilt_for_another_pyamplicol(self): + self.compile('1') + self.compile('2') + self.assertEqual(len(self.compiled), 2) + + def test_cache_is_rebuilt_for_another_restriction_card(self): + card = os.path.join(self.path, 'restrict_default.dat') + with open(card, 'w') as stream: + stream.write('a') + self.compile('1') + with open(card, 'w') as stream: + stream.write('b') + self.compile('1') + self.assertEqual(len(self.compiled), 2) + + if __name__ == '__main__': unittest.main() From 09ff3817f160be25416f1b0a98047f2550eb5939 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 19:59:48 +0100 Subject: [PATCH 66/88] pyamplicol: read exact-runtime payloads through payloads.pacbin when packed --- .../src/pyamplicol/artifacts/manifest.py | 111 ++++++++++++++++++ .../pyamplicol/runtime/_evaluator_payloads.py | 10 +- pyamplicol/src/pyamplicol/runtime/backend.py | 9 +- .../runtime/eager_exact/_contracts.py | 31 +++-- .../runtime/eager_exact/_executor.py | 11 +- .../pyamplicol/runtime/eager_exact/_plan.py | 11 +- .../runtime/recurrence_exact/_color.py | 10 +- .../runtime/recurrence_exact/_executor.py | 12 +- .../src/pyamplicol/runtime/symbolica_exact.py | 11 +- .../tests/unit/test_artifact_payloads.py | 78 ++++++++++++ 10 files changed, 249 insertions(+), 45 deletions(-) create mode 100644 pyamplicol/tests/unit/test_artifact_payloads.py diff --git a/pyamplicol/src/pyamplicol/artifacts/manifest.py b/pyamplicol/src/pyamplicol/artifacts/manifest.py index c8e19ada7d..a64266f0f2 100644 --- a/pyamplicol/src/pyamplicol/artifacts/manifest.py +++ b/pyamplicol/src/pyamplicol/artifacts/manifest.py @@ -3,6 +3,7 @@ from __future__ import annotations +import hashlib import importlib import json import re @@ -23,6 +24,12 @@ verify_native_module, ) from pyamplicol.api.errors import ArtifactError, CompatibilityError +from pyamplicol.generation.evaluator_container import ( + PacbinError, + PacbinMember, + PacbinMemberKind, + PacbinReader, +) from .security import ( confined_path, @@ -37,6 +44,8 @@ ARTIFACT_IDENTITY_KIND = "pyamplicol-runtime-payload-identity" ARTIFACT_IDENTITY_SCHEMA_VERSION = 1 PORTABLE_64LE_TARGET = "portable-64le" +PAYLOAD_CONTAINER_PATH = "payloads.pacbin" +PAYLOAD_CONTAINER_ROLE = "payload-container" _SHA256 = re.compile(r"^[a-f0-9]{64}$") _GIT_REVISION = re.compile(r"^[a-f0-9]{40}$") _PUBLIC_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:+,~-]{0,254}$") @@ -906,6 +915,97 @@ def _validate_target_compatibility( ) +class ArtifactPayloads: + """Read declared payloads stored as files or packed in payloads.pacbin. + + The container index is checked against the payload declarations once; + each packed read opens its own stream and verifies the member digest. + """ + + def __init__(self, manifest: ArtifactManifest) -> None: + self.root = manifest.root + self._records = {record.path: record for record in manifest.payloads} + self.container_path: Path | None = None + self._members: Mapping[str, PacbinMember] = {} + containers = [ + record + for record in manifest.payloads + if record.role == PAYLOAD_CONTAINER_ROLE + ] + if not containers: + return + if len(containers) != 1 or containers[0].path != PAYLOAD_CONTAINER_PATH: + raise ArtifactError( + "an artifact declares at most one payload container, at " + f"{PAYLOAD_CONTAINER_PATH!r}" + ) + container = containers[0] + path = confined_path(self.root, container.path) + try: + with PacbinReader.open(path, verify_payloads=False) as reader: + index = reader.index + except (OSError, PacbinError) as exc: + raise ArtifactError( + f"could not open payload container {container.path}: {exc}" + ) from exc + if index.file_size != container.size_bytes: + raise ArtifactError("payload container size disagrees with its declaration") + for member in index.members: + record = self._records.get(member.logical_path) + if ( + record is None + or member.kind is not PacbinMemberKind.DECLARED_PAYLOAD + or record.role == PAYLOAD_CONTAINER_ROLE + or record.executable + or member.length != record.size_bytes + or member.sha256 != record.sha256 + ): + raise ArtifactError( + f"invalid payload container member {member.logical_path!r}" + ) + self.container_path = path + self._members = {member.logical_path: member for member in index.members} + + def is_packed(self, relative: str) -> bool: + """Whether ``relative`` is stored in the payload container.""" + + return normalize_relative_path(relative) in self._members + + def read_bytes(self, relative: str) -> bytes: + """The bytes of one payload, wherever it is stored.""" + + normalized = normalize_relative_path(relative) + member = self._members.get(normalized) + if member is None or self.container_path is None: + return confined_path(self.root, normalized).read_bytes() + with self.container_path.open("rb") as stream: + stream.seek(member.offset) + payload = stream.read(member.length) + if len(payload) != member.length: + raise ArtifactError(f"packed payload is truncated: {normalized}") + if hashlib.sha256(payload).hexdigest() != member.sha256: + raise ArtifactError(f"packed payload digest mismatch: {normalized}") + return payload + + def read_text(self, relative: str) -> str: + """The UTF-8 text of one payload, wherever it is stored.""" + + return self.read_bytes(relative).decode("utf-8") + + def verify_packed(self) -> None: + """Stream and verify the digest of every packed payload.""" + + if self.container_path is None: + return + try: + with PacbinReader.open(self.container_path, verify_payloads=True): + pass + except (OSError, PacbinError) as exc: + raise ArtifactError( + f"payload container verification failed: {exc}" + ) from exc + + def validate_payloads( manifest: ArtifactManifest, *, @@ -915,7 +1015,15 @@ def validate_payloads( raise ArtifactError("artifact manifest identity digest mismatch") _validate_target_compatibility(manifest, expected_target=expected_target) seen = {record.path for record in manifest.payloads} + payloads = ArtifactPayloads(manifest) + payloads.verify_packed() for record in manifest.payloads: + if payloads.is_packed(record.path): + if confined_path(manifest.root, record.path, must_exist=False).exists(): + raise ArtifactError( + f"packed artifact payload is also stored as a file: {record.path}" + ) + continue path = confined_path(manifest.root, record.path) if path.stat().st_size != record.size_bytes: raise ArtifactError(f"artifact payload size mismatch: {record.path}") @@ -1000,9 +1108,12 @@ def load_manifest( "ARTIFACT_IDENTITY_KIND", "ARTIFACT_IDENTITY_SCHEMA_VERSION", "MANIFEST_NAME", + "PAYLOAD_CONTAINER_PATH", + "PAYLOAD_CONTAINER_ROLE", "PORTABLE_64LE_TARGET", "SCHEMA_VERSION", "ArtifactManifest", + "ArtifactPayloads", "PayloadRecord", "artifact_identity_extension", "canonical_manifest_bytes", diff --git a/pyamplicol/src/pyamplicol/runtime/_evaluator_payloads.py b/pyamplicol/src/pyamplicol/runtime/_evaluator_payloads.py index a0270b1fd4..788c14c37f 100644 --- a/pyamplicol/src/pyamplicol/runtime/_evaluator_payloads.py +++ b/pyamplicol/src/pyamplicol/runtime/_evaluator_payloads.py @@ -8,7 +8,11 @@ from pathlib import Path from pyamplicol.api.errors import ArtifactError, CompatibilityError -from pyamplicol.artifacts.manifest import ArtifactManifest, PayloadRecord +from pyamplicol.artifacts.manifest import ( + ArtifactManifest, + ArtifactPayloads, + PayloadRecord, +) from pyamplicol.artifacts.security import confined_path, normalize_relative_path from pyamplicol.generation.evaluator_container import ( PacbinError, @@ -35,6 +39,8 @@ class ExactEvaluatorPayloadResolver: def __init__(self, manifest: ArtifactManifest) -> None: self._root = manifest.root self._loose = {record.path: record for record in manifest.payloads} + # declared payloads, loose or packed in payloads.pacbin + self.payloads = ArtifactPayloads(manifest) self._container_path: Path | None = None self._members: Mapping[str, PacbinMember] = {} @@ -152,7 +158,7 @@ def read_exact_state( record = self._loose.get(normalized) if record is not None: try: - return confined_path(self._root, normalized).read_bytes() + return self.payloads.read_bytes(normalized) except OSError as exc: raise CompatibilityError( f"could not read retained Symbolica evaluator state " diff --git a/pyamplicol/src/pyamplicol/runtime/backend.py b/pyamplicol/src/pyamplicol/runtime/backend.py index f5bae308f3..d5b8fdf0f8 100644 --- a/pyamplicol/src/pyamplicol/runtime/backend.py +++ b/pyamplicol/src/pyamplicol/runtime/backend.py @@ -45,6 +45,7 @@ WarmUpResult, ) from pyamplicol.artifacts import MANIFEST_NAME, ArtifactManifest, load_manifest +from pyamplicol.artifacts.manifest import ArtifactPayloads from pyamplicol.reporting import ( ProgressEnd, ProgressSink, @@ -405,12 +406,12 @@ def _has_reusable_runtime_selector_contract( ) -> bool: relative = str(process["physics_path"]) try: - payload = json.loads((manifest.root / relative).read_text(encoding="utf-8")) + payload = json.loads(ArtifactPayloads(manifest).read_bytes(relative)) except OSError as exc: raise ArtifactError( f"could not read runtime selector metadata {relative}: {exc}" ) from exc - except json.JSONDecodeError as exc: + except (UnicodeDecodeError, json.JSONDecodeError) as exc: raise ArtifactError( f"invalid runtime selector metadata {relative}: {exc}" ) from exc @@ -1622,8 +1623,8 @@ def validation_momenta(self) -> Momenta | None: f"process {process_id!r} must declare exactly one validation point" ) try: - raw = json.loads((manifest.root / payloads[0].path).read_text("utf-8")) - except (OSError, json.JSONDecodeError) as exc: + raw = json.loads(ArtifactPayloads(manifest).read_bytes(payloads[0].path)) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: raise ArtifactError( f"could not read validation point for process {process_id!r}: {exc}" ) from exc diff --git a/pyamplicol/src/pyamplicol/runtime/eager_exact/_contracts.py b/pyamplicol/src/pyamplicol/runtime/eager_exact/_contracts.py index b39c42f549..b3c4c0d6c0 100644 --- a/pyamplicol/src/pyamplicol/runtime/eager_exact/_contracts.py +++ b/pyamplicol/src/pyamplicol/runtime/eager_exact/_contracts.py @@ -15,8 +15,12 @@ PROCESS_ARTIFACT_SCHEMA_VERSION, ) from pyamplicol.api.errors import ArtifactError, CompatibilityError, EvaluationError -from pyamplicol.artifacts.manifest import ArtifactManifest, PayloadRecord -from pyamplicol.artifacts.security import confined_path, normalize_relative_path +from pyamplicol.artifacts.manifest import ( + ArtifactManifest, + ArtifactPayloads, + PayloadRecord, +) +from pyamplicol.artifacts.security import normalize_relative_path from pyamplicol.generation.eager_tables import ( EAGER_PLAN_ABI, EAGER_RUNTIME_CAPABILITY, @@ -180,10 +184,12 @@ def _integer(value: object, context: str, *, minimum: int = 0) -> int: return value -def _read_json(path: Path, context: str) -> Mapping[str, object]: +def _read_payload_json( + payloads: ArtifactPayloads, path: str, context: str +) -> Mapping[str, object]: try: - value = json.loads(path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError) as exc: + value = json.loads(payloads.read_bytes(path)) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: raise ArtifactError(f"could not load {context}: {exc}") from exc return _mapping(value, context) @@ -199,10 +205,19 @@ def _joined_payload_path(prefix: str, member: str) -> str: @dataclass(frozen=True, slots=True) class _PayloadIndex: records: Mapping[str, PayloadRecord] + files: ArtifactPayloads @classmethod def from_manifest(cls, manifest: ArtifactManifest) -> _PayloadIndex: - return cls({record.path: record for record in manifest.payloads}) + return cls( + {record.path: record for record in manifest.payloads}, + ArtifactPayloads(manifest), + ) + + def read_json(self, path: str, context: str) -> Mapping[str, object]: + """Load one declared JSON payload, loose or packed.""" + + return _read_payload_json(self.files, path, context) def require( self, @@ -469,9 +484,9 @@ def _load_table( ) global_path = _joined_payload_path(process_prefix, path) payloads.require(global_path, role="evaluator-state", process_id=process_id) - table_path = confined_path(process_root, path) + del process_root try: - content = table_path.read_bytes() + content = payloads.files.read_bytes(global_path) except OSError as exc: raise ArtifactError(f"could not read {context}: {exc}") from exc if len(content) != count * expected_size: diff --git a/pyamplicol/src/pyamplicol/runtime/eager_exact/_executor.py b/pyamplicol/src/pyamplicol/runtime/eager_exact/_executor.py index b7a0b58119..4bcf324354 100644 --- a/pyamplicol/src/pyamplicol/runtime/eager_exact/_executor.py +++ b/pyamplicol/src/pyamplicol/runtime/eager_exact/_executor.py @@ -12,7 +12,6 @@ from pyamplicol.api.protocols import Momenta from pyamplicol.api.results import ResolvedEvaluation from pyamplicol.artifacts import load_manifest -from pyamplicol.artifacts.security import confined_path from pyamplicol.runtime._color_topology_exact import ( apply_exact_color_replay_input_mapping, parse_exact_color_topology_replay, @@ -29,7 +28,6 @@ _KernelLoader, _mapping, _PayloadIndex, - _read_json, ) from pyamplicol.runtime.eager_exact._execution import _evaluate_point from pyamplicol.runtime.eager_exact._plan import _EagerExactPlan @@ -86,12 +84,11 @@ def __init__( payloads.require( physics_path, role="runtime-physics", process_id=representative_id ) - execution = _read_json( - confined_path(self._artifact, execution_records[0].path), - "eager execution metadata", + execution = payloads.read_json( + execution_records[0].path, "eager execution metadata" ) - representative_physics = _read_json( - confined_path(self._artifact, physics_path), "runtime physics metadata" + representative_physics = payloads.read_json( + physics_path, "runtime physics metadata" ) axes = native_physics_axes(native_runtime, representative_physics) physics = axes.public_physics diff --git a/pyamplicol/src/pyamplicol/runtime/eager_exact/_plan.py b/pyamplicol/src/pyamplicol/runtime/eager_exact/_plan.py index 853fe79a80..719b958715 100644 --- a/pyamplicol/src/pyamplicol/runtime/eager_exact/_plan.py +++ b/pyamplicol/src/pyamplicol/runtime/eager_exact/_plan.py @@ -11,7 +11,7 @@ from pyamplicol.api.errors import ArtifactError, CompatibilityError, EvaluationError from pyamplicol.artifacts.manifest import ArtifactManifest -from pyamplicol.artifacts.security import confined_path, normalize_relative_path +from pyamplicol.artifacts.security import normalize_relative_path from pyamplicol.generation.eager_tables import ( EAGER_KERNEL_ABI, EAGER_OUTPUT_FACTOR_NONE, @@ -48,7 +48,6 @@ _load_table, _mapping, _PayloadIndex, - _read_json, _sequence, _validate_execution_header, ) @@ -409,9 +408,7 @@ def load( pack_path = normalize_relative_path(pack_path) payload_root_name = normalize_relative_path(payload_root_name) payloads.require(pack_path, role="evaluator-manifest", process_id=None) - pack_payload = _read_json( - confined_path(artifact_root, pack_path), "eager kernel pack" - ) + pack_payload = payloads.read_json(pack_path, "eager kernel pack") raw_backend = pack_payload.get("backend") if raw_backend not in _SUPPORTED_PREPARED_BACKENDS: raise CompatibilityError( @@ -1109,9 +1106,7 @@ def _load_exact_kernel_pack( pack_path = normalize_relative_path(pack_path) payload_root_name = normalize_relative_path(payload_root_name) payloads.require(pack_path, role="evaluator-manifest", process_id=None) - pack_payload = _read_json( - confined_path(artifact_root, pack_path), "eager kernel pack" - ) + pack_payload = payloads.read_json(pack_path, "eager kernel pack") raw_backend = pack_payload.get("backend") if raw_backend not in _SUPPORTED_PREPARED_BACKENDS: raise CompatibilityError(f"unsupported prepared eager backend {raw_backend!r}") diff --git a/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_color.py b/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_color.py index 588fed2cb1..7b9e5d2724 100644 --- a/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_color.py +++ b/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_color.py @@ -12,8 +12,12 @@ from pathlib import Path, PurePosixPath from pyamplicol.api.errors import ArtifactError, CompatibilityError -from pyamplicol.artifacts.manifest import ArtifactManifest, PayloadRecord -from pyamplicol.artifacts.security import confined_path, normalize_relative_path +from pyamplicol.artifacts.manifest import ( + ArtifactManifest, + ArtifactPayloads, + PayloadRecord, +) +from pyamplicol.artifacts.security import normalize_relative_path from pyamplicol.color.symmetric_group import ( _inverse_permutation, _lexicographic_permutation_rank, @@ -252,7 +256,7 @@ def _load_recurrence_color_contraction( "recurrence color-contraction payload disagrees with execution metadata" ) try: - payload = confined_path(artifact_root, logical_path).read_bytes() + payload = ArtifactPayloads(manifest).read_bytes(logical_path) except OSError as exc: raise ArtifactError( f"could not read recurrence color-contraction payload: {exc}" diff --git a/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_executor.py b/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_executor.py index 8bde998036..930ea0c5d6 100644 --- a/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_executor.py +++ b/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_executor.py @@ -12,7 +12,6 @@ from pyamplicol.api.protocols import Momenta from pyamplicol.api.results import ResolvedEvaluation from pyamplicol.artifacts import load_manifest -from pyamplicol.artifacts.security import confined_path from pyamplicol.runtime._evaluator_payloads import ExactEvaluatorPayloadResolver from pyamplicol.runtime._native_selection import ( native_physics_axes, @@ -24,7 +23,6 @@ _KernelLoader, _mapping, _PayloadIndex, - _read_json, _sequence, ) from pyamplicol.runtime.symbolica_exact import ( @@ -176,13 +174,11 @@ def __init__( role="runtime-physics", process_id=representative_id, ) - execution = _read_json( - confined_path(self._artifact, execution_records[0].path), - "recurrence execution metadata", + execution = payloads.read_json( + execution_records[0].path, "recurrence execution metadata" ) - representative_physics = _read_json( - confined_path(self._artifact, physics_path), - "runtime physics metadata", + representative_physics = payloads.read_json( + physics_path, "runtime physics metadata" ) axes = native_physics_axes(native_runtime, representative_physics) if execution.get("key") != selection.representative_process_key: diff --git a/pyamplicol/src/pyamplicol/runtime/symbolica_exact.py b/pyamplicol/src/pyamplicol/runtime/symbolica_exact.py index 30f357a2bc..1aef0fcbe1 100644 --- a/pyamplicol/src/pyamplicol/runtime/symbolica_exact.py +++ b/pyamplicol/src/pyamplicol/runtime/symbolica_exact.py @@ -27,7 +27,8 @@ from pyamplicol.api.protocols import Momenta from pyamplicol.api.results import ResolvedEvaluation from pyamplicol.artifacts import load_manifest -from pyamplicol.artifacts.security import confined_path, normalize_relative_path +from pyamplicol.artifacts.manifest import ArtifactPayloads +from pyamplicol.artifacts.security import normalize_relative_path from pyamplicol.runtime._color_topology_exact import ( apply_exact_color_replay_input_mapping, parse_exact_color_topology_replay, @@ -508,12 +509,12 @@ def __init__(self, artifact: Path, process_id: str, native_runtime: Any) -> None ) try: self._execution = json.loads( - (artifact / records[0].path).read_text(encoding="utf-8") + self._payloads.payloads.read_text(records[0].path) ) representative_physics = json.loads( - confined_path(artifact, physics_path).read_text(encoding="utf-8") + self._payloads.payloads.read_text(physics_path) ) - except (OSError, json.JSONDecodeError) as exc: + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: raise ArtifactError( f"could not load exact-runtime metadata: {exc}" ) from exc @@ -841,7 +842,7 @@ def _derive_model_parameters( "exact parameter derivation requires one compiled model" ) model = json.loads( - confined_path(self._artifact, model_records[0].path).read_text() + ArtifactPayloads(manifest).read_text(model_records[0].path) )["ir"] self._exact_compiled_model = model definitions = { diff --git a/pyamplicol/tests/unit/test_artifact_payloads.py b/pyamplicol/tests/unit/test_artifact_payloads.py new file mode 100644 index 0000000000..eb5618abaa --- /dev/null +++ b/pyamplicol/tests/unit/test_artifact_payloads.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: 0BSD +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from pyamplicol import ArtifactError +from pyamplicol.artifacts import load_manifest +from pyamplicol.artifacts.manifest import ArtifactPayloads, validate_payloads +from pyamplicol.generation.evaluator_container import PacbinReader +from pyamplicol.runtime._evaluator_payloads import ExactEvaluatorPayloadResolver +from pyamplicol.runtime.eager_exact._contracts import _PayloadIndex +from tests.unit.test_artifacts import _build_packable + +_EXECUTION = "processes/dd_to_z/execution.json" +_BINDING = "processes/dd_to_z/recurrence-binding.bin" + + +@pytest.mark.parametrize("pack", [False, True]) +def test_payloads_resolve_loose_and_packed_layouts(tmp_path: Path, pack: bool) -> None: + root = tmp_path / "artifact" + _build_packable(root, pack=pack) + manifest = load_manifest(root, verify_payloads=True) + payloads = ArtifactPayloads(manifest) + + assert payloads.is_packed(_EXECUTION) is pack + assert (payloads.container_path is not None) is pack + assert json.loads(payloads.read_text(_EXECUTION)) == {"plan": [1, 2, 3]} + assert payloads.read_bytes(_BINDING) == b"binding" + # nested containers and top-level payloads stay loose + assert not payloads.is_packed("processes/dd_to_z/eager-runtime.pacbin") + assert ( + payloads.read_bytes("physics/process.json") + == (root / "physics/process.json").read_bytes() + ) + assert _PayloadIndex.from_manifest(manifest).read_json(_EXECUTION, "execution") == { + "plan": [1, 2, 3] + } + assert ExactEvaluatorPayloadResolver(manifest).payloads.read_bytes(_BINDING) == ( + b"binding" + ) + + +def _corrupt_member(root: Path, logical_path: str) -> None: + container = root / "payloads.pacbin" + with PacbinReader.open(container, verify_payloads=False) as reader: + offset = reader.member(logical_path).offset + data = bytearray(container.read_bytes()) + data[offset] ^= 0xFF + container.chmod(0o644) + container.write_bytes(bytes(data)) + + +def test_packed_payload_reads_verify_member_digests(tmp_path: Path) -> None: + root = tmp_path / "artifact" + _build_packable(root, pack=True) + _corrupt_member(root, _BINDING) + manifest = load_manifest(root) + payloads = ArtifactPayloads(manifest) + + assert json.loads(payloads.read_text(_EXECUTION)) == {"plan": [1, 2, 3]} + with pytest.raises(ArtifactError, match="digest mismatch"): + payloads.read_bytes(_BINDING) + with pytest.raises(ArtifactError, match="payload container verification"): + validate_payloads(manifest) + + +def test_packed_payloads_may_not_also_be_files(tmp_path: Path) -> None: + root = tmp_path / "artifact" + _build_packable(root, pack=True) + shadow = root / _EXECUTION + shadow.parent.mkdir(parents=True, exist_ok=True) + shadow.write_text('{"plan": []}', encoding="utf-8") + + with pytest.raises(ArtifactError, match="also stored as a file"): + load_manifest(root, verify_payloads=True) From 059924adc16c6d457f2fb9203eb6f236ed0c6257 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:00:05 +0100 Subject: [PATCH 67/88] madspace: freeze the step maxima independently of the thread timing The freeze compared count_accepted, which is scaled by the first maximum at the time of the step decision and so depends on the thread timing. Count the expected accepted events of each step on its own, like _freeze_count. --- .../madspace/driver/channel_generator.hpp | 3 ++ madspace/src/driver/channel_generator.cpp | 20 ++++++++- madspace/tests/test_weighting_steps.py | 45 ++++++++++++++++++- 3 files changed, 65 insertions(+), 3 deletions(-) diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 2f661edb14..a958713702 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -334,9 +334,12 @@ class ChannelEventGenerator { bool max_weight_fixed = false; std::vector large_weights; double count_accepted = 0.; + // like _freeze_count, the expected accepted events of this step alone + double freeze_count = 0.; }; void apply_truncation_budget(); + void apply_step_truncation_budget(std::size_t step, std::size_t count_target); void update_status_integrals(); void unweight_step_file(MixMaxRandom& rand_gen); StepChainResult step_chain(EventBuffer& weight_buffer, std::size_t index) const; diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index c27bba0fae..26902eb102 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -751,6 +751,7 @@ void ChannelEventGenerator::clear_events() { } state.large_weights.clear(); state.count_accepted = 0.; + state.freeze_count = 0.; } for (auto& hist : _histograms) { std::fill(hist.bin_values.begin(), hist.bin_values.end(), 0.); @@ -1228,8 +1229,9 @@ void ChannelEventGenerator::update_step_max_weight( ) { auto& state = _step_states.at(step - 1); auto& settings = _step_settings.at(step - 1); + // not count_accepted, it depends on the first maximum at the time of the call if (state.max_weight_fixed || - state.count_accepted > settings.freeze_max_weight_after) { + state.freeze_count > settings.freeze_max_weight_after) { return; } double min_nonzero = 0.; @@ -1245,8 +1247,21 @@ void ChannelEventGenerator::update_step_max_weight( if (state.max_weight == 0) { state.max_weight = min_nonzero; } + std::sort(state.large_weights.begin(), state.large_weights.end(), std::greater{}); + apply_step_truncation_budget(step, count_target); + if (state.max_weight > 0) { + for (double value : values) { + state.freeze_count += std::min(1., std::abs(value) / state.max_weight); + } + } +} + +void ChannelEventGenerator::apply_step_truncation_budget( + std::size_t step, std::size_t count_target +) { + auto& state = _step_states.at(step - 1); + auto& settings = _step_settings.at(step - 1); auto& large_weights = state.large_weights; - std::sort(large_weights.begin(), large_weights.end(), std::greater{}); // same truncation as for the first step, see apply_truncation_budget() double max_truncation = settings.max_overweight_truncation * static_cast(std::min(count_target, settings.freeze_max_weight_after)); @@ -1266,6 +1281,7 @@ void ChannelEventGenerator::update_step_max_weight( for (std::size_t i = step - 1; i < _step_states.size(); ++i) { _step_states.at(i).count_accepted *= factor; } + state.freeze_count *= factor; state.max_weight = value; _unweighted_count = 0; _unweighted_accept_count = 0; diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py index 8a3cc78651..84477eeb9d 100644 --- a/madspace/tests/test_weighting_steps.py +++ b/madspace/tests/test_weighting_steps.py @@ -77,6 +77,8 @@ def generate( override_outputs=(), integrand_outputs=("matrix_element",), step_channel=-1, + step_kwargs=None, + config_kwargs=None, ): """Survey and generate, returns the status, the combined events and the channels. step_me and step_outputs can be lists, one entry per step.""" @@ -93,6 +95,8 @@ def generate( config = make_config() config.target_count = target_count config.cpu_batch_size = 500 + for key, value in (config_kwargs or {}).items(): + setattr(config, key, value) channels = [ ms.ChannelEventGenerator( contexts=[context], @@ -107,7 +111,12 @@ def generate( for i, pt_min in enumerate(pt_mins) ] steps = [ - ms.WeightingStep([me], mode=mode, override_outputs=list(override_outputs)) + ms.WeightingStep( + [me], + mode=mode, + override_outputs=list(override_outputs), + **(step_kwargs or {}), + ) for me, mode in zip(step_mes, step_modes) ] groups = [ms.WeightingGroup(list(grouped), steps)] if steps else [] @@ -197,6 +206,40 @@ def test_reproducible(tmp_path, step_modes): assert status_a.mean == status_b.mean +def heavy_tailed_me(momenta): + return surrogate_me(momenta) * ( + 1.0 + 0.05 / (0.002 + np.abs(np.sin(momenta[:, 3, 1] / 7.0))) + ) + + +def test_reproducible_step_freeze(tmp_path): + """The step maximum freezes while the first one still rises, in an order + relative to the step decisions that depends on the threads.""" + results = [ + generate( + str(tmp_path / str(threads)), + heavy_tailed_me, + ["unweight"], + threads=threads, + pt_mins=(20.0, 40.0, 60.0), + grouped=(0, 1, 2), + step_kwargs={"freeze_max_weight_after": 300}, + config_kwargs={ + "freeze_max_weight_after": 100000, + "max_overweight_truncation": 0.002, + }, + ) + for threads in (1, 2, 4, 8) + ] + status_a, events_a, channels_a = results[0] + for status_b, events_b, channels_b in results[1:]: + assert [c.step_max_weight(1) for c in channels_a] == [ + c.step_max_weight(1) for c in channels_b + ] + assert np.array_equal(events_a, events_b) + assert status_a.mean == status_b.mean + + def test_index_outputs(tmp_path): """A step fills in indices the integrand matrix element does not return, and replaces the others only with override_outputs.""" From dafb6e02aa4131511efcd21ed7793e826951ea5c Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:03:50 +0100 Subject: [PATCH 68/88] pyamplicol: s-channel resonances selected at launch, every compatible configuration The resolved physics lists every s-channel resonance a tree current can carry whose mass and width are not zero in the model, without the threshold and width filters of the default parameters: the launch decides with the parameter card which go on shell. The phase-space channels follow every pairwise nested or disjoint set of colour-singlet resonances, the non-resonant one first and the larger ones next, enumerated once per process and lazily up to the caps. A coloured resonance is listed in channels[].colored_resonances for the launch to check; a singlet decaying to coloured and colourless particles, or never to one complete colour line, is not followed and is warned about. --- pyamplicol/docs/user/umami-api.md | 17 +- pyamplicol/src/pyamplicol/artifacts/umami.py | 210 ++++++++------ .../generation/s_channel_resonances.py | 37 ++- .../unit/test_umami_phasespace_partners.py | 263 ++++++++++++++++-- 4 files changed, 403 insertions(+), 124 deletions(-) diff --git a/pyamplicol/docs/user/umami-api.md b/pyamplicol/docs/user/umami-api.md index 898c79326e..be01f2f5a8 100644 --- a/pyamplicol/docs/user/umami-api.md +++ b/pyamplicol/docs/user/umami-api.md @@ -215,12 +215,17 @@ A provider's metadata contains: A selector lists its phase-space channels in `multichannels`; they differ in where its colour singlets attach to the colour lines (after the antiquark ends, in every order and distribution, as in AmpliCol), in the non-cyclic - order of three or more lines, and in which resonances of the process they - follow, with a non-resonant one always included. A UMAMI channel collects - the selectors linked by shared phase-space channels. Weight the channels of a - selector by a partition of unity, for example their normalised mapping - densities, rather than adding them. A coloured resonance that can go on shell - is rejected for now. + order of three or more lines, and in which compatible (nested or disjoint) + set of colour-singlet resonances of the process they follow, with a + non-resonant one always included. The candidates are every resonance whose + mass and width are not zero in the model: drop those that cannot go on shell + with the parameter-card values, and merge the channels that then agree. A + UMAMI channel collects the selectors linked by shared phase-space channels. + Weight the channels of a selector by a partition of unity, for example their + normalised mapping densities, rather than adding them. +- `channels[].colored_resonances`: the coloured resonances (`pdg`, decay + `legs`, `color`) of its selectors, which the phase-space channels do not + follow; reject the parameters if one of them can go on shell. - `pdg_ids`, `color_orders`, `color_flows` and `helicities`: shared physical configuration tables. Indices are zero-based and colour-flow references can be absent for contracted colour. diff --git a/pyamplicol/src/pyamplicol/artifacts/umami.py b/pyamplicol/src/pyamplicol/artifacts/umami.py index 6501e701f3..7300b9d854 100644 --- a/pyamplicol/src/pyamplicol/artifacts/umami.py +++ b/pyamplicol/src/pyamplicol/artifacts/umami.py @@ -13,11 +13,12 @@ from __future__ import annotations import json +import logging import math from collections.abc import Mapping, Sequence from contextlib import suppress from dataclasses import dataclass -from itertools import combinations_with_replacement, permutations +from itertools import combinations_with_replacement, islice, permutations from typing import TYPE_CHECKING, Any from .api_bundle import ApiBundlePayload @@ -25,6 +26,8 @@ if TYPE_CHECKING: from ..generation.process_grouping import GroupMember +_LOGGER = logging.getLogger(__name__) + FLAVOR_BLIND_ASSUMPTION = ( "flavor-blind matrix elements: ungenerated process-group members are " "evaluated through their relabelled representative (assumed, not proven)" @@ -349,10 +352,10 @@ def _relabelled_runtime(runtime, relabel): # Phase-space partners of one colour order beyond this many are dropped: any -# subset gives correct channel weights. Every resonance configuration keeps one. +# subset gives correct channel weights, as each selector lists its own. _MAX_PARTNERS = 64 -# More resonance candidates than this in one process are not enumerated. -_MAX_RESONANCES = 12 +# Resonance configurations of one process beyond this many are not enumerated. +_MAX_CONFIGURATIONS = 1024 def _roles(pdgs, incoming, particles): @@ -383,36 +386,45 @@ def _line_blocks(word_slots, roles): return blocks -def _configurations(resonances): - """The maximal sets of resonances that can be mapped together, pairwise - nested or disjoint, and the empty set for the non-resonant region.""" - if len(resonances) > _MAX_RESONANCES: - raise ValueError( - f"{len(resonances)} candidate s-channel resonances are too many to " - "enumerate their phase-space configurations" - ) +def _classify_resonances(resonances, roles): + """Split the candidates ``(pdg, legs, color)`` of a process into the colour + singlets a phase-space channel can follow, ``(pdg, legs)``, the coloured + ones and the singlets decaying to coloured and colourless particles.""" + mappable, coloured, mixed = [], [], [] + for pdg, legs, color in resonances: + if color != 1: + coloured.append((pdg, legs, color)) + elif len({roles[slot] == 1 for slot in legs}) > 1: + mixed.append((pdg, legs)) + else: + mappable.append((pdg, legs)) + return mappable, coloured, mixed + + +def _configurations(resonances, limit=_MAX_CONFIGURATIONS): + """Every set of resonances that can be mapped together, pairwise nested or + disjoint: the non-resonant one first, then the larger ones first. Returns + ``(configurations, complete)``, enumerating at most ``limit`` of them.""" def compatible(a, b): return a < b or b < a or not a & b - families = [ - bits - for bits in range(1, 1 << len(resonances)) - if all( - compatible(resonances[i][1], resonances[j][1]) - for i in range(len(resonances)) - for j in range(i + 1, len(resonances)) - if bits >> i & 1 and bits >> j & 1 - ) - ] - maximal = [ - bits - for bits in families - if not any(other != bits and other & bits == bits for other in families) - ] - return [()] + [ - tuple(r for i, r in enumerate(resonances) if bits >> i & 1) for bits in maximal - ] + found, truncated = [], False + + def extend(chosen, start): + nonlocal truncated + found.append(tuple(chosen)) + for index in range(start, len(resonances)): + candidate = resonances[index] + if all(compatible(candidate[1], r[1]) for r in chosen): + if len(found) >= limit: + truncated = True + return + extend([*chosen, candidate], index + 1) + + extend([], 0) + found[1:] = sorted(found[1:], key=len, reverse=True) + return found, not truncated def _resonance_legs(resonance, configuration): @@ -430,19 +442,16 @@ def _resonance_legs(resonance, configuration): def _placements(word_slots, blocks, units): - """Phase-space orders of one colour order: as in AmpliCol, the colour - singlet units follow the antiquark ends of the lines, in every order and - every distribution over the lines, and with three or more lines every - non-cyclic order of the lines is a partner too (the same colour tensor).""" + """Phase-space orders of one colour order, generated lazily: as in AmpliCol, + the colour singlet units follow the antiquark ends of the lines, in every + order and every distribution over the lines, and with three or more lines + every non-cyclic order of the lines is a partner too (the same colour + tensor).""" if blocks is None: - return [list(word_slots) + [slot for unit in units for slot in unit]] - line_orders = [blocks] + ( - [[blocks[0], *rest] for rest in permutations(blocks[1:])][1:] - if len(blocks) > 2 - else [] - ) - orders = [] - for lines in line_orders: + yield list(word_slots) + [slot for unit in units for slot in unit] + return + for rest in permutations(blocks[1:]): + lines = [blocks[0], *rest] for unit_order in permutations(units): for cuts in combinations_with_replacement( range(len(units) + 1), len(lines) - 1 @@ -453,39 +462,30 @@ def _placements(word_slots, blocks, units): order += line for unit in unit_order[bounds[index] : bounds[index + 1]]: order += unit - orders.append(order) - return orders + yield order -def _phasespace_partners(word_slots, singlet_slots, roles, resonances, process): +def _phasespace_partners(word_slots, singlet_slots, roles, configurations): """The phase-space orders and resonance configurations integrating one colour order together, the first one its own: ``[(order, configuration)]`` - with ``configuration`` a tuple of ``(pdg, legs)``.""" + with ``configuration`` one of the process's ``configurations`` of + ``(pdg, legs)``. A resonance decaying to coloured particles applies only if + they are a complete colour line of this order. At most ``_MAX_PARTNERS`` + configurations are kept, in their order, and as many placements of each + as fit.""" blocks = _line_blocks(word_slots, roles) if word_slots else None - applicable = [] - for pdg, legs, color in resonances: - if color != 1: - raise ValueError( - f"{process}: the coloured s-channel resonance {pdg} can go on shell, " - "which [amplicol] phase-space mappings do not follow yet; generate " - "the process with on-shell decays or with the madmatrix backend" - ) - coloured = {slot for slot in legs if roles[slot] != 1} - if not coloured: - applicable.append((pdg, legs)) - elif coloured != legs: - raise ValueError( - f"{process}: the s-channel resonance {pdg} decays to coloured and " - "colourless particles, which [amplicol] phase-space mappings do not " - "follow yet" - ) - elif blocks is not None and any(set(block) == legs for block in blocks): - # a colour singlet decaying to a complete colour line of this order - applicable.append((pdg, legs)) - configurations = _configurations(applicable) - per_configuration = max(1, _MAX_PARTNERS // len(configurations)) + lines = {frozenset(block) for block in blocks or ()} + applicable = [ + configuration + for configuration in configurations + if all( + legs in lines or all(roles[slot] == 1 for slot in legs) + for _, legs in configuration + ) + ][:_MAX_PARTNERS] + per_configuration = max(1, _MAX_PARTNERS // len(applicable)) partners = [] - for configuration in configurations: + for configuration in applicable: top = [ r for r in configuration @@ -499,8 +499,8 @@ def _phasespace_partners(word_slots, singlet_slots, roles, resonances, process): covered = {slot for unit in units for slot in unit} units += [[slot] for slot in singlet_slots if slot not in covered] units.sort(key=min) - orders = _placements(word_slots, blocks, units) - partners += [(order, configuration) for order in orders[:per_configuration]] + orders = islice(_placements(word_slots, blocks, units), per_configuration) + partners += [(order, configuration) for order in orders] return partners @@ -602,7 +602,8 @@ def _add_terms(entry, terms): def _group_phasespace_channels(data, selectors, phasespace_channels): """One UMAMI channel per set of contributions linked by shared phase-space channels: its contributions are integrated together, each over the - phase-space channels listed in its ``multichannels``.""" + phase-space channels listed in its ``multichannels``. The coloured + resonances of its contributions are listed for the launch to check.""" parent = {key: key for key in phasespace_channels} def root(key): @@ -611,16 +612,19 @@ def root(key): key = parent[key] return key - for _, keys in selectors: + for _, keys, _ in selectors: for key in keys[1:]: parent[root(key)] = root(keys[0]) groups: dict[tuple, dict] = {} - for entry, keys in selectors: - group = groups.setdefault(root(keys[0]), {"lookup": {}, "processes": []}) + for entry, keys, colored in selectors: + group = groups.setdefault( + root(keys[0]), {"lookup": {}, "processes": [], "colored": []} + ) entry["multichannels"] = [ group["lookup"].setdefault(key, len(group["lookup"])) for key in keys ] group["processes"].append(entry) + group["colored"] += [r for r in colored if r not in group["colored"]] data["channels"] = [] for group in groups.values(): channels = [phasespace_channels[key] for key in group["lookup"]] @@ -628,6 +632,7 @@ def root(key): { "phasespace_order": channels[0]["phasespace_order"], "phasespace_channels": channels, + "colored_resonances": group["colored"], "processes": group["processes"], } ) @@ -849,9 +854,9 @@ def _build_provider( orbit_entries: dict[tuple, tuple[dict, list[int]]] = {} channel_patterns: dict[tuple, list[int]] = {} # (order, resonances) of every phase-space channel, and the exported - # contributions with the channels that integrate them + # contributions with the channels integrating them and coloured resonances phasespace_channels: dict[tuple, dict] = {} - selectors: list[tuple[dict, list[tuple]]] = [] + selectors: list[tuple[dict, list[tuple], list[dict]]] = [] physical_count = 0 runtime_helicities: list[list[list[int]]] = [] runtime_indices: dict[str, int] = {} @@ -934,6 +939,17 @@ def _build_provider( (pdg, frozenset(relation[slot] for slot in legs), color) for pdg, legs, color in resonances ] + mappable, coloured, unmapped = _classify_resonances(resonances, roles) + configurations, complete = _configurations(mappable) + if not complete and process.color_accuracy == "lc": + _LOGGER.warning( + "%s: only the first %d s-channel resonance configurations are " + "followed by the [amplicol] phase space", + source_id, + len(configurations), + ) + # the legs of a complete colour line in some colour order of the process + line_legs = set() colors = sorted( _records(physics.get("color_components")), key=lambda color: tuple(color.get("word", ())), @@ -969,6 +985,10 @@ def _build_provider( [int(external[slot]["label"]) for slot in word_slots] ) color_index = None if flow is None else _intern(data["color_flows"], flow) + if flow is not None: + coloured_slots = [slot for slot in order if roles[slot] != 1] + blocks = _line_blocks(coloured_slots, roles) if coloured_slots else None + line_legs.update(frozenset(block) for block in blocks or ()) rep_color, perm = replay.get(color_id, (color_id, tuple(identity))) runtime = { "process_index": process_index, @@ -1078,8 +1098,7 @@ def _build_provider( [slot for slot in order if roles[slot] != 1], [slot for slot in order if roles[slot] == 1], roles, - resonances, - source_id, + configurations, ) ) # Contributions share a phase-space channel only with the same kind @@ -1120,8 +1139,39 @@ def _build_provider( ], }, ) - selectors.append((entry, keys)) + # not followed by the channels: the launch rejects those that go on shell + colored_resonances = ( + [] + if flow is None + else [ + { + "pdg": pdg, + "legs": sorted(relabel[slot] for slot in legs), + "color": color, + } + for pdg, legs, color in coloured + ] + ) + selectors.append((entry, keys, colored_resonances)) orbit_entries[orbit] = (entry, aligned_order, runtime, relabel) + unmapped += [ + (pdg, legs) + for pdg, legs in mappable + if legs not in line_legs and any(roles[slot] != 1 for slot in legs) + ] + if unmapped and process.color_accuracy == "lc": + _LOGGER.warning( + "%s: the [amplicol] phase space does not follow the s-channel " + "resonances %s, which decay to coloured and colourless particles " + "or not to one complete colour line", + source_id, + ", ".join( + f"{pdg} > {' '.join(str(pdgs[slot]) for slot in sorted(legs))}" + for pdg, legs in sorted( + unmapped, key=lambda r: (r[0], sorted(r[1])) + ) + ), + ) _group_phasespace_channels(data, selectors, phasespace_channels) _merge_flavour_orbits(data) _export_helicities(data, runtime_helicities) diff --git a/pyamplicol/src/pyamplicol/generation/s_channel_resonances.py b/pyamplicol/src/pyamplicol/generation/s_channel_resonances.py index 4131c5c1d9..7f52d380bf 100644 --- a/pyamplicol/src/pyamplicol/generation/s_channel_resonances.py +++ b/pyamplicol/src/pyamplicol/generation/s_channel_resonances.py @@ -1,17 +1,18 @@ # SPDX-License-Identifier: 0BSD -"""The unstable s-channel propagators a process can put on shell. +"""The unstable s-channel propagators a process may put on shell. A phase-space mapping has to follow these resonances, or it undersamples their peaks. This is a species-level reachability sweep over the subsets of external legs, like the DAG prepass: it ignores helicity and colour, keeps the coupling -orders, and may overestimate. +orders, and may overestimate. Masses and widths are only known at launch, so +every candidate whose mass and width are not zero in the model is kept. """ from __future__ import annotations from collections.abc import Mapping -from ..models.base import CouplingOrders, Model +from ..models.base import CouplingOrders, Model, Particle from ..processes.ir import CanonicalProcessIR from .dag_reachability import ( _combine_coupling_order_tuples, @@ -27,12 +28,12 @@ def s_channel_resonances( model: Model, max_coupling_orders: Mapping[str, int] | None = None, ) -> list[dict[str, object]]: - """``[{"pdg", "labels", "color"}]`` of every unstable particle that a tree - current of two or more final-state legs can carry, with a mass above their - threshold and a complementary current to close the amplitude with. The + """``[{"pdg", "labels", "color"}]`` of every possibly unstable particle (see + ``_possibly_unstable``) that a tree current of two or more final-state legs + can carry, with a complementary current to close the amplitude with. The pdg carries the summed charge of the outgoing legs it decays to.""" - if not any(p.width > 0.0 and p.mass > 0.0 for p in model.particles.values()): + if not any(_possibly_unstable(model, p) for p in model.particles.values()): return [] limits = _normalize_coupling_order_limits(max_coupling_orders) legs = process.legs @@ -75,11 +76,9 @@ def s_channel_resonances( if mask & ~final or mask.bit_count() < 2: continue decay_legs = [leg for leg in legs if mask >> (leg.label - 1) & 1] - threshold = sum(float(model.mass(int(leg.pdg))) for leg in decay_legs) rest = species.get(every & ~mask, {}) for pdg, orders in particles.items(): - particle = model.particle(pdg) - if particle.width <= 0.0 or particle.mass <= threshold: + if not _possibly_unstable(model, model.particle(pdg)): continue closing = rest.get(model.anti_particle(pdg), ()) if not any( @@ -103,3 +102,21 @@ def s_channel_resonances( }, ) return [found[key] for key in sorted(found)] + + +def _possibly_unstable(model: Model, particle: Particle) -> bool: + """Whether neither the mass nor the width of a particle is zero in the + model: a zero value without a parameter (ZERO in UFO) is final.""" + + def zero(kind: str) -> bool: + if getattr(particle, kind) != 0.0: + return False + provider = getattr(model, f"runtime_{kind}_parameter_name", None) + if not callable(provider): + return True + try: + return provider(particle.pdg) is None + except KeyError: + return False + + return not zero("mass") and not zero("width") diff --git a/pyamplicol/tests/unit/test_umami_phasespace_partners.py b/pyamplicol/tests/unit/test_umami_phasespace_partners.py index 4fb27d05cf..db997f6d85 100644 --- a/pyamplicol/tests/unit/test_umami_phasespace_partners.py +++ b/pyamplicol/tests/unit/test_umami_phasespace_partners.py @@ -5,13 +5,20 @@ from __future__ import annotations -import pytest +import logging +import time +from dataclasses import replace +from types import SimpleNamespace from pyamplicol.artifacts.umami import ( + _MAX_PARTNERS, _configurations, _line_blocks, _phasespace_partners, ) +from pyamplicol.generation.s_channel_resonances import s_channel_resonances +from pyamplicol.models.base import Particle +from tests.unit.test_umami_metadata import export, process # all-outgoing colour representations Q, QBAR, G, S = 3, -3, 8, 1 @@ -27,7 +34,7 @@ def test_line_blocks_follow_the_word(): def test_singlets_follow_the_antiquark_in_every_order(): # q q~ > e+ e- g, the word q g q~ crossed: slots 1 (q~ in) ... 0 (q in) roles = [QBAR, Q, S, S, G] - partners = _phasespace_partners([1, 4, 0], [2, 3], roles, [], "q q~ > e+ e- g") + partners = _phasespace_partners([1, 4, 0], [2, 3], roles, [()]) assert [order for order, _ in partners] == [[1, 4, 0, 2, 3], [1, 4, 0, 3, 2]] assert all(configuration == () for _, configuration in partners) @@ -35,13 +42,13 @@ def test_singlets_follow_the_antiquark_in_every_order(): def test_singlets_are_distributed_over_the_lines(): # two lines and one singlet: after the first or the second antiquark roles = [Q, QBAR, Q, QBAR, S] - partners = _phasespace_partners([0, 1, 2, 3], [4], roles, [], "4q + W") + partners = _phasespace_partners([0, 1, 2, 3], [4], roles, [()]) assert [order for order, _ in partners] == [[0, 1, 2, 3, 4], [0, 1, 4, 2, 3]] def test_three_lines_take_both_cyclic_orders(): roles = [Q, QBAR] * 3 - partners = _phasespace_partners(list(range(6)), [], roles, [], "6q") + partners = _phasespace_partners(list(range(6)), [], roles, [()]) assert [order for order, _ in partners] == [ [0, 1, 2, 3, 4, 5], [0, 1, 4, 5, 2, 3], @@ -50,8 +57,9 @@ def test_three_lines_take_both_cyclic_orders(): def test_a_resonance_keeps_its_decay_products_together(): roles = [QBAR, Q, S, S, G] - z = (23, frozenset({2, 3}), 1) - partners = _phasespace_partners([1, 4, 0], [2, 3], roles, [z], "q q~ > e+ e- g") + z = (23, frozenset({2, 3})) + configurations, _ = _configurations([z]) + partners = _phasespace_partners([1, 4, 0], [2, 3], roles, configurations) # the non-resonant placements and the Z, with its legs as one unit assert partners == [ ([1, 4, 0, 2, 3], ()), @@ -62,40 +70,239 @@ def test_a_resonance_keeps_its_decay_products_together(): def test_a_colour_singlet_decaying_to_a_complete_line(): roles = [Q, QBAR, Q, QBAR, G] - z = (23, frozenset({2, 3}), 1) + configurations, _ = _configurations([(23, frozenset({2, 3}))]) # the Z decays to the line 2 3 of this order, but not to legs of two lines - assert len(_phasespace_partners([0, 1, 2, 3], [], roles[:4], [z], "4q")) == 2 - assert len(_phasespace_partners([0, 3, 2, 1], [], roles[:4], [z], "4q")) == 1 + assert len(_phasespace_partners([0, 1, 2, 3], [], roles, configurations)) == 2 + assert len(_phasespace_partners([0, 3, 2, 1], [], roles, configurations)) == 1 def test_overlapping_resonances_are_alternatives(): # e+ e- e+ ve: a Z from either positron, with the W from the other z1, z2 = (23, frozenset({2, 4})), (23, frozenset({3, 4})) w1, w2 = (24, frozenset({2, 5})), (24, frozenset({3, 5})) - configurations = _configurations([z1, z2, w1, w2]) - assert configurations[0] == () - assert sorted(map(set, configurations[1:]), key=str) == sorted( - [{z1, w2}, {z2, w1}], key=str - ) + configurations, complete = _configurations([z1, z2, w1, w2]) + assert complete + assert configurations == [(), (z1, w2), (z2, w1), (z1,), (z2,), (w1,), (w2,)] def test_nested_resonances_stay_contiguous(): roles = [QBAR, Q, S, S, S, S] - h = (25, frozenset({2, 3, 4, 5}), 1) - z1, z2 = (23, frozenset({2, 4}), 1), (23, frozenset({3, 5}), 1) - partners = _phasespace_partners([1, 0], [2, 3, 4, 5], roles, [h, z1, z2], "4l") + h = (25, frozenset({2, 3, 4, 5})) + z1, z2 = (23, frozenset({2, 4})), (23, frozenset({3, 5})) + configurations, _ = _configurations([h, z1, z2]) + partners = _phasespace_partners([1, 0], [2, 3, 4, 5], roles, configurations) resonant = [order for order, configuration in partners if len(configuration) == 3] assert resonant == [[1, 0, 2, 4, 3, 5]] -@pytest.mark.parametrize( - "resonance, message", - [ - ((6, frozenset({2, 3, 4}), 3), "coloured s-channel resonance 6"), - ((25, frozenset({2, 4}), 1), "coloured and colourless"), - ], -) -def test_unsupported_resonances_raise(resonance, message): - roles = [QBAR, Q, Q, S, S] - with pytest.raises(ValueError, match=message): - _phasespace_partners([1, 2, 0], [3, 4], roles, [resonance], "process") +def test_every_compatible_set_of_resonances_is_a_configuration(): + # H > Z Z > 4l: each resonance alone and every nested or disjoint set + h = (25, frozenset({2, 3, 4, 5})) + z1, z2 = (23, frozenset({2, 4})), (23, frozenset({3, 5})) + configurations, complete = _configurations([h, z1, z2]) + assert complete + assert configurations == [ + (), + (h, z1, z2), + (h, z1), + (h, z2), + (z1, z2), + (h,), + (z1,), + (z2,), + ] + + +def test_many_resonances_are_capped_not_rejected(): + # 24 pairwise disjoint resonances have 2^24 configurations + resonances = [(23, frozenset({2 * i, 2 * i + 1})) for i in range(24)] + started = time.perf_counter() + configurations, complete = _configurations(resonances, limit=100) + assert time.perf_counter() - started < 1 + assert not complete + assert len(configurations) == 100 + assert configurations[0] == () + assert [len(c) for c in configurations[1:]] == sorted( + (len(c) for c in configurations[1:]), reverse=True + ) + assert _configurations(resonances, limit=100) == (configurations, complete) + + +def test_partners_are_truncated_deterministically(): + # eight singlets: 8! orders per configuration, enumerated only up to the cap + roles = [QBAR, Q] + [S] * 8 + singlets = list(range(2, 10)) + resonances = [(23, frozenset({2 * i, 2 * i + 1})) for i in range(1, 5)] + configurations, _ = _configurations(resonances) + started = time.perf_counter() + partners = _phasespace_partners([1, 0], singlets, roles, configurations) + assert time.perf_counter() - started < 0.5 + assert len(partners) <= _MAX_PARTNERS + assert partners[0] == ([1, 0, *singlets], ()) + # every configuration keeps its first placements, the largest ones first + assert [c for c in dict.fromkeys(c for _, c in partners)] == configurations + assert partners == _phasespace_partners([1, 0], singlets, roles, configurations) + + +def test_unmapped_singlet_resonances_are_dropped_with_one_warning(caplog): + # q q~ > s s g: an s decaying to an s and the gluon, and a coloured one + p = process(pdgs=(991, -991, 995, 995, 990), words=((2, 5, 1),)) + p.physics["extensions"]["s_channel_resonances"] = [ + {"pdg": 23, "labels": [3, 4], "color": 1}, + {"pdg": 25, "labels": [3, 5], "color": 1}, + {"pdg": 25, "labels": [4, 5], "color": 1}, + {"pdg": 6, "labels": [4, 5], "color": 3}, + ] + with caplog.at_level(logging.WARNING, logger="pyamplicol.artifacts.umami"): + (data,) = export(p) + (warning,) = caplog.records + assert "25 > 995 990, 25 > 995 990" in warning.getMessage() + (channel,) = data["channels"] + assert { + tuple(r["pdg"] for r in ps["resonances"]) + for ps in channel["phasespace_channels"] + } == {(), (23,)} + # a coloured resonance is not followed; the launch checks it + assert [(r["pdg"], r["color"]) for r in channel["colored_resonances"]] == [(6, 3)] + + +def test_singlet_resonance_inside_a_gluon_trace_is_dropped_with_a_warning(caplog): + p = process() + p.physics["extensions"]["s_channel_resonances"] = [ + {"pdg": 25, "labels": [3, 4], "color": 1} + ] + with caplog.at_level(logging.WARNING, logger="pyamplicol.artifacts.umami"): + (data,) = export(p) + (warning,) = caplog.records + assert "25 > 990 990" in warning.getMessage() + (channel,) = data["channels"] + assert all(not ps["resonances"] for ps in channel["phasespace_channels"]) + + +def test_line_resonance_of_some_colour_orders_is_not_warned(caplog): + p = process(pdgs=(991, -991, 992, -992), words=((2, 1, 3, 4), (2, 4, 3, 1))) + p.physics["extensions"]["s_channel_resonances"] = [ + {"pdg": 23, "labels": [3, 4], "color": 1} + ] + sectors = ( + { + "id": 0, + "kind": "open-lines", + "word_labels": [2, 1, 3, 4], + "open_color_lines": [ + { + "fundamental_label": 2, + "adjoint_labels": [], + "antifundamental_label": 1, + }, + { + "fundamental_label": 3, + "adjoint_labels": [], + "antifundamental_label": 4, + }, + ], + }, + { + "id": 1, + "kind": "open-lines", + "word_labels": [2, 4, 3, 1], + "open_color_lines": [ + { + "fundamental_label": 2, + "adjoint_labels": [], + "antifundamental_label": 4, + }, + { + "fundamental_label": 3, + "adjoint_labels": [], + "antifundamental_label": 1, + }, + ], + }, + ) + with caplog.at_level(logging.WARNING, logger="pyamplicol.artifacts.umami"): + (data,) = export(replace(p, color_sectors=sectors)) + assert not caplog.records + resonant = [ + ps + for c in data["channels"] + for ps in c["phasespace_channels"] + if ps["resonances"] + ] + assert [r["pdg"] for ps in resonant for r in ps["resonances"]] == [23] + + +class _Model: + """f f~ > l- l+ through any vector of ``particles``.""" + + def __init__(self, particles, parameters): + self.particles = {p.pdg: p for p in particles} + self.parameters = parameters + + def vertices_accepting(self, left, right, color_accuracy): + if {left, right} != {11, -11}: + return [] + return [ + SimpleNamespace(particles=(left, right, pdg)) + for pdg in self.particles + if abs(pdg) != 11 + ] + + def skip_duplicate_vertex_orientation(self, vertex): + return False + + def vertex_coupling_orders(self, vertex): + return (("QED", 1),) + + def particle(self, pdg): + return self.particles[pdg] + + def anti_particle(self, pdg): + return self.particles[pdg].anti_pdg + + def charge(self, pdg): + return self.particles[pdg].charge + + def color_rep(self, pdg): + return self.particles[pdg].color_rep + + def runtime_mass_parameter_name(self, pdg): + return self.parameters.get((pdg, "mass")) + + def runtime_width_parameter_name(self, pdg): + return self.parameters.get((pdg, "width")) + + +def test_candidates_are_kept_unless_their_mass_or_width_is_zero_in_the_model(): + def vector(pdg, mass, width): + return Particle(pdg, pdg, 2, 4, 1, mass, width) + + particles = [ + Particle(11, -11, 1, 4, 1, 50.0, 0.0, -1.0), + Particle(-11, 11, 1, 4, 1, 50.0, 0.0, 1.0), + vector(22, 0.0, 0.0), # massless + vector(23, 91.0, 2.5), + vector(97, 60.0, 1.0), # below the threshold of the default masses + vector(98, 200.0, 0.0), # no width in the model + vector(99, 200.0, 0.0), # zero width on the parameter card + ] + parameters = {(11, "mass"): "ML", (-11, "mass"): "ML", (99, "width"): "WX"} + parameters |= {(pdg, kind): kind for pdg in (23, 97, 98, 99) for kind in ("mass",)} + parameters |= {(pdg, "width"): "W" for pdg in (23, 97)} + + def leg(label, pdg, outgoing, side): + return SimpleNamespace( + label=label, pdg=pdg, outgoing_pdg=outgoing, is_final=side == "final" + ) + + ir = SimpleNamespace( + color_accuracy="lc", + legs=( + leg(1, 11, -11, "initial"), + leg(2, -11, 11, "initial"), + leg(3, 11, 11, "final"), + leg(4, -11, -11, "final"), + ), + ) + found = s_channel_resonances(ir, _Model(particles, parameters)) + assert found == [{"pdg": pdg, "labels": [3, 4], "color": 1} for pdg in (23, 97, 99)] From ba4ee238e8f99d0ce14fbb6b835def2b13d25127 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:04:22 +0100 Subject: [PATCH 69/88] madspace: give non-finite weighting step ratios the weight 0 A non-finite ratio, or one with an earlier matrix element of 0, could become the maximum of a step and then reject every event, or abort a reweighting step. Such events now get the weight 0, are counted per channel in GeneratorStatus.count_invalid_step_values and a warning is logged after the survey and generate(). --- .../madspace/driver/channel_generator.hpp | 5 +++ .../madspace/driver/event_generator.hpp | 3 ++ .../madspace/driver/generator_data.hpp | 3 ++ madspace/src/driver/channel_generator.cpp | 3 ++ madspace/src/driver/event_generator.cpp | 21 ++++++++++ madspace/src/driver/event_utils.cpp | 3 ++ madspace/src/driver/generator_data.cpp | 1 + madspace/src/driver/weighting_pipeline.cpp | 25 +++++++++--- madspace/src/python/madspace.cpp | 5 +++ madspace/tests/test_weighting_steps.py | 40 +++++++++++++++++++ 10 files changed, 104 insertions(+), 5 deletions(-) diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index a958713702..c76334638b 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -229,6 +229,11 @@ class ChannelEventGenerator { void update_step_max_weight( std::size_t step, const std::vector& values, std::size_t count_target ); + /// Count `count` events given the weight 0 by a step, see + /// @ref GeneratorStatus::count_invalid_step_values. + void add_invalid_step_values(std::size_t count) { + _status.count_invalid_step_values += count; + } /// Count `count` events accepted by step `step`, which passed the earlier /// unweightings with the maxima `max_weight` (first step) and /// `max_weights` (the steps before). diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index eaaf1d511a..ebed03d570 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -159,6 +159,8 @@ class EventGenerator { std::vector _channel_optimizing; std::vector _channel_integral_fractions; std::vector _context_job_counts; + // invalid step values per channel that were already warned about + std::vector _warned_invalid_step_values; // True while a channel has a steady-state batch dispatched but not yet fully // committed; keeps next_batch_event_count() from double-counting in-flight work. std::vector _channel_batch_pending; @@ -302,6 +304,7 @@ class EventGenerator { std::size_t iter ); + void warn_invalid_step_values(); void print_gen_init(); void print_gen_update(bool done); void print_gen_update_pretty(bool done); diff --git a/madspace/include/madspace/driver/generator_data.hpp b/madspace/include/madspace/driver/generator_data.hpp index 447a5aec5f..02e4c52e82 100644 --- a/madspace/include/madspace/driver/generator_data.hpp +++ b/madspace/include/madspace/driver/generator_data.hpp @@ -251,6 +251,9 @@ struct GeneratorStatus { bool optimized; /// Whether generation of this subprocess is complete. bool done; + /// Events given the weight 0 by a weighting step, because their matrix + /// element ratio was not finite or the earlier matrix element was 0. + std::size_t count_invalid_step_values = 0; }; /// One binned observable, accumulated during event generation; see @ref diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 26902eb102..4d2e8b9faa 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -1237,6 +1237,9 @@ void ChannelEventGenerator::update_step_max_weight( double min_nonzero = 0.; for (double value : values) { double abs_value = std::abs(value); + if (!std::isfinite(abs_value)) { + continue; + } if (abs_value != 0 && (min_nonzero == 0 || abs_value < min_nonzero)) { min_nonzero = abs_value; } diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index e81c56e090..1733eeced4 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -50,6 +50,7 @@ EventGenerator::EventGenerator( _channel_optimizing(channels.size()), _channel_integral_fractions(channels.size(), 1.), _context_job_counts(contexts.size()), + _warned_invalid_step_values(channels.size()), _channel_batch_pending(channels.size()), _channel_batch_dispatch_done(channels.size()), _channel_gen_order(channels.size()), @@ -258,6 +259,7 @@ void EventGenerator::run_generate(GenerateMode mode) { } } print_gen_update(true); + warn_invalid_step_values(); } // Same commit-ordering approach as generate(). @@ -403,6 +405,25 @@ void EventGenerator::survey(std::size_t survey_pass) { drain_pipeline(); } print_survey_update(true, done_event_count, total_event_count, iter - 1); + warn_invalid_step_values(); +} + +void EventGenerator::warn_invalid_step_values() { + for (auto [channel, warned] : zip(_channels, _warned_invalid_step_values)) { + std::size_t count = channel->status().count_invalid_step_values; + if (count > warned) { + Logger::warning( + std::format( + "channel '{}': {} events got the weight 0 in the weighting steps, " + "as their matrix element ratio was not finite or the earlier " + "matrix element was 0", + channel->status().name, + count - warned + ) + ); + warned = count; + } + } } // Every job counted in *in_flight* posts exactly one result (see diff --git a/madspace/src/driver/event_utils.cpp b/madspace/src/driver/event_utils.cpp index bc358e7e6f..af661685a1 100644 --- a/madspace/src/driver/event_utils.cpp +++ b/madspace/src/driver/event_utils.cpp @@ -132,6 +132,7 @@ void madspace::sum_channel_status( std::size_t total_count_after_cuts = 0, total_count_after_cuts_opt = 0; std::size_t total_integ_count = 0; std::size_t iterations = 0; + std::size_t count_invalid_step_values = 0; bool optimized = true; const ChannelEventGenerator* bad_channel = nullptr; for (auto& channel : channels) { @@ -166,6 +167,7 @@ void madspace::sum_channel_status( total_count_after_cuts_opt += chan_status.count_after_cuts_opt; total_integ_count += integral.count; iterations = std::max(chan_status.iterations, iterations); + count_invalid_step_values += chan_status.count_invalid_step_values; if (!chan_status.optimized) { optimized = false; } @@ -206,5 +208,6 @@ void madspace::sum_channel_status( status.count_after_cuts = total_count_after_cuts; status.count_after_cuts_opt = total_count_after_cuts_opt; status.iterations = iterations; + status.count_invalid_step_values = count_invalid_step_values; status.optimized = optimized; } diff --git a/madspace/src/driver/generator_data.cpp b/madspace/src/driver/generator_data.cpp index 1b3679ed61..b2d3696767 100644 --- a/madspace/src/driver/generator_data.cpp +++ b/madspace/src/driver/generator_data.cpp @@ -76,6 +76,7 @@ void madspace::to_json(nlohmann::json& j, const GeneratorStatus& status) { {"iters_without_improvement", status.iters_without_improvement}, {"optimized", status.optimized}, {"done", status.done}, + {"count_invalid_step_values", status.count_invalid_step_values}, }; } diff --git a/madspace/src/driver/weighting_pipeline.cpp b/madspace/src/driver/weighting_pipeline.cpp index 6cda38adb4..c3b8b59fe9 100644 --- a/madspace/src/driver/weighting_pipeline.cpp +++ b/madspace/src/driver/weighting_pipeline.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include "madspace/util.hpp" @@ -463,17 +464,31 @@ void WeightingPipeline::decide(Chunk& chunk, std::size_t step) { auto& settings = channel->weighting_steps().at(step - 1); std::size_t count = chunk.size(); std::vector ratios(count); + std::size_t invalid_count = 0; for (std::size_t i = 0; i < count; ++i) { double previous = chunk.matrix_elements.at(i); double& value = chunk.step_matrix_elements.at(i); + double ratio = value; + if (!settings.ratio) { + ratio = previous != 0. ? value / previous + : value != 0. + ? std::numeric_limits::quiet_NaN() + : 0.; + } + // the event gets the weight 0, so that it cannot set the maximum + if (!std::isfinite(ratio)) { + ratio = 0.; + ++invalid_count; + } if (settings.ratio) { // the implied matrix element is the reference of the next step - ratios.at(i) = value; - value *= previous; - } else { - ratios.at(i) = previous != 0. ? value / previous : 0.; + value = ratio * previous; } - chunk.ratio_products.at(i) *= ratios.at(i); + ratios.at(i) = ratio; + chunk.ratio_products.at(i) *= ratio; + } + if (invalid_count > 0) { + channel->add_invalid_step_values(invalid_count); } std::vector keep; if (settings.unweight) { diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 580e4e9c4e..33a1d2b538 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -3002,6 +3002,11 @@ PYBIND11_MODULE(_madspace_py, m) { ) .def_readwrite( "done", &GeneratorStatus::done, pydoc::doc("GeneratorStatus::done") + ) + .def_readwrite( + "count_invalid_step_values", + &GeneratorStatus::count_invalid_step_values, + pydoc::doc("GeneratorStatus::count_invalid_step_values") ); py::classh(m, "Histogram", pydoc::doc("Histogram")) diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py index 84477eeb9d..79594f708b 100644 --- a/madspace/tests/test_weighting_steps.py +++ b/madspace/tests/test_weighting_steps.py @@ -240,6 +240,46 @@ def test_reproducible_step_freeze(tmp_path): assert status_a.mean == status_b.mean +def broken_me(momenta, invalid=(np.nan, np.inf)): + """true_me, with invalid values in two regions.""" + px = momenta[:, 3, 1] + me = np.where(px > 100.0, invalid[0], true_me(momenta)) + return np.where(px < -100.0, invalid[1], me) + + +@pytest.mark.parametrize("step_modes", [["unweight"], ["reweight"]]) +def test_invalid_step_values(tmp_path, capfd, step_modes): + """Non-finite ratios count as weight 0, like a step matrix element of 0, + and are counted and warned about.""" + results = [ + generate( + str(tmp_path / name), + surrogate_me, + step_modes, + step_me=step_me, + pt_mins=(20.0, 40.0), + grouped=(0, 1), + ) + for name, step_me in [ + ("broken", broken_me), + ("zero", lambda momenta: broken_me(momenta, (0.0, 0.0))), + ] + ] + (status_broken, events_broken, channels), (status_zero, events_zero, _) = results + assert np.array_equal(events_broken, events_zero) + assert status_broken.mean == status_zero.mean + assert status_zero.count_invalid_step_values == 0 + invalid_counts = [c.status().count_invalid_step_values for c in channels] + assert all(count > 0 for count in invalid_counts) + assert status_broken.count_invalid_step_values == sum(invalid_counts) + if step_modes == ["unweight"]: + assert all(np.isfinite(c.step_max_weight(1)) for c in channels) + # after the survey and after generate() + warnings = [line for line in capfd.readouterr().out.splitlines() if "WARN" in line] + assert all("got the weight 0 in the weighting steps" in w for w in warnings) + assert {w.split("'")[1] for w in warnings} == {"0", "1"} + + def test_index_outputs(tmp_path): """A step fills in indices the integrand matrix element does not return, and replaces the others only with override_outputs.""" From 0ce842d542e4fe9969ccc7a6b6ced000139993fd Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:05:32 +0100 Subject: [PATCH 70/88] pyamplicol: inspect artifacts whose payloads are packed in payloads.pacbin --- .../src/pyamplicol/artifacts/inspection.py | 55 ++++++++++----- .../tests/unit/test_artifact_payloads.py | 67 +++++++++++++++++++ ...est_recurrence_direct_artifact_metadata.py | 8 +-- 3 files changed, 111 insertions(+), 19 deletions(-) diff --git a/pyamplicol/src/pyamplicol/artifacts/inspection.py b/pyamplicol/src/pyamplicol/artifacts/inspection.py index 242348fb64..cc5a5e1a3c 100644 --- a/pyamplicol/src/pyamplicol/artifacts/inspection.py +++ b/pyamplicol/src/pyamplicol/artifacts/inspection.py @@ -30,7 +30,7 @@ from pyamplicol.api.errors import ArtifactError from pyamplicol.generation.process_grouping import GroupMember, ProcessGroupMap -from .manifest import MANIFEST_NAME, ArtifactManifest, load_manifest +from .manifest import MANIFEST_NAME, ArtifactManifest, ArtifactPayloads, load_manifest from .process_groups import initial_process_groups, read_process_groups from .security import confined_path, normalize_relative_path @@ -321,12 +321,30 @@ def _string(value: object, context: str) -> str: return value -def _json_mapping(path: Path, context: str) -> Mapping[str, object]: +# the payload reader of the last inspected manifest, which it keeps alive +_PAYLOAD_READER: list[tuple[ArtifactManifest, ArtifactPayloads]] = [] + + +def _payloads(manifest: ArtifactManifest) -> ArtifactPayloads: + if not _PAYLOAD_READER or _PAYLOAD_READER[0][0] is not manifest: + _PAYLOAD_READER[:] = [(manifest, ArtifactPayloads(manifest))] + return _PAYLOAD_READER[0][1] + + +def _read_payload(manifest: ArtifactManifest, path: Path) -> bytes: + """Read a payload at ``path`` below the root, loose or packed.""" + + return _payloads(manifest).read_bytes(path.relative_to(manifest.root).as_posix()) + + +def _json_mapping( + manifest: ArtifactManifest, path: Path, context: str +) -> Mapping[str, object]: try: - value = json.loads(path.read_text(encoding="utf-8")) + value = json.loads(_read_payload(manifest, path)) except OSError as exc: raise ArtifactError(f"cannot read {context}: {exc}") from exc - except json.JSONDecodeError as exc: + except (UnicodeDecodeError, json.JSONDecodeError) as exc: raise ArtifactError(f"invalid {context}: {exc}") from exc return _mapping(value, context) @@ -342,6 +360,8 @@ def _artifact_path( raise except (TypeError, ValueError) as exc: raise ArtifactError(f"{context} is not a valid artifact path") from exc + if _payloads(manifest).is_packed(normalized): + return manifest.root / normalized return confined_path(manifest.root, normalized) @@ -355,7 +375,9 @@ def _execution_paths(manifest: ArtifactManifest) -> Mapping[str, Path]: evaluator_relative, "runtime.evaluator_manifest_path", ) - evaluator_set = _json_mapping(evaluator_path, "runtime evaluator manifest") + evaluator_set = _json_mapping( + manifest, evaluator_path, "runtime evaluator manifest" + ) records = _sequence( evaluator_set.get("processes"), "runtime evaluator manifest.processes", @@ -504,7 +526,7 @@ def _table_rows( f"{context}.path", ) try: - payload = path.read_bytes() + payload = _read_payload(manifest, path) except OSError as exc: raise ArtifactError(f"cannot read {context} payload: {exc}") from exc expected_size = count * row_size @@ -905,7 +927,7 @@ def _eager_execution_inspection( ), "eager execution.kernel_pack.manifest_path", ) - pack = _json_mapping(pack_path, "prepared eager kernel pack") + pack = _json_mapping(manifest, pack_path, "prepared eager kernel pack") kernels = _sequence(pack.get("kernels"), "prepared eager kernel pack.kernels") kernel_ids = { _integer( @@ -1000,7 +1022,7 @@ def _eager_v3_execution_inspection( ), "eager execution.kernel_pack.manifest_path", ) - pack = _json_mapping(pack_path, "prepared eager kernel pack") + pack = _json_mapping(manifest, pack_path, "prepared eager kernel pack") kernels = _sequence(pack.get("kernels"), "prepared eager kernel pack.kernels") invocation_count = _integer( summary.get("invocation_count"), @@ -1202,7 +1224,7 @@ def _recurrence_execution_inspection( ), "recurrence execution.kernel_pack.manifest_path", ) - pack = _json_mapping(pack_path, "prepared recurrence kernel pack") + pack = _json_mapping(manifest, pack_path, "prepared recurrence kernel pack") kernels = _sequence( pack.get("kernels"), "prepared recurrence kernel pack.kernels", @@ -1949,7 +1971,9 @@ def _execution_inspection( selected_color_capability == ON_THE_FLY_CONTRACTED_COLOR_RUNTIME_CAPABILITY ) expected_color_accuracies = {"nlc", "full"} if contracted_color else {"lc"} - execution = _json_mapping(execution_path, "process execution manifest") + execution = _json_mapping( + manifest, execution_path, "process execution manifest" + ) expected_root_fields = { "schema_version", "kind", @@ -2268,7 +2292,9 @@ def _execution_inspection( "recurrence direct-arena execution is missing required " f"capabilities: {missing}" ) - execution = _json_mapping(execution_path, "process execution manifest") + execution = _json_mapping( + manifest, execution_path, "process execution manifest" + ) kind = _string(execution.get("kind"), "process execution manifest.kind") if kind != _RECURRENCE_RUNTIME_KIND: raise ArtifactError( @@ -2284,7 +2310,7 @@ def _execution_inspection( execution_mode="compiled", native_profile_phases=_COMPILED_PROFILE_PHASES, ) - execution = _json_mapping(execution_path, "process execution manifest") + execution = _json_mapping(manifest, execution_path, "process execution manifest") kind = _string(execution.get("kind"), "process execution manifest.kind") if kind == _EAGER_RUNTIME_KIND: return _eager_execution_inspection(manifest, execution, execution_path) @@ -2365,14 +2391,13 @@ def _physics_inspection( execution: _ExecutionInspection, ) -> _PhysicsInspection: relative = str(process["physics_path"]) - path = manifest.root / relative try: - payload = json.loads(path.read_text(encoding="utf-8")) + payload = json.loads(_payloads(manifest).read_bytes(relative)) except OSError as exc: raise ArtifactError( f"cannot read runtime physics metadata {relative}: {exc}" ) from exc - except json.JSONDecodeError as exc: + except (UnicodeDecodeError, json.JSONDecodeError) as exc: raise ArtifactError( f"invalid runtime physics metadata {relative}: {exc}" ) from exc diff --git a/pyamplicol/tests/unit/test_artifact_payloads.py b/pyamplicol/tests/unit/test_artifact_payloads.py index eb5618abaa..0ef010ae56 100644 --- a/pyamplicol/tests/unit/test_artifact_payloads.py +++ b/pyamplicol/tests/unit/test_artifact_payloads.py @@ -2,6 +2,7 @@ from __future__ import annotations import json +import shutil from pathlib import Path import pytest @@ -76,3 +77,69 @@ def test_packed_payloads_may_not_also_be_files(tmp_path: Path) -> None: with pytest.raises(ArtifactError, match="also stored as a file"): load_manifest(root, verify_payloads=True) + + +def _pack_artifact(source: Path, root: Path) -> None: + """Copy ``source`` and move its packable payloads into payloads.pacbin.""" + + from pyamplicol.artifacts.manifest import compute_artifact_id + from pyamplicol.artifacts.security import sha256_file + from pyamplicol.generation.artifact_writer import _packs_payload + from pyamplicol.generation.evaluator_container import ( + PacbinMemberKind, + PacbinMemberSource, + write_pacbin, + ) + + shutil.copytree(source, root) + raw = json.loads((root / "artifact.json").read_text(encoding="utf-8")) + packed = [ + record["path"] for record in raw["payloads"] if _packs_payload(record["path"]) + ] + assert packed + container = root / "payloads.pacbin" + write_pacbin( + container, + [ + PacbinMemberSource(path, PacbinMemberKind.DECLARED_PAYLOAD, root / path) + for path in packed + ], + ) + for path in packed: + (root / path).unlink() + raw["payloads"] = sorted( + [ + *raw["payloads"], + { + "executable": False, + "media_type": "application/octet-stream", + "path": "payloads.pacbin", + "role": "payload-container", + "sha256": sha256_file(container), + "size_bytes": container.stat().st_size, + }, + ], + key=lambda record: record["path"], + ) + raw["artifact_id"] = compute_artifact_id(raw) + (root / "artifact.json").write_text(json.dumps(raw), encoding="utf-8") + + +def test_inspection_reads_packed_payloads(tmp_path: Path) -> None: + from pyamplicol.artifacts import inspect_artifact + + source = ( + Path(__file__).resolve().parents[2] + / "src/pyamplicol/assets/selftest/portable-64le/artifact" + ) + root = tmp_path / "packed" + _pack_artifact(source, root) + assert ArtifactPayloads(load_manifest(root)).is_packed( + "processes/d_dbar_to_z/physics.json" + ) + assert not (root / "processes/d_dbar_to_z/physics.json").exists() + + loose = inspect_artifact(source) + packed = inspect_artifact(root) + assert packed.integrity == "verified" + assert packed.processes == loose.processes diff --git a/pyamplicol/tests/unit/test_recurrence_direct_artifact_metadata.py b/pyamplicol/tests/unit/test_recurrence_direct_artifact_metadata.py index f6c0f377f3..ac845c70b9 100644 --- a/pyamplicol/tests/unit/test_recurrence_direct_artifact_metadata.py +++ b/pyamplicol/tests/unit/test_recurrence_direct_artifact_metadata.py @@ -1084,7 +1084,7 @@ def test_direct_recurrence_inspection_exposes_arena_counters(tmp_path: Path) -> _write_prepared_pack(tmp_path) execution = _execution_manifest() result = inspection._recurrence_execution_inspection( - SimpleNamespace(root=tmp_path), + SimpleNamespace(root=tmp_path, payloads=()), execution, ) @@ -1133,7 +1133,7 @@ def test_direct_recurrence_inspection_exposes_contracted_color_summary( } result = inspection._recurrence_execution_inspection( - SimpleNamespace(root=tmp_path), + SimpleNamespace(root=tmp_path, payloads=()), execution, ) @@ -1160,7 +1160,7 @@ def test_inspection_requires_v2_metadata_without_packet_fallback( with pytest.raises(ArtifactError, match="regenerate the recurrence artifact"): inspection._recurrence_execution_inspection( - SimpleNamespace(root=tmp_path), + SimpleNamespace(root=tmp_path, payloads=()), execution, ) @@ -1174,7 +1174,7 @@ def test_v2_inspection_tolerates_pending_direct_counter_wiring( execution = _execution_manifest(process) result = inspection._recurrence_execution_inspection( - SimpleNamespace(root=tmp_path), + SimpleNamespace(root=tmp_path, payloads=()), execution, ) From 57920b2d2156bb512f74087bd7f8c84c84a4148a Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:05:32 +0100 Subject: [PATCH 71/88] pyamplicol: take the payload container constants from the manifest module --- pyamplicol/src/pyamplicol/artifacts/writer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyamplicol/src/pyamplicol/artifacts/writer.py b/pyamplicol/src/pyamplicol/artifacts/writer.py index ce2b22ddeb..26aa455d5a 100644 --- a/pyamplicol/src/pyamplicol/artifacts/writer.py +++ b/pyamplicol/src/pyamplicol/artifacts/writer.py @@ -17,6 +17,8 @@ from .manifest import ( ARTIFACT_IDENTITY_EXTENSION, MANIFEST_NAME, + PAYLOAD_CONTAINER_PATH, + PAYLOAD_CONTAINER_ROLE, PayloadRecord, artifact_identity_extension, canonical_manifest_bytes, @@ -32,8 +34,6 @@ ) from .transaction import ArtifactTransaction, ArtifactWriteMode -PAYLOAD_CONTAINER_PATH = "payloads.pacbin" -PAYLOAD_CONTAINER_ROLE = "payload-container" # packable payloads are appended here while staging, then moved into the container _PAYLOAD_SPOOL_NAME = ".payloads.spool" From c9adec2a3b986fa282b366fb55b2b670fadf2e77 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:06:16 +0100 Subject: [PATCH 72/88] mg7: [amplicol] resonances that go on shell are selected at launch The phase-space channels of a subprocess follow the resonances found at output that go on shell with the param card: a positive width and a mass above the masses of the decay legs. The others are dropped, the channels that then agree merged and the flavours' channels and LHE resonances renumbered. A coloured resonance is an error only when it goes on shell. --- madgraph/iolibs/export_amplicol.py | 1 + madgraph/mg7/subprocess_amplicol.py | 57 +++++++++++++- .../interface/test_amplicol_syntax.py | 75 +++++++++++++++++++ 3 files changed, 130 insertions(+), 3 deletions(-) diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index fbbabff300..1e091e3877 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -403,6 +403,7 @@ def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power, 'phasespace_channels': [ {'order': ps['phasespace_order'], 'resonances': ps['resonances']} for ps in phasespace], + 'colored_resonances': channel.get('colored_resonances', []), 'masses': provider['masses'], 'supports_alpha_s': provider['supports_alpha_s'], 'flavors': flavors, diff --git a/madgraph/mg7/subprocess_amplicol.py b/madgraph/mg7/subprocess_amplicol.py index cbc2277ad9..cb93718e8d 100644 --- a/madgraph/mg7/subprocess_amplicol.py +++ b/madgraph/mg7/subprocess_amplicol.py @@ -4,9 +4,11 @@ integrated at leading colour over a common set of colour-ordered phase-space channels. They differ in where the colour singlets attach to the colour lines and in the s-channel resonances they map out, and are weighted by the densities -of their mappings. A second pass re- or unweights the events to full colour with -the ratio of the full-colour to the leading-colour matrix element summed over -all colour orderings, at the helicity of the event. +of their mappings. Resonances that cannot go on shell with the param card +masses and widths are dropped at launch. A second pass re- or +unweights the events to full colour with the ratio of the full-colour to the +leading-colour matrix element summed over all colour orderings, at the +helicity of the event. """ from __future__ import annotations @@ -135,6 +137,53 @@ def channel_topologies(meta: dict, process: MadgraphProcess) -> list[tuple]: ] +def active_resonances(meta: dict, get_mass, get_width) -> dict: + """The subprocess meta without the s-channel resonances that cannot go on + shell: a resonance is active if its width is positive and its mass above + the masses of its decay legs. Phase-space channels that then agree are + merged, and the flavours renumbered to them.""" + leg_masses = [ + get_mass(pid) for pid in clean_pids(meta["incoming"] + meta["outgoing"]) + ] + + def active(resonance): + pdg = abs(resonance["pdg"]) + return get_width(pdg) > 0 and get_mass(pdg) > sum( + leg_masses[leg] for leg in resonance["legs"] + ) + + for resonance in meta.get("colored_resonances", []): + if active(resonance): + raise RuntimeError( + f"the coloured s-channel resonance {resonance['pdg']} can go on " + "shell, which [amplicol] phase-space mappings do not follow yet; " + "generate the process with on-shell decays or with the madmatrix " + "backend" + ) + channels, lookup, remap, kept = [], {}, [], [] + for channel in meta["phasespace_channels"]: + keep = [i for i, r in enumerate(channel["resonances"]) if active(r)] + resonances = [channel["resonances"][i] for i in keep] + key = ( + tuple(channel["order"]), + tuple((abs(r["pdg"]), tuple(r["legs"])) for r in resonances), + ) + if key not in lookup: + lookup[key] = len(channels) + channels.append({**channel, "resonances": resonances}) + remap.append(lookup[key]) + kept.append(keep) + flavors = [] + for flavor in meta["flavors"]: + indices, pdgs = [], [] + for index, resonance_pdgs in zip(flavor["channels"], flavor["resonance_pdgs"]): + if remap[index] not in indices: + indices.append(remap[index]) + pdgs.append([resonance_pdgs[i] for i in kept[index]]) + flavors.append({**flavor, "channels": indices, "resonance_pdgs": pdgs}) + return {**meta, "phasespace_channels": channels, "flavors": flavors} + + def lhe_diagrams(meta: dict) -> tuple[list[tuple[int, tuple]], list[list[float]]]: """The diagrams that the LHE writer takes the resonances of an event from, one per phase-space channel and resonance charges, and per flavour a mask @@ -159,6 +208,7 @@ def lhe_diagrams(meta: dict) -> tuple[list[tuple[int, tuple]], list[list[float]] def amplicol_subproc_args(meta: dict, process: MadgraphProcess) -> ms.SubprocArgs: """LHE completion of an AmpliCol subprocess: its diagrams are the phase-space channels, with the resonances they map out.""" + meta = active_resonances(meta, process.get_mass, process.get_width) diagrams, _ = lhe_diagrams(meta) if diagrams: topologies = channel_topologies(meta, process) @@ -207,6 +257,7 @@ def amplicol_subproc_args(meta: dict, process: MadgraphProcess) -> ms.SubprocArg class AmplicolSubprocess(MadgraphSubprocess): def __init__(self, process, meta, subproc_id, unmerged_meta=None): super().__init__(process, meta, subproc_id, unmerged_meta) + self.meta = active_resonances(self.meta, process.get_mass, process.get_width) self.check_flavor_blind_cuts() def check_flavor_blind_cuts(self): diff --git a/tests/unit_tests/interface/test_amplicol_syntax.py b/tests/unit_tests/interface/test_amplicol_syntax.py index 84804ddf5d..9ee4f4945c 100644 --- a/tests/unit_tests/interface/test_amplicol_syntax.py +++ b/tests/unit_tests/interface/test_amplicol_syntax.py @@ -232,6 +232,7 @@ def test_resonance_charges_split_the_selector_rows(self): {'phasespace_order': [0, 4, 1, 2, 3], 'resonances': []}, {'phasespace_order': [0, 4, 1, 2, 3], 'resonances': [{'pdg': 24, 'legs': [2, 3]}]}], + 'colored_resonances': [{'pdg': 6, 'legs': [2, 3, 4], 'color': 3}], 'processes': [ {'multichannels': [0, 1], 'color_flows': 0, 'matrix_elements': [{'pdg_ids': index, 'factor': 1} @@ -252,3 +253,77 @@ def test_resonance_charges_split_the_selector_rows(self): self.assertEqual(entries[0]['phasespace_channels'][1], {'order': [0, 4, 1, 2, 3], 'resonances': [{'pdg': 24, 'legs': [2, 3]}]}) + # the launch rejects a coloured resonance that goes on shell + self.assertEqual(entries[0]['colored_resonances'], + [{'pdg': 6, 'legs': [2, 3, 4], 'color': 3}]) + + +class AmplicolLaunchResonancesTest(unittest.TestCase): + """The s-channel resonances of an [amplicol] subprocess are fixed at output; + the launch keeps those that go on shell with the param card values.""" + + # q q~ > e- e+ mu- mu+ through Z Z, Z and H > Z Z + order = [0, 2, 3, 4, 5, 1] + z_e, z_mu = {'pdg': 23, 'legs': [2, 3]}, {'pdg': 23, 'legs': [4, 5]} + h = {'pdg': 25, 'legs': [2, 3, 4, 5]} + + def setUp(self): + try: + from madgraph.mg7 import subprocess_amplicol + except ImportError: + self.skipTest('madspace is not installed') + self.module = subprocess_amplicol + self.meta = { + 'incoming': [2, -2], 'outgoing': [11, -11, 13, -13], + 'phasespace_channels': [ + {'order': self.order, 'resonances': []}, + {'order': self.order, 'resonances': [self.z_e, self.z_mu]}, + {'order': self.order, 'resonances': [self.h, self.z_e]}, + {'order': self.order, 'resonances': [self.h]}, + ], + 'colored_resonances': [{'pdg': 6, 'legs': [2, 3], 'color': 3}], + 'flavors': [ + {'channels': [0, 1, 2, 3], + 'resonance_pdgs': [[], [23, 23], [25, 23], [25]]}, + {'channels': [3, 1], 'resonance_pdgs': [[25], [23, 23]]}, + ], + } + + def active(self, masses, widths): + return self.module.active_resonances( + self.meta, lambda pid: masses.get(pid, 0.), lambda pid: widths.get(pid, 0.)) + + def test_all_resonances_on_shell(self): + meta = self.active({23: 91., 25: 125.}, {23: 2.5, 25: 4e-3}) + self.assertEqual(meta['phasespace_channels'], self.meta['phasespace_channels']) + self.assertEqual(meta['flavors'], self.meta['flavors']) + + def test_stable_resonance_is_dropped_and_channels_merged(self): + # without a Higgs width, H > Z Z becomes Z and H the non-resonant channel + meta = self.active({23: 91., 25: 125.}, {23: 2.5}) + self.assertEqual(meta['phasespace_channels'], [ + {'order': self.order, 'resonances': []}, + {'order': self.order, 'resonances': [self.z_e, self.z_mu]}, + {'order': self.order, 'resonances': [self.z_e]}, + ]) + self.assertEqual([(f['channels'], f['resonance_pdgs']) for f in meta['flavors']], + [([0, 1, 2], [[], [23, 23], [23]]), ([0, 1], [[], [23, 23]])]) + # the LHE diagram masks follow the merged channels + diagrams, masks = self.module.lhe_diagrams(meta) + self.assertEqual(diagrams, [(0, ()), (1, (23, 23)), (2, (23,))]) + self.assertEqual(masks, [[1., 1., 1.], [1., 1., 0.]]) + + def test_resonance_below_its_decay_threshold_is_dropped(self): + # mu+ mu- of 2 x 70 GeV: neither Z > mu+ mu- nor H > 4l are on shell + meta = self.active({13: 70., 23: 91., 25: 125.}, {23: 2.5, 25: 4e-3}) + self.assertEqual( + [c['resonances'] for c in meta['phasespace_channels']], [[], [self.z_e]]) + self.assertEqual([f['channels'] for f in meta['flavors']], [[0, 1], [0, 1]]) + self.assertEqual(self.meta['flavors'][1]['channels'], [3, 1]) + + def test_coloured_resonance_on_shell_is_an_error(self): + masses, widths = {6: 173., 23: 91.}, {23: 2.5} + self.active(masses, widths) + widths[6] = 1.5 + with self.assertRaisesRegex(RuntimeError, 'coloured s-channel resonance 6'): + self.active(masses, widths) From 4a2909603bf39e6e707d4f39989b9efe79630955 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:15:25 +0100 Subject: [PATCH 73/88] madspace: apply the truncation budget of a smaller target to the step maxima set_target_count() only re-applied the budget to the first maximum, so a step maximum chosen under generate()'s generous round-one target stayed too low for a channel with a small target. EventGenerator now re-applies it to the steps at the end of each round, with the steps drained so that it does not depend on the thread timing. --- .../madspace/driver/channel_generator.hpp | 4 ++ madspace/src/driver/channel_generator.cpp | 11 +++++ madspace/src/driver/event_generator.cpp | 4 ++ madspace/src/python/madspace.cpp | 11 +++++ madspace/tests/test_weighting_steps.py | 47 +++++++++++++++++-- 5 files changed, 73 insertions(+), 4 deletions(-) diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index c76334638b..547ce793bf 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -234,6 +234,10 @@ class ChannelEventGenerator { void add_invalid_step_values(std::size_t count) { _status.count_invalid_step_values += count; } + /// Apply the truncation budget of the current target count to the maxima + /// of the steps, which are updated with the target count of each job. Like + /// @ref set_target_count for the first maximum. + void apply_step_truncation_budgets(); /// Count `count` events accepted by step `step`, which passed the earlier /// unweightings with the maxima `max_weight` (first step) and /// `max_weights` (the steps before). diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 4d2e8b9faa..b926ea56d7 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -1259,6 +1259,17 @@ void ChannelEventGenerator::update_step_max_weight( } } +void ChannelEventGenerator::apply_step_truncation_budgets() { + for (std::size_t step = 1; + auto [state, settings] : zip(_step_states, _step_settings)) { + if (settings.unweight && !state.max_weight_fixed && state.max_weight != 0 && + state.freeze_count <= settings.freeze_max_weight_after) { + apply_step_truncation_budget(step, _status.count_target); + } + ++step; + } +} + void ChannelEventGenerator::apply_step_truncation_budget( std::size_t step, std::size_t count_target ) { diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index 1733eeced4..934ffd497b 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -249,6 +249,10 @@ void EventGenerator::run_generate(GenerateMode mode) { } update_integral_fractions(); + // only with the steps drained, they are not in the commit order of the jobs + for (auto& channel : _channels) { + channel->apply_step_truncation_budgets(); + } update_counts(); if (_status.done) { diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 33a1d2b538..c77d17ee06 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -3515,6 +3515,17 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("step"), pydoc::doc("ChannelEventGenerator::step_max_weight") ) + .def( + "set_target_count", + &ChannelEventGenerator::set_target_count, + py::arg("target_count"), + pydoc::doc("ChannelEventGenerator::set_target_count") + ) + .def( + "apply_step_truncation_budgets", + &ChannelEventGenerator::apply_step_truncation_budgets, + pydoc::doc("ChannelEventGenerator::apply_step_truncation_budgets") + ) .def( "set_fixed_step_max_weight", &ChannelEventGenerator::set_fixed_step_max_weight, diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py index 79594f708b..0e67613d1e 100644 --- a/madspace/tests/test_weighting_steps.py +++ b/madspace/tests/test_weighting_steps.py @@ -79,12 +79,18 @@ def generate( step_channel=-1, step_kwargs=None, config_kwargs=None, + survey=True, ): """Survey and generate, returns the status, the combined events and the - channels. step_me and step_outputs can be lists, one entry per step.""" + channels. step_me and step_outputs can be lists, one entry per step, and + integrand_me one entry per channel.""" os.makedirs(run_dir, exist_ok=True) context, alphas_file = make_context(run_dir, threads) - me0 = python_matrix_element(context, integrand_me, integrand_outputs) + if not isinstance(integrand_me, list): + integrand_me = [integrand_me] * len(pt_mins) + me0s = [ + python_matrix_element(context, me, integrand_outputs) for me in integrand_me + ] if not isinstance(step_me, list): step_me = [step_me] * len(step_modes) step_outputs = [step_outputs] * len(step_modes) @@ -108,7 +114,7 @@ def generate( name=str(i), histograms=None, ) - for i, pt_min in enumerate(pt_mins) + for i, (pt_min, me0) in enumerate(zip(pt_mins, me0s)) ] steps = [ ms.WeightingStep( @@ -123,7 +129,8 @@ def generate( generator = ms.EventGenerator( [context], channels, seed=seed, config=config, groups=groups ) - generator.survey() + if survey: + generator.survey() generator.generate() output_file = os.path.join(run_dir, "combined.npy") generator.combine_to_compact_npy(output_file) @@ -240,6 +247,38 @@ def test_reproducible_step_freeze(tmp_path): assert status_a.mean == status_b.mean +def test_step_budget_of_smaller_target(tmp_path): + """A smaller target applies its truncation budget to the step maxima too, + and rescales the unweighted count accordingly. EventGenerator does this at + the end of each round, as generate() starts with a generous target.""" + small = lambda momenta: surrogate_me(momenta) * (momenta[:, 3, 1] > 300.0) + _, _, channels = generate( + str(tmp_path), + [surrogate_me, small], + ["unweight"], + step_me=heavy_tailed_me, + pt_mins=(20.0, 20.0), + grouped=(0, 1), + target_count=20000, + survey=False, + step_kwargs={"freeze_max_weight_after": 10**6}, + ) + channel = channels[1] + target = channel.status().count_target + assert target < 5000 + # the budget of the final target is already applied + max_weight, count = channel.step_max_weight(1), channel.status().count_unweighted + channel.apply_step_truncation_budgets() + assert channel.step_max_weight(1) == max_weight + channel.set_target_count(target // 20) + channel.apply_step_truncation_budgets() + new_max_weight = channel.step_max_weight(1) + assert new_max_weight > max_weight + assert channel.status().count_unweighted == pytest.approx( + count * max_weight / new_max_weight, rel=1e-12 + ) + + def broken_me(momenta, invalid=(np.nan, np.inf)): """true_me, with invalid values in two regions.""" px = momenta[:, 3, 1] From 2ae747e877f4861b8eeb26452cf45916e73ac56f Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:15:34 +0100 Subject: [PATCH 74/88] madspace: uniform mapping channel weights if no density is valid The fallback put all of the weight on the generating channel, so the weights were not a function of the event alone, and it compared a final channel index with a mapping index. Use uniform weights over the masked channels instead, which makes the channel_index argument of density_channel_weights and MappingChannelWeights unused, so it is removed. --- .../madspace/phasespace/channel_weights.hpp | 5 ++-- madspace/instruction_set.yaml | 5 +--- madspace/src/kernels/multichannel.hpp | 18 ++++++------- madspace/src/phasespace/channel_weights.cpp | 5 +--- madspace/src/phasespace/integrand.cpp | 8 +----- .../tests/test_mapping_channel_weights.py | 25 +++++++++++++++---- 6 files changed, 34 insertions(+), 32 deletions(-) diff --git a/madspace/include/madspace/phasespace/channel_weights.hpp b/madspace/include/madspace/phasespace/channel_weights.hpp index e66af441e6..b699f948b5 100644 --- a/madspace/include/madspace/phasespace/channel_weights.hpp +++ b/madspace/include/madspace/phasespace/channel_weights.hpp @@ -125,7 +125,8 @@ class SubchannelWeights : public FunctionGenerator { * discrete choices, and @f$m_{fk}@f$ is 1 if channel @f$k@f$ integrates * flavour @f$f@f$, else 0. The adaptive samplers do not enter, so the weights * of one channel do not depend on the state of the others. A channel whose - * mapping cannot produce the event gets the weight 0. + * mapping cannot produce the event gets the weight 0. If no channel of the + * flavour has a valid density, the weights are uniform over these channels. * * Optionally, it also samples a diagram for the event record: diagram * @f$d@f$ belongs to channel @f$c_d@f$ and is drawn with a probability @@ -138,8 +139,6 @@ class SubchannelWeights : public FunctionGenerator { * - `x1`, `x2` – `float`, shape `(batch,)` – the momentum fractions returned * by the mapping that generated the event. * - `flavor` – `int`, shape `(batch,)` – the flavour index. - * - `channel_index` – `int`, shape `(batch,)` – the channel that generated the - * event. * - `mirror_index` – `int`, shape `(batch,)` – the orientation of the beams. * Present only if the mappings are built with `mirror_beams`. * diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index 8c3b55b552..2d1408345e 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -2876,14 +2876,11 @@ density_channel_weights: - name: mask type: [float, n] desc: "1 for the channels that integrate the flavor of the event, else 0" - - name: channel_index - type: [int] - desc: "channel that generated the event" outputs: - name: channel_weights type: [float, n] desc: "weight per channel" - desc: "Normalised channel weights proportional to the masked channel densities, all of the weight on the generating channel if none is positive. Building block of @ref MappingChannelWeights." + desc: "Normalised channel weights proportional to the masked channel densities, uniform over the masked channels if none is positive. Building block of @ref MappingChannelWeights." interpolate_pdf: inputs: diff --git a/madspace/src/kernels/multichannel.hpp b/madspace/src/kernels/multichannel.hpp index cb3c22f2b2..25891594ac 100644 --- a/madspace/src/kernels/multichannel.hpp +++ b/madspace/src/kernels/multichannel.hpp @@ -46,28 +46,28 @@ KERNELSPEC void kernel_collect_channel_weights( template KERNELSPEC void kernel_density_channel_weights( - FIn densities, - FIn mask, - IIn channel_index, - FOut channel_weights + FIn densities, FIn mask, FOut channel_weights ) { // A density that is not a finite positive number (an inverse mapping that // failed) counts as zero: that channel cannot have produced the event. - FVal norm(0.); + FVal norm(0.), mask_count(0.); for (std::size_t i = 0; i < channel_weights.size(); ++i) { auto density = densities[i]; - auto valid = (density > 0.) & (density < 1e300) & (mask[i] > 0.); + auto masked = mask[i] > 0.; + auto valid = (density > 0.) & (density < 1e300) & masked; channel_weights[i] = where(valid, density, FVal(0.)); norm = norm + channel_weights[i]; + mask_count = mask_count + where(masked, FVal(1.), FVal(0.)); } - // the generating channel always has a positive density; if its inverse - // failed numerically, it keeps the whole weight + // if all inverses failed numerically, uniform over the masked channels, so + // that the weights stay a function of the event alone auto degenerate = norm == 0.; auto safe_norm = where(degenerate, FVal(1.), norm); + auto uniform = 1. / where(mask_count > 0., mask_count, FVal(1.)); for (std::size_t i = 0; i < channel_weights.size(); ++i) { channel_weights[i] = where( degenerate, - where(i == channel_index, FVal(1.), 0.), + where(mask[i] > 0., uniform, FVal(0.)), channel_weights[i] / safe_norm ); } diff --git a/madspace/src/phasespace/channel_weights.cpp b/madspace/src/phasespace/channel_weights.cpp index 0907c4614f..e3f309fa71 100644 --- a/madspace/src/phasespace/channel_weights.cpp +++ b/madspace/src/phasespace/channel_weights.cpp @@ -232,7 +232,6 @@ MappingChannelWeights::MappingChannelWeights( {"x1", batch_float}, {"x2", batch_float}, {"flavor", batch_int}, - {"channel_index", batch_int}, }; if (mappings.at(0).mirror_beams()) { args.push_back("mirror_index", batch_int); @@ -327,9 +326,7 @@ NamedVector MappingChannelWeights::build_function_impl( } Value flavor = args.at("flavor"); Value channel_weights = fb.density_channel_weights( - fb.stack(densities), - fb.gather_vector(flavor, _flavor_masks), - args.at("channel_index") + fb.stack(densities), fb.gather_vector(flavor, _flavor_masks) ); NamedVector ret{{"channel_weights", channel_weights}}; if (!_diagram_channels.empty()) { diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index 0e55b2c1aa..4a13ed1a5c 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -825,13 +825,7 @@ NamedVector Integrand::build_common_part( Value mapping_diagram_acc; if (_mapping_chan_weights && (channel_count > 1 || _mapping_chan_weights->has_diagrams())) { - ValueVec cw_args{ - momenta_acc, - x1_acc, - x2_acc, - flavor_id, - fb.batch_gather(indices_acc, args.at("chan_index")) - }; + ValueVec cw_args{momenta_acc, x1_acc, x2_acc, flavor_id}; if (_mapping_chan_weights->mirror_beams()) { cw_args.push_back(args.at("mirror_index_acc")); } diff --git a/madspace/tests/test_mapping_channel_weights.py b/madspace/tests/test_mapping_channel_weights.py index e233ed1b7f..eca5dfe122 100644 --- a/madspace/tests/test_mapping_channel_weights.py +++ b/madspace/tests/test_mapping_channel_weights.py @@ -105,10 +105,7 @@ def test_weights_are_the_normalised_masked_densities(flavor): p, x1, x2 = p[ok], x1[ok], x2[ok] n = p.shape[0] flavors = torch.full((n,), flavor, dtype=torch.int32) - channels = torch.full((n,), 1, dtype=torch.int32) - alpha, diagram = ( - torch.from_dlpack(t) for t in runtime.call([p, x1, x2, flavors, channels]) - ) + alpha, diagram = (torch.from_dlpack(t) for t in runtime.call([p, x1, x2, flavors])) expected = torch.stack( [density(RESONANT, p, x1, x2), density(PLAIN, p, x1, x2)], dim=1 ) * torch.tensor(masks[flavor]) @@ -136,9 +133,27 @@ def test_channels_without_the_flavor_get_no_weight(): x1[ok], x2[ok], torch.zeros(n, dtype=torch.int32), - torch.ones(n, dtype=torch.int32), ] )[0] ) torch.testing.assert_close(alpha[:, 1], torch.ones(n, dtype=alpha.dtype)) torch.testing.assert_close(alpha[:, 0], torch.zeros(n, dtype=alpha.dtype)) + + +def test_uniform_weights_without_valid_densities(): + """If no inverse gives a valid density, the weights are uniform over the + channels of the flavor: a function of the event alone.""" + weights = ms.MappingChannelWeights([RESONANT, PLAIN, PLAIN], [[1.0, 0.0, 1.0]]) + runtime = ms.FunctionRuntime(weights.function(), default_context()) + p, x1, x2, w = sample(PLAIN, 1000, 5) + ok = w > 0 + n = int(ok.sum()) + flavors = torch.zeros(n, dtype=torch.int32) + p, x1, x2 = p[ok], x1[ok], x2[ok] + alpha = torch.from_dlpack(runtime.call([p, x1, x2, flavors])[0]) + assert not torch.allclose(alpha[:, 0], alpha[:, 2]) + # inverse mappings that fail numerically + p = torch.full_like(p, float("nan")) + alpha = torch.from_dlpack(runtime.call([p, x1, x2, flavors])[0]) + expected = torch.tensor([0.5, 0.0, 0.5], dtype=alpha.dtype).expand(n, 3) + torch.testing.assert_close(alpha, expected) From d5d9a4eb7e4edd42ac6801baa5cde898b415751c Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:18:18 +0100 Subject: [PATCH 75/88] madspace: first-pass unweighting settings per weighting group WeightingGroup takes optional first_pass_max_overweight_truncation and first_pass_freeze_max_weight_after that replace the GeneratorConfig values for the first unweighting of the channels of the group only, so that looser first-pass settings do not leak into the channels without steps. They are stored in the group JSON if set. --- .../madspace/driver/channel_generator.hpp | 22 ++++++++ .../madspace/driver/weighting_pipeline.hpp | 6 +++ madspace/src/driver/channel_generator.cpp | 13 +++++ madspace/src/driver/event_generator.cpp | 8 ++- madspace/src/driver/weighting_pipeline.cpp | 24 ++++++++- madspace/src/python/madspace.cpp | 33 +++++++++++- madspace/tests/test_weighting_steps.py | 50 +++++++++++++++++++ 7 files changed, 151 insertions(+), 5 deletions(-) diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 547ce793bf..f96f6d7245 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -209,6 +209,27 @@ class ChannelEventGenerator { void set_weighting_steps(const std::vector& steps); /// The re- and unweighting steps. const std::vector& weighting_steps() const { return _step_settings; } + /// Replace @ref GeneratorConfig::max_overweight_truncation and + /// @ref GeneratorConfig::freeze_max_weight_after of the first unweighting, + /// where set. + void set_first_pass_settings( + std::optional max_overweight_truncation, + std::optional freeze_max_weight_after + ); + /// The truncation of the first unweighting, see + /// @ref GeneratorConfig::max_overweight_truncation. + double max_overweight_truncation() const { + return _config.max_overweight_truncation; + } + /// The freeze count of the first unweighting, see + /// @ref GeneratorConfig::freeze_max_weight_after. + std::size_t freeze_max_weight_after() const { + return _config.freeze_max_weight_after; + } + /// The freeze count set with @ref set_first_pass_settings, if any. + std::optional first_pass_freeze_max_weight_after() const { + return _first_pass_freeze_max_weight_after; + } /// Maximum weight of step `step`, counted from 1. double step_max_weight(std::size_t step) const { return _step_states.at(step - 1).max_weight; @@ -448,6 +469,7 @@ class ChannelEventGenerator { double _freeze_count = 0.; StepIntegral _final_integral; std::size_t _event_epoch = 0; + std::optional _first_pass_freeze_max_weight_after; friend void to_json(nlohmann::json& j, const ChannelEventGenerator& channel); }; diff --git a/madspace/include/madspace/driver/weighting_pipeline.hpp b/madspace/include/madspace/driver/weighting_pipeline.hpp index a61a6c6a47..0fe3db2d9a 100644 --- a/madspace/include/madspace/driver/weighting_pipeline.hpp +++ b/madspace/include/madspace/driver/weighting_pipeline.hpp @@ -22,6 +22,12 @@ struct WeightingGroup { std::vector channel_indices; /// The steps, in order. std::vector steps; + /// If set, replaces @ref GeneratorConfig::max_overweight_truncation for + /// the first unweighting of the channels of the group. + std::optional first_pass_max_overweight_truncation; + /// If set, replaces @ref GeneratorConfig::freeze_max_weight_after for the + /// first unweighting of the channels of the group. + std::optional first_pass_freeze_max_weight_after; }; /// JSON text of `groups`, for example to store them in a gridpack. The matrix diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index b926ea56d7..3d7e04095f 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -1211,6 +1211,19 @@ void ChannelEventGenerator::set_weighting_steps( } } +void ChannelEventGenerator::set_first_pass_settings( + std::optional max_overweight_truncation, + std::optional freeze_max_weight_after +) { + if (max_overweight_truncation) { + _config.max_overweight_truncation = *max_overweight_truncation; + } + if (freeze_max_weight_after) { + _config.freeze_max_weight_after = *freeze_max_weight_after; + _first_pass_freeze_max_weight_after = freeze_max_weight_after; + } +} + void ChannelEventGenerator::set_fixed_step_max_weight( std::size_t step, double max_weight ) { diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index 934ffd497b..273d4c9886 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -138,7 +138,8 @@ void EventGenerator::run_generate(GenerateMode mode) { if (channel->status().count_opt == 0 || mode == GenerateMode::fix_max_weights) { channel->set_target_count( mode == GenerateMode::fix_max_weights - ? _config.freeze_max_weight_after + ? channel->first_pass_freeze_max_weight_after() + .value_or(_config.freeze_max_weight_after) : _config.target_count ); } @@ -662,7 +663,10 @@ void EventGenerator::update_integral_fractions(bool final_step) { zip(_channels, counts, _channel_integral_fractions)) { if (_generate_mode == GenerateMode::fix_max_weights) { // the maxima of all steps have to freeze - std::size_t freeze_count = _config.freeze_max_weight_after; + std::size_t freeze_count = + channel->first_pass_freeze_max_weight_after().value_or( + _config.freeze_max_weight_after + ); for (auto& step : channel->weighting_steps()) { freeze_count = std::max(freeze_count, step.freeze_max_weight_after); } diff --git a/madspace/src/driver/weighting_pipeline.cpp b/madspace/src/driver/weighting_pipeline.cpp index c3b8b59fe9..c76f935a62 100644 --- a/madspace/src/driver/weighting_pipeline.cpp +++ b/madspace/src/driver/weighting_pipeline.cpp @@ -146,7 +146,16 @@ madspace::weighting_groups_to_json(const std::vector& groups) { {"override_outputs", override_outputs}, }); } - j.push_back({{"channel_indices", group.channel_indices}, {"steps", steps}}); + json group_json{{"channel_indices", group.channel_indices}, {"steps", steps}}; + if (group.first_pass_max_overweight_truncation) { + group_json["first_pass_max_overweight_truncation"] = + *group.first_pass_max_overweight_truncation; + } + if (group.first_pass_freeze_max_weight_after) { + group_json["first_pass_freeze_max_weight_after"] = + *group.first_pass_freeze_max_weight_after; + } + j.push_back(group_json); } return j.dump(); } @@ -159,6 +168,14 @@ madspace::weighting_groups_from_json(const std::string& text) { .channel_indices = group_json.at("channel_indices").get>() }; + if (group_json.contains("first_pass_max_overweight_truncation")) { + group.first_pass_max_overweight_truncation = + group_json.at("first_pass_max_overweight_truncation").get(); + } + if (group_json.contains("first_pass_freeze_max_weight_after")) { + group.first_pass_freeze_max_weight_after = + group_json.at("first_pass_freeze_max_weight_after").get(); + } for (auto& step_json : group_json.at("steps")) { WeightingStep step{ .mode = step_json.at("mode").get() == "unweight" @@ -359,6 +376,11 @@ WeightingPipeline::WeightingPipeline( ); for (auto channel_index : group_config.channel_indices) { channels.at(channel_index)->set_weighting_steps(group.settings); + channels.at(channel_index) + ->set_first_pass_settings( + group_config.first_pass_max_overweight_truncation, + group_config.first_pass_freeze_max_weight_after + ); _next_sequences.at(channel_index).assign(step_count + 1, 0); _ready.at(channel_index).resize(step_count + 1); } diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index c77d17ee06..8ef84e3e13 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -3515,6 +3515,16 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("step"), pydoc::doc("ChannelEventGenerator::step_max_weight") ) + .def( + "max_overweight_truncation", + &ChannelEventGenerator::max_overweight_truncation, + pydoc::doc("ChannelEventGenerator::max_overweight_truncation") + ) + .def( + "freeze_max_weight_after", + &ChannelEventGenerator::freeze_max_weight_after, + pydoc::doc("ChannelEventGenerator::freeze_max_weight_after") + ) .def( "set_target_count", &ChannelEventGenerator::set_target_count, @@ -3599,15 +3609,34 @@ PYBIND11_MODULE(_madspace_py, m) { py::classh(m, "WeightingGroup", pydoc::doc("WeightingGroup")) .def( py::init([](const std::vector& channel_indices, - const std::vector& steps) { + const std::vector& steps, + std::optional first_pass_max_overweight_truncation, + std::optional first_pass_freeze_max_weight_after) { return WeightingGroup{ - .channel_indices = channel_indices, .steps = steps + .channel_indices = channel_indices, + .steps = steps, + .first_pass_max_overweight_truncation = + first_pass_max_overweight_truncation, + .first_pass_freeze_max_weight_after = + first_pass_freeze_max_weight_after, }; }), py::arg("channel_indices"), py::arg("steps"), + py::arg("first_pass_max_overweight_truncation") = py::none(), + py::arg("first_pass_freeze_max_weight_after") = py::none(), pydoc::doc("WeightingGroup") ) + .def_readwrite( + "first_pass_max_overweight_truncation", + &WeightingGroup::first_pass_max_overweight_truncation, + pydoc::doc("WeightingGroup::first_pass_max_overweight_truncation") + ) + .def_readwrite( + "first_pass_freeze_max_weight_after", + &WeightingGroup::first_pass_freeze_max_weight_after, + pydoc::doc("WeightingGroup::first_pass_freeze_max_weight_after") + ) .def_readwrite( "channel_indices", &WeightingGroup::channel_indices, diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py index 0e67613d1e..a2a08b31d8 100644 --- a/madspace/tests/test_weighting_steps.py +++ b/madspace/tests/test_weighting_steps.py @@ -467,6 +467,56 @@ def test_ratio_matrix_element_rules(tmp_path): assert loaded[0].steps[0].matrix_elements[0].channel() == 3 +def test_first_pass_settings_of_group(tmp_path): + """The first unweighting of the channels of a group uses the settings of + the group, the other channels the global ones.""" + context, alphas_file = make_context(str(tmp_path), 2) + surrogate = python_matrix_element(context, surrogate_me) + full = python_matrix_element(context, true_me) + config = make_config() + channels = [ + ms.ChannelEventGenerator( + [context], + make_integrand(alphas_file, surrogate), + str(tmp_path / f"events.{name}.npy"), + str(tmp_path / f"weights.{name}.npy"), + config, + 0, + name, + None, + ) + for name in ("group", "plain") + ] + group = ms.WeightingGroup( + [0], + [ms.WeightingStep([full], freeze_max_weight_after=300)], + first_pass_max_overweight_truncation=0.2, + first_pass_freeze_max_weight_after=2000, + ) + generator = ms.EventGenerator( + [context], channels, seed=3, config=config, groups=[group] + ) + assert [c.max_overweight_truncation() for c in channels] == [0.2, 0.05] + assert [c.freeze_max_weight_after() for c in channels] == [2000, 1000] + generator.survey() + # each channel runs until its own maximum freezes + generator.fix_max_weights() + assert [c.status().count_target for c in channels] == [2000, 1000] + + # stored with the group, old JSON without them still loads + text = ms.weighting_groups_to_json([group]) + loaded = ms.weighting_groups_from_json(text)[0] + assert loaded.first_pass_max_overweight_truncation == 0.2 + assert loaded.first_pass_freeze_max_weight_after == 2000 + old = json.loads(text) + del old[0]["first_pass_max_overweight_truncation"] + del old[0]["first_pass_freeze_max_weight_after"] + loaded = ms.weighting_groups_from_json(json.dumps(old))[0] + assert loaded.first_pass_max_overweight_truncation is None + assert loaded.first_pass_freeze_max_weight_after is None + assert "first_pass" not in ms.weighting_groups_to_json([loaded]) + + def test_invalid_groups(tmp_path): context, alphas_file = make_context(str(tmp_path), 1) me = python_matrix_element(context, true_me) From 5f542acbb7825913bd1e55f81dceede35670e41a Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:20:01 +0100 Subject: [PATCH 76/88] madspace: EventStream status without the inactive channels Channels without fixed maximum weights return no events, but their integral was still added to the cross section of status(). --- madspace/include/madspace/driver/event_stream.hpp | 5 +++-- madspace/src/driver/event_stream.cpp | 9 ++++++++- madspace/tests/test_weighting_steps.py | 15 +++++++++++++++ 3 files changed, 26 insertions(+), 3 deletions(-) diff --git a/madspace/include/madspace/driver/event_stream.hpp b/madspace/include/madspace/driver/event_stream.hpp index 7839f886df..535154a2f4 100644 --- a/madspace/include/madspace/driver/event_stream.hpp +++ b/madspace/include/madspace/driver/event_stream.hpp @@ -82,8 +82,9 @@ class EventStream { /// Stop the coordinator thread and discard the running jobs. Called by /// the destructor. void close(); - /// Combined status over every channel. `count_unweighted` is the number of - /// events returned so far. + /// Combined status over the channels with fixed maximum weights, the others + /// return no events. `count_unweighted` is the number of events returned + /// so far. GeneratorStatus status() const; /// Per-channel status; `count_unweighted` is the number of events of the /// channel returned so far. diff --git a/madspace/src/driver/event_stream.cpp b/madspace/src/driver/event_stream.cpp index 8779a10a91..0403617856 100644 --- a/madspace/src/driver/event_stream.cpp +++ b/madspace/src/driver/event_stream.cpp @@ -488,7 +488,14 @@ std::vector EventStream::next_events(std::size_t count) { GeneratorStatus EventStream::status() const { std::unique_lock lock(_mutex); GeneratorStatus status{}; - sum_channel_status(_channels, status); + // the inactive channels return no events + std::vector> active_channels; + for (auto [channel, state] : zip(_channels, _channel_states)) { + if (state.active) { + active_channels.push_back(channel); + } + } + sum_channel_status(active_channels, status); status.count_unweighted = _event_count; status.count_target = 0; status.optimized = true; diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py index a2a08b31d8..d242c0467b 100644 --- a/madspace/tests/test_weighting_steps.py +++ b/madspace/tests/test_weighting_steps.py @@ -707,3 +707,18 @@ def test_event_stream_cross_section(fixed_gridpack): assert abs(status.mean - stream_status.mean) < 4 * np.hypot( status.error, stream_status.error ) + + +def test_event_stream_status_without_inactive_channels(fixed_gridpack): + """A channel without fixed step maxima returns no events, so it does not + count in the cross section of the stream.""" + gridpack = dict(fixed_gridpack) + channel_json = json.loads(gridpack["channels"][1]) + del channel_json["step_max_weights"] + gridpack["channels"] = [gridpack["channels"][0], json.dumps(channel_json)] + stream = make_stream(gridpack) + integral = json.loads(gridpack["channels"][0])["integral"] + status = stream.status() + assert status.mean == pytest.approx(integral["mean"], rel=1e-12) + assert status.error == pytest.approx(integral["error"], rel=1e-12) + assert stream.channel_probabilities()[1] == 0.0 From c116dd6f19c9b6aedea066d9c6044022f4ddb540 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:39:37 +0100 Subject: [PATCH 77/88] mg7: first-pass unweighting settings per weighting group, not for the whole run --- docs/source/gridpacks.rst | 3 +- docs/source/run_card.rst | 9 +- .../iolibs/template_files/mg7/event_stream.py | 5 +- .../iolibs/template_files/mg7/gridpack.py | 13 +- .../template_files/mg7/gridpack_setup.py | 21 ++- madgraph/mg7/madgraph_process.py | 33 ++--- madgraph/mg7/subprocess_amplicol.py | 6 + madgraph/mg7/subprocess_madmatrix.py | 7 + .../interface/test_mg7_weighting_passes.py | 127 ++++++++++++++++++ 9 files changed, 184 insertions(+), 40 deletions(-) create mode 100644 tests/unit_tests/interface/test_mg7_weighting_passes.py diff --git a/docs/source/gridpacks.rst b/docs/source/gridpacks.rst index 2a63f5972b..315d33ee4b 100644 --- a/docs/source/gridpacks.rst +++ b/docs/source/gridpacks.rst @@ -175,7 +175,8 @@ directory. Both are reset to ``""`` when the gridpack is created. With ``[amplicol]`` processes unweighted to full colour, ``max_overweight_truncation`` and ``freeze_max_weight_after`` set the second pass, which gives the events, and the -``first_pass_`` options the leading-colour unweighting, as in the main run. The colour +``first_pass_`` options the leading-colour unweighting of these processes, as in the main +run. The other processes are unweighted once, as set by ``[generation]``. The colour accuracy and the mode are fixed when the gridpack is created. The pyamplicol artifacts the libraries read are copied to ``SubProcesses/``, and ``include_source = true`` is not supported for these processes. diff --git a/docs/source/run_card.rst b/docs/source/run_card.rst index 3bbd01ecac..f8e23ca2aa 100644 --- a/docs/source/run_card.rst +++ b/docs/source/run_card.rst @@ -403,10 +403,11 @@ The surrogate does not return channel weights, so it needs ``sde_strategy = overweights. ``"reweight"`` keeps the ratio in the event weights. ``first_pass_max_overweight_truncation`` (``0.01``), ``first_pass_freeze_max_weight_after`` (``10000``) - ``max_overweight_truncation`` and ``freeze_max_weight_after`` of the unweighting with - the surrogate, in unweight mode. The ones of ``[generation]`` are used for the second - pass, which gives the final events. In reweight mode, ``[generation]`` sets the - unweighting with the surrogate. + ``max_overweight_truncation`` and ``freeze_max_weight_after`` of the first + unweighting, with the surrogate, in unweight mode. The ones of ``[generation]`` set + the second unweighting with the full matrix element, which gives the final events, + and the unweighting of the subprocesses without a surrogate. In reweight mode, the + events are unweighted once, with the surrogate, as set by ``[generation]``. ``param_card`` (``""``) The param card a ``.so`` surrogate is initialized with. An empty string uses diff --git a/madgraph/iolibs/template_files/mg7/event_stream.py b/madgraph/iolibs/template_files/mg7/event_stream.py index b4f8c71564..62bfcfcb62 100644 --- a/madgraph/iolibs/template_files/mg7/event_stream.py +++ b/madgraph/iolibs/template_files/mg7/event_stream.py @@ -190,8 +190,9 @@ def pick(value, default): load_context_data([self._context], [backend], madspace_data) channel_generators = load_channels([self._context], config) # the maxima of all passes are fixed, so their settings do not matter - groups, _ = load_weighting_groups( - gen_args["max_overweight_truncation"], gen_args["freeze_max_weight_after"] + groups = load_weighting_groups( + gen_args["max_overweight_truncation"], gen_args["freeze_max_weight_after"], + gen_args["max_overweight_truncation"], gen_args["freeze_max_weight_after"], ) self._stream = ms.EventStream( self._context, channel_generators, self.seed, load_lhe_completer(), config, diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 1cadef9463..5d704bec26 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -222,15 +222,12 @@ def resolve_dir(cli_value, card_value): # set up generator configuration config = ms.GeneratorConfig() config.target_count = args.events - groups, first_pass = load_weighting_groups( - args.max_overweight_truncation, args.freeze_max_weight_after + groups = load_weighting_groups( + args.max_overweight_truncation, args.freeze_max_weight_after, + args.first_pass_max_overweight_truncation, args.first_pass_freeze_max_weight_after, ) - if first_pass: - config.max_overweight_truncation = args.first_pass_max_overweight_truncation - config.freeze_max_weight_after = args.first_pass_freeze_max_weight_after - else: - config.max_overweight_truncation = args.max_overweight_truncation - config.freeze_max_weight_after = args.freeze_max_weight_after + config.max_overweight_truncation = args.max_overweight_truncation + config.freeze_max_weight_after = args.freeze_max_weight_after config.cpu_batch_size = args.cpu_batch_size config.gpu_batch_size = args.gpu_batch_size config.verbosity = resolve_verbosity(args.verbosity) diff --git a/madgraph/iolibs/template_files/mg7/gridpack_setup.py b/madgraph/iolibs/template_files/mg7/gridpack_setup.py index 74976c7baa..89fce519f6 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack_setup.py +++ b/madgraph/iolibs/template_files/mg7/gridpack_setup.py @@ -139,19 +139,30 @@ def load_channels(contexts, config, file_dir=None): ] -def load_weighting_groups(max_overweight_truncation, freeze_max_weight_after): +def load_weighting_groups( + max_overweight_truncation, freeze_max_weight_after, + first_pass_max_overweight_truncation, first_pass_freeze_max_weight_after, +): """The re- and unweighting steps after the first unweighting stored in - data/groups.json, with the unweighting settings of the final pass, and - whether a step unweights: the first unweighting is then a first pass.""" + data/groups.json, with the unweighting settings of the final pass. The + first unweighting of a group with an unweighting step is a first pass.""" with open(gridpack_path("data", "groups.json")) as f: groups = json.load(f) - unweights = False for group in groups: + unweights = False for step in group["steps"]: step["max_overweight_truncation"] = max_overweight_truncation step["freeze_max_weight_after"] = freeze_max_weight_after unweights |= step["mode"] == "unweight" - return ms.weighting_groups_from_json(json.dumps(groups)), unweights + # the card replaces the stored values, which older gridpacks lack + group.pop("first_pass_max_overweight_truncation", None) + group.pop("first_pass_freeze_max_weight_after", None) + if unweights: + group["first_pass_max_overweight_truncation"] = ( + first_pass_max_overweight_truncation + ) + group["first_pass_freeze_max_weight_after"] = first_pass_freeze_max_weight_after + return ms.weighting_groups_from_json(json.dumps(groups)) def load_lhe_completer(): diff --git a/madgraph/mg7/madgraph_process.py b/madgraph/mg7/madgraph_process.py index 8e77069c27..bb69a53ba6 100644 --- a/madgraph/mg7/madgraph_process.py +++ b/madgraph/mg7/madgraph_process.py @@ -417,8 +417,19 @@ def weighting_groups(self, subprocess_channels: list[list[int]]) -> list: groups = [] for subproc, channel_indices in zip(self.subprocesses, subprocess_channels): steps = subproc.weighting_steps() - if channel_indices and steps: - groups.append(ms.WeightingGroup(channel_indices, steps)) + if not (channel_indices and steps): + continue + group = ms.WeightingGroup(channel_indices, steps) + # [generation] sets the last unweighting, the first_pass_ ones the earlier + first_pass = subproc.first_pass_settings() + if first_pass is not None: + group.first_pass_max_overweight_truncation = ( + first_pass["first_pass_max_overweight_truncation"] + ) + group.first_pass_freeze_max_weight_after = ( + first_pass["first_pass_freeze_max_weight_after"] + ) + groups.append(group) return groups @staticmethod @@ -1319,24 +1330,6 @@ def init_generator_config(self) -> None: cfg.vegas_damping = vegas_args["damping"] cfg.max_overweight_truncation = gen_args["max_overweight_truncation"] cfg.freeze_max_weight_after = gen_args["freeze_max_weight_after"] - first_pass = None - if self.surrogate is not None and self.surrogate["mode"] == "unweight": - first_pass = self.surrogate - amplicol = self.run_card["amplicol"] - if ( - first_pass is None and amplicol["mode"] == "unweight" - and amplicol["color_accuracy"] != "lc" - and any(is_amplicol(meta) for meta in self.subprocess_data) - ): - first_pass = amplicol - if first_pass is not None: - # [generation] applies to the second pass, which gives the events - cfg.max_overweight_truncation = ( - first_pass["first_pass_max_overweight_truncation"] - ) - cfg.freeze_max_weight_after = ( - first_pass["first_pass_freeze_max_weight_after"] - ) cfg.start_batch_size = vegas_args["start_batch_size"] cfg.max_batch_size = vegas_args["max_batch_size"] cfg.survey_min_iters = gen_args["survey_min_iters"] diff --git a/madgraph/mg7/subprocess_amplicol.py b/madgraph/mg7/subprocess_amplicol.py index cce1635cde..45339c6514 100644 --- a/madgraph/mg7/subprocess_amplicol.py +++ b/madgraph/mg7/subprocess_amplicol.py @@ -436,3 +436,9 @@ def weighting_steps(self) -> list: max_overweight_truncation=gen_args["max_overweight_truncation"], freeze_max_weight_after=gen_args["freeze_max_weight_after"], )] + + def first_pass_settings(self) -> dict | None: + amplicol = self.process.run_card["amplicol"] + if self.ratio_matrix_element is None or amplicol["mode"] != "unweight": + return None + return amplicol diff --git a/madgraph/mg7/subprocess_madmatrix.py b/madgraph/mg7/subprocess_madmatrix.py index f77eaccdc7..c28f412cfe 100644 --- a/madgraph/mg7/subprocess_madmatrix.py +++ b/madgraph/mg7/subprocess_madmatrix.py @@ -452,6 +452,13 @@ def weighting_steps(self) -> list: freeze_max_weight_after=gen_args["freeze_max_weight_after"], )] + def first_pass_settings(self) -> dict | None: + """The run card section with the first_pass_ unweighting settings, if + the events are unweighted again after the first unweighting.""" + if not self.surrogates or self.process.surrogate["mode"] != "unweight": + return None + return self.process.surrogate + def flavor_pids(self) -> list[list[int]]: """PDG ids of the external particles for every flavor of the integrand, in the order of its flavor index.""" diff --git a/tests/unit_tests/interface/test_mg7_weighting_passes.py b/tests/unit_tests/interface/test_mg7_weighting_passes.py new file mode 100644 index 0000000000..39d3cafcf0 --- /dev/null +++ b/tests/unit_tests/interface/test_mg7_weighting_passes.py @@ -0,0 +1,127 @@ +################################################################################ +# +# Copyright (c) 2026 The MadGraph5_aMC@NLO Development team and Contributors +# +# This file is a part of the MadGraph5_aMC@NLO project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph5_aMC@NLO license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""The [surrogate] and [amplicol] second passes of an mg7 run: the run card +settings of each pass and the checks done at launch.""" + +import importlib.util +import json +import os +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +from madgraph.mg7 import bootstrap as mg7_bootstrap + + +@unittest.skipUnless(mg7_bootstrap.madspace_is_installed(), + 'madspace is not installed') +class FirstPassSettingsTest(unittest.TestCase): + """The first_pass_ settings apply to the subprocesses unweighted again in a + second pass only, the [generation] ones to all others.""" + + def setUp(self): + import madgraph.various.banner as banner_mod + from madgraph.mg7 import madgraph_process + from madgraph.mg7.subprocess_amplicol import AmplicolSubprocess + from madgraph.mg7.subprocess_madmatrix import MadgraphSubprocess + self.ms = madgraph_process.ms + self.process = object.__new__(madgraph_process.MadgraphProcess) + run_card = banner_mod.RunCardMG7() + run_card['generation']['max_overweight_truncation'] = 0.001 + run_card['generation']['freeze_max_weight_after'] = 500 + run_card['amplicol']['mode'] = 'unweight' + run_card['amplicol']['first_pass_max_overweight_truncation'] = 0.02 + run_card['amplicol']['first_pass_freeze_max_weight_after'] = 70 + self.process.run_card = run_card + self.process.surrogate = None + self.process.subprocess_data = [] + step = self.ms.WeightingStep([], mode='unweight') + + amplicol = object.__new__(AmplicolSubprocess) + amplicol.process = self.process + amplicol.ratio_matrix_element = object() + amplicol.weighting_steps = lambda: [step] + madmatrix = object.__new__(MadgraphSubprocess) + madmatrix.process = self.process + madmatrix.surrogates = [] + self.process.subprocesses = [amplicol, madmatrix] + + def test_mixed_output(self): + groups = self.process.weighting_groups([[0, 1], [2]]) + self.assertEqual(len(groups), 1) + self.assertEqual(list(groups[0].channel_indices), [0, 1]) + self.assertAlmostEqual(groups[0].first_pass_max_overweight_truncation, 0.02) + self.assertEqual(groups[0].first_pass_freeze_max_weight_after, 70) + + self.process.init_generator_config() + cfg = self.process.event_generator_config + self.assertAlmostEqual(cfg.max_overweight_truncation, 0.001) + self.assertEqual(cfg.freeze_max_weight_after, 500) + + def test_reweight_mode_has_no_first_pass(self): + self.process.run_card['amplicol']['mode'] = 'reweight' + groups = self.process.weighting_groups([[0, 1], [2]]) + self.assertIsNone(groups[0].first_pass_max_overweight_truncation) + self.assertIsNone(groups[0].first_pass_freeze_max_weight_after) + + +_TEMPLATES = (Path(__file__).resolve().parents[3] / 'madgraph' / 'iolibs' + / 'template_files' / 'mg7') + + +def _template_module(name): + spec = importlib.util.spec_from_file_location( + f'mg7_template_{name}', _TEMPLATES / f'{name}.py') + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@unittest.skipUnless(mg7_bootstrap.madspace_is_installed(), + 'madspace is not installed') +class GridpackFirstPassTest(unittest.TestCase): + """A gridpack takes the first_pass_ settings from its card, for the groups + that unweight in a later step, also if its groups.json predates them.""" + + def load(self, groups): + setup = _template_module('gridpack_setup') + with tempfile.TemporaryDirectory() as tmp: + os.mkdir(os.path.join(tmp, 'data')) + with open(os.path.join(tmp, 'data', 'groups.json'), 'w') as f: + json.dump(groups, f) + with mock.patch.object(setup, 'gridpack_path', + lambda *parts: os.path.join(tmp, *parts)): + return setup.load_weighting_groups(0.001, 500, 0.02, 70) + + def test_old_and_new_groups(self): + ms = _template_module('gridpack_setup').ms + step = ms.WeightingStep([], mode='unweight') + stored = json.loads(ms.weighting_groups_to_json([ + ms.WeightingGroup([0], [step], 0.5, 3), ms.WeightingGroup([1], [step]), + ])) + reweight = json.loads(ms.weighting_groups_to_json([ + ms.WeightingGroup([2], [ms.WeightingStep([], mode='reweight')], 0.5, 3) + ])) + groups = self.load(stored + reweight) + for group in groups[:2]: + self.assertAlmostEqual(group.first_pass_max_overweight_truncation, 0.02) + self.assertEqual(group.first_pass_freeze_max_weight_after, 70) + self.assertAlmostEqual(group.steps[0].max_overweight_truncation, 0.001) + self.assertIsNone(groups[2].first_pass_max_overweight_truncation) + + +if __name__ == '__main__': + unittest.main() From 7628853fe177390d8692bae25bf075c085b1f465 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:40:02 +0100 Subject: [PATCH 78/88] mg7: refuse a [surrogate] with [amplicol] subprocesses --- docs/source/run_card.rst | 2 +- madgraph/mg7/madgraph_process.py | 5 +++++ .../interface/test_mg7_weighting_passes.py | 19 +++++++++++++++++++ 3 files changed, 25 insertions(+), 1 deletion(-) diff --git a/docs/source/run_card.rst b/docs/source/run_card.rst index f8e23ca2aa..339dcb8248 100644 --- a/docs/source/run_card.rst +++ b/docs/source/run_card.rst @@ -380,7 +380,7 @@ its optimization use the surrogate. The events are then reweighted or unweighted time with the full matrix element, which is only evaluated for the events that pass the first unweighting. The cross section and the events are those of the full matrix element. The surrogate does not return channel weights, so it needs ``sde_strategy = -"denominators"``. +"denominators"``. It is not supported for outputs with ``[amplicol]`` subprocesses. ``function`` (``""``) The surrogate. An empty string disables it. diff --git a/madgraph/mg7/madgraph_process.py b/madgraph/mg7/madgraph_process.py index bb69a53ba6..ac27e661aa 100644 --- a/madgraph/mg7/madgraph_process.py +++ b/madgraph/mg7/madgraph_process.py @@ -391,6 +391,11 @@ def init_surrogate(self) -> None: if not settings["function"]: return path = settings["function"] + if any(is_amplicol(meta) for meta in self.subprocess_data): + raise RuntimeError( + "[surrogate] is not supported with [amplicol] subprocesses: " + "set [surrogate] function = \"\" or use the madmatrix backend" + ) if path.endswith(".py") and not os.path.isfile(path): raise RuntimeError(f"[surrogate] function = {path!r}: file not found") if self.run_card["phasespace"]["sde_strategy"] != "denominators": diff --git a/tests/unit_tests/interface/test_mg7_weighting_passes.py b/tests/unit_tests/interface/test_mg7_weighting_passes.py index 39d3cafcf0..7eb0296821 100644 --- a/tests/unit_tests/interface/test_mg7_weighting_passes.py +++ b/tests/unit_tests/interface/test_mg7_weighting_passes.py @@ -78,6 +78,25 @@ def test_reweight_mode_has_no_first_pass(self): self.assertIsNone(groups[0].first_pass_freeze_max_weight_after) +@unittest.skipUnless(mg7_bootstrap.madspace_is_installed(), + 'madspace is not installed') +class SurrogateWithAmplicolTest(unittest.TestCase): + + def test_is_refused(self): + import madgraph.various.banner as banner_mod + from madgraph.mg7 import madgraph_process + process = object.__new__(madgraph_process.MadgraphProcess) + process.run_card = banner_mod.RunCardMG7() + process.run_card['surrogate']['function'] = 'surrogate.py' + process.subprocess_data = [{'me_path': 'a'}, {'backend': 'amplicol'}] + with self.assertRaisesRegex(RuntimeError, r'\[surrogate\] is not supported'): + process.init_surrogate() + # without [amplicol] subprocesses, the next check is reached + process.subprocess_data = process.subprocess_data[:1] + with self.assertRaisesRegex(RuntimeError, 'file not found'): + process.init_surrogate() + + _TEMPLATES = (Path(__file__).resolve().parents[3] / 'madgraph' / 'iolibs' / 'template_files' / 'mg7') From cc5232ffa551099133bf0d1c2af55385b6bdb8dd Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:42:25 +0100 Subject: [PATCH 79/88] mg7: event stream docs: reweighted [amplicol] events carry ratio weights --- docs/source/event_stream.rst | 21 ++- .../iolibs/template_files/mg7/event_stream.py | 11 +- madspace/tests/test_event_stream_reweight.py | 132 ++++++++++++++++++ 3 files changed, 153 insertions(+), 11 deletions(-) create mode 100644 madspace/tests/test_event_stream_reweight.py diff --git a/docs/source/event_stream.rst b/docs/source/event_stream.rst index 306c7b14ca..cfe49193b1 100644 --- a/docs/source/event_stream.rst +++ b/docs/source/event_stream.rst @@ -68,7 +68,7 @@ with one array per LHE field. The particle fields have the shape cross_section = stream.integral print(f"cross section: {cross_section:.1f} +- {stream.integral_error:.1f} pb") -The event weights are ``1``, except for overweight events, so they have to be normalized. +The event weights are not normalized to the cross section (see :ref:`event-stream-weights`). Scaling them to the cross section gives the differential cross section in pb/GeV: .. code-block:: python @@ -110,16 +110,23 @@ Events have different numbers of particles if the process has several subprocess includes decays. The particle fields are then padded with zeros. ``particle_count`` holds the number of particles of every event and ``particle_mask`` marks the valid entries. +.. _event-stream-weights: + Weights and cross section ------------------------- -Every event weight is ``1``, or ``-1`` for an event with a negative weight. The exception are -overweight events, whose weight exceeded the fixed maximum weight of their channel. They -keep the ratio of their weight to the maximum weight, so their weight is larger than ``1``. -Treating them as ``1`` would bias the distributions in the phase-space regions where they -occur. +After the unweighting, every event weight is ``1``, or ``-1`` for an event with a negative +weight. The exception are overweight events, whose weight exceeded the fixed maximum weight +of their channel. They keep the ratio of their weight to the maximum weight, so their weight +is larger than ``1``. Treating them as ``1`` would bias the distributions in the phase-space +regions where they occur. + +The events of ``[amplicol]`` subprocesses are then reweighted to full colour in the default +``mode = "reweight"``: their weight is multiplied by the ratio of the full-colour to the +leading-colour matrix element, so it is not ``1``. With ``mode = "unweight"``, they are +unweighted a second time and have unit weights apart from overweights. -To normalize a sample to the cross section, scale its weights by +In all cases, to normalize a sample to the cross section, scale its weights by ``stream.integral / sum(weights)``. The cross section estimate improves while events are generated: it combines the integrals stored in the gridpack with the weighted events generated by the stream so far. Read it after generating the sample, as in the example diff --git a/madgraph/iolibs/template_files/mg7/event_stream.py b/madgraph/iolibs/template_files/mg7/event_stream.py index 62bfcfcb62..3fade84fe3 100644 --- a/madgraph/iolibs/template_files/mg7/event_stream.py +++ b/madgraph/iolibs/template_files/mg7/event_stream.py @@ -18,9 +18,11 @@ batch.px # shape (10000, max_particle_count) Event weights are in units of the maximum weight: +-1, except for overweight -events. To normalize a sample to the cross section, scale its weights by -``stream.integral / sum(weights)``. The event sequence is reproducible from the seed, -independent of the thread count and of how it is split into batches. +events and for events reweighted in a later step (the default [amplicol] +mode = "reweight"), which carry the ratio of the matrix elements. To normalize +a sample to the cross section, scale its weights by +``stream.integral / sum(weights)``. The event sequence is reproducible from the +seed, independent of the thread count and of how it is split into batches. """ import dataclasses @@ -62,7 +64,8 @@ class EventBatch: #: Process ID of the event (IDPRUP). process_id: Any - #: Event weight, +-1 or larger in magnitude for overweight events (XWGTUP). + #: Event weight (XWGTUP), +-1 or larger in magnitude for overweight events, times + #: the matrix-element ratio of a reweighting step. Not normalized. weight: Any #: Renormalization scale in GeV (SCALUP). scale: Any diff --git a/madspace/tests/test_event_stream_reweight.py b/madspace/tests/test_event_stream_reweight.py new file mode 100644 index 0000000000..2babcce75b --- /dev/null +++ b/madspace/tests/test_event_stream_reweight.py @@ -0,0 +1,132 @@ +"""EventStream with a reweighting step after the first unweighting, like the +default [amplicol] mode = "reweight": the streamed events carry the ratio of the +step matrix element to the integrand one, not unit weights. Uses the setup of +test_weighting_steps. +""" + +import os + +import numpy as np +import pytest +from test_event_stream import CM_ENERGY, lhe_completer, make_config, make_context +from test_integrand_me_outputs import python_matrix_element +from test_weighting_steps import FULL_ME_OUTPUTS, generate, make_integrand, true_me + +import madspace as ms + +PT_MINS = (20.0, 60.0) + + +def biased_me(momenta): + """Does not average to the true one, so that the ratio changes the norm.""" + return true_me(momenta) * (0.6 + 0.8 * momenta[:, 2, 0] / CM_ENERGY) + + +def full_me_with_indices(momenta): + zeros = np.zeros(len(momenta), np.int32) + return true_me(momenta), zeros, zeros, zeros + + +def load_matrix_elements(context): + return ( + python_matrix_element(context, biased_me), + python_matrix_element(context, full_me_with_indices, FULL_ME_OUTPUTS), + ) + + +@pytest.fixture(scope="module") +def reweight_gridpack(tmp_path_factory): + """Two grouped channels with fixed maxima and a reweighting step.""" + config = make_config() + run_dir = str(tmp_path_factory.mktemp("reweight")) + context, alphas_file = make_context(run_dir, 2) + biased, full = load_matrix_elements(context) + groups = [ms.WeightingGroup([0, 1], [ms.WeightingStep([full], mode="reweight")])] + channels = [ + ms.ChannelEventGenerator( + contexts=[context], + integrand=make_integrand(alphas_file, biased, pt_min), + event_file=os.path.join(run_dir, f"events.{i}.npy"), + weight_file=os.path.join(run_dir, f"weights.{i}.npy"), + config=config, + subprocess_index=0, + name=str(i), + histograms=None, + ) + for i, pt_min in enumerate(PT_MINS) + ] + generator = ms.EventGenerator( + [context], channels, seed=1, config=config, groups=groups + ) + generator.survey() + generator.fix_max_weights() + globals_file = os.path.join(run_dir, "globals.tar") + context.save_globals(globals_file) + return { + "channels": [channel.to_json(True) for channel in channels], + "groups": ms.weighting_groups_to_json(groups), + "globals": globals_file, + } + + +def make_stream(gridpack, thread_count=4, seed=5): + context = ms.Context(ms.cpu_device(), thread_count) + context.load_globals(gridpack["globals"]) + load_matrix_elements(context) + config = make_config() + channels = [ + ms.ChannelEventGenerator.load_json( + channel_json, [context], event_file="", weight_file="", config=config + ) + for channel_json in gridpack["channels"] + ] + groups = ms.weighting_groups_from_json(gridpack["groups"]) + return ms.EventStream( + context, channels, seed, lhe_completer(), config, groups=groups + ) + + +def test_event_stream_reweight(reweight_gridpack, tmp_path): + stream = make_stream(reweight_gridpack) + batch = stream.next_batch(20000) + weights = batch["weight"] + # the ratio of the full to the biased matrix element, times the first-pass + # weight (1, or more for the overweights) + ratio = 1.0 / (0.6 + 0.8 * batch["energy"][:, 2] / CM_ENERGY) + assert np.mean(np.isclose(weights, ratio, rtol=1e-6)) > 0.9 + assert np.all(weights >= ratio * (1.0 - 1e-6)) + + # mean weight x first-pass cross section = full cross section + first_pass, _, _ = generate( + str(tmp_path / "biased"), + biased_me, + pt_mins=PT_MINS, + grouped=(0, 1), + target_count=1000, + ) + full, full_events, _ = generate( + str(tmp_path / "full"), + true_me, + pt_mins=PT_MINS, + grouped=(0, 1), + target_count=20000, + ) + status = stream.status() + mean_weight_error = np.std(weights) / np.sqrt(len(weights)) + xs_error = np.hypot( + mean_weight_error * first_pass.mean, np.mean(weights) * first_pass.error + ) + assert abs(np.mean(weights) * first_pass.mean - full.mean) < 4 * np.hypot( + xs_error, full.error + ) + assert abs(status.mean - full.mean) < 4 * np.hypot(status.error, full.error) + + # the weighted events follow the full matrix element + energies = [ + np.sum(w * e) / np.sum(w) + for w, e in ( + (weights, batch["energy"][:, 2]), + (full_events["weight"], full_events["part3_energy"]), + ) + ] + assert energies[0] == pytest.approx(energies[1], rel=0.01) From 08d6e0ad79ea4512df85eca6e9d58f4dd3ee2458 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:44:24 +0100 Subject: [PATCH 80/88] mg7: event stream keeps the unread events of an iterator chunk, refuses batch_size <= 0 --- docs/source/event_stream.rst | 3 +- .../iolibs/template_files/mg7/event_stream.py | 48 +++++++- .../iolibs/test_mg7_event_stream.py | 105 ++++++++++++++++++ 3 files changed, 150 insertions(+), 6 deletions(-) create mode 100644 tests/unit_tests/iolibs/test_mg7_event_stream.py diff --git a/docs/source/event_stream.rst b/docs/source/event_stream.rst index cfe49193b1..c3ce6391c8 100644 --- a/docs/source/event_stream.rst +++ b/docs/source/event_stream.rst @@ -98,7 +98,8 @@ Output formats There are three ways to get events from a stream: - Iterating over the stream yields one ``madspace.LHEEvent`` at a time. The events are - completed in chunks of 1000 internally. + completed in chunks of 1000 internally. If the loop is left early, the rest of the chunk + is returned first by the next call, so no event is skipped. - :py:meth:`next_events ` returns a list of a given number of ``LHEEvent`` objects. - :py:meth:`batches ` and diff --git a/madgraph/iolibs/template_files/mg7/event_stream.py b/madgraph/iolibs/template_files/mg7/event_stream.py index 3fade84fe3..ac0f59ce73 100644 --- a/madgraph/iolibs/template_files/mg7/event_stream.py +++ b/madgraph/iolibs/template_files/mg7/event_stream.py @@ -25,6 +25,7 @@ seed, independent of the thread count and of how it is split into batches. """ +import collections import dataclasses import json import sys @@ -108,6 +109,27 @@ def __len__(self) -> int: return len(self.weight) +def _events_to_arrays(events: list, template: dict) -> dict: + """The fields of ms.LHEEvent objects as arrays like those of + ms.EventStream.next_batch, with the dtypes and padding of ``template``.""" + arrays = {} + for key, value in template.items(): + if key == "particle_count": + data = [len(event.particles) for event in events] + elif value.ndim == 1: + data = [getattr(event, key) for event in events] + else: + data = np.zeros((len(events), value.shape[1]), value.dtype) + for row, event in zip(data, events): + row[: len(event.particles)] = [getattr(p, key) for p in event.particles] + arrays[key] = np.asarray(data, value.dtype) + return arrays + + +def _concatenate(first: dict, second: dict) -> dict: + return {key: np.concatenate([first[key], second[key]]) for key in first} + + class EventStream: """Endless stream of unweighted events from this gridpack. @@ -145,6 +167,8 @@ def __init__( ignore_source_hash: bool = False, ): self._stream = None + # events fetched for the iterator but not returned yet + self._pending = collections.deque() run_card = load_run_card() run_args = run_card["run"] gen_args = run_card["generation"] @@ -204,20 +228,32 @@ def pick(value, default): def __iter__(self) -> Iterator["ms.LHEEvent"]: """Endless iterator over the events as ms.LHEEvent objects. It takes - events from the stream in chunks of 1000; those of the last chunk that - are not iterated over are dropped.""" + events from the stream in chunks of 1000. The events of a chunk that are + not iterated over are kept: the next calls return them first, so no + event is lost when the loop is left early.""" while True: - yield from self._stream.next_events(_ITER_CHUNK) + if not self._pending: + self._pending.extend(self._stream.next_events(_ITER_CHUNK)) + yield self._pending.popleft() + + def _take_pending(self, count: int) -> list: + return [self._pending.popleft() for _ in range(min(count, len(self._pending)))] def next_events(self, count: int) -> list: """The next ``count`` events as ms.LHEEvent objects.""" - return self._stream.next_events(count) + events = self._take_pending(count) + if len(events) < count: + events.extend(self._stream.next_events(count - len(events))) + return events def next_batch(self, batch_size: int, output_type: str = "numpy") -> EventBatch: """The next ``batch_size`` events as an :class:`EventBatch`.""" if output_type not in ("numpy", "torch"): raise ValueError('output_type must be "numpy" or "torch"') - arrays = self._stream.next_batch(batch_size) + pending = self._take_pending(batch_size) + arrays = self._stream.next_batch(batch_size - len(pending)) + if pending: + arrays = _concatenate(_events_to_arrays(pending, arrays), arrays) counts = arrays["particle_count"] max_count = arrays["pdg_id"].shape[1] arrays["particle_mask"] = np.arange(max_count)[None, :] < counts[:, None] @@ -234,6 +270,8 @@ def batches( ) -> Iterator[EventBatch]: """Iterator over batches of ``batch_size`` events, endless unless ``max_events`` is given (the last batch may then be smaller).""" + if batch_size <= 0: + raise ValueError("batch_size must be positive") remaining = max_events while remaining is None or remaining > 0: size = batch_size if remaining is None else min(batch_size, remaining) diff --git a/tests/unit_tests/iolibs/test_mg7_event_stream.py b/tests/unit_tests/iolibs/test_mg7_event_stream.py new file mode 100644 index 0000000000..518b0b0e94 --- /dev/null +++ b/tests/unit_tests/iolibs/test_mg7_event_stream.py @@ -0,0 +1,105 @@ +"""The Python EventStream of the mg7 gridpacks (bin/event_stream.py), on a fake +madspace stream: no event is lost between the ways of reading the stream.""" + +import collections +import importlib.util +import sys +import types +import unittest +from pathlib import Path +from unittest import mock + +import numpy as np + +_TEMPLATE = (Path(__file__).resolve().parents[3] / "madgraph" / "iolibs" + / "template_files" / "mg7" / "event_stream.py") + + +def _event_stream_module(): + setup = types.ModuleType("gridpack_setup") + setup.__getattr__ = lambda name: None + spec = importlib.util.spec_from_file_location("mg7_event_stream", _TEMPLATE) + module = importlib.util.module_from_spec(spec) + with mock.patch.dict(sys.modules, {"madspace": types.ModuleType("madspace"), + "gridpack_setup": setup}): + spec.loader.exec_module(module) + return module + + +Particle = collections.namedtuple("Particle", "pdg_id px") +Event = collections.namedtuple("Event", "weight particles") + + +class FakeStream: + """Event i has weight i and i % 3 + 1 particles.""" + + def __init__(self): + self.count = 0 + + def _events(self, count): + events = [ + Event(float(i), [Particle(21, float(i + j)) for j in range(i % 3 + 1)]) + for i in range(self.count, self.count + count) + ] + self.count += count + return events + + def next_events(self, count): + return self._events(count) + + def close(self): + pass + + def next_batch(self, count): + events = self._events(count) + px = np.zeros((count, 3)) + for row, event in zip(px, events): + row[: len(event.particles)] = [p.px for p in event.particles] + return { + "weight": np.array([e.weight for e in events]), + "particle_count": np.array([len(e.particles) for e in events], np.int32), + "pdg_id": np.where(px != 0.0, 21, 0).astype(np.int32), + "px": px, + } + + +class TestEventStream(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.module = _event_stream_module() + # the fake stream has only some of the LHE fields + cls.module.EventBatch = types.SimpleNamespace + + def make_stream(self): + stream = object.__new__(self.module.EventStream) + stream._stream = FakeStream() + stream._pending = collections.deque() + return stream + + def test_no_event_is_lost_after_leaving_the_loop(self): + stream = self.make_stream() + for event in stream: + if event.weight == 4.0: + break + self.assertEqual([e.weight for e in stream.next_events(3)], [5.0, 6.0, 7.0]) + batch = stream.next_batch(4) + self.assertEqual(list(batch.weight), [8.0, 9.0, 10.0, 11.0]) + self.assertEqual(list(batch.particle_count), [3, 1, 2, 3]) + self.assertEqual(list(batch.px[0]), [8.0, 9.0, 10.0]) + self.assertEqual(list(batch.px[2]), [10.0, 11.0, 0.0]) + self.assertEqual(batch.particle_mask[2].tolist(), [True, True, False]) + self.assertEqual(batch.pdg_id.dtype, np.int32) + # the rest of the chunk, then the stream itself + batch = stream.next_batch(1000) + self.assertEqual(list(batch.weight), [float(i) for i in range(12, 1012)]) + self.assertEqual(next(iter(stream)).weight, 1012.0) + + def test_batches_needs_a_positive_size(self): + stream = self.make_stream() + for size in (0, -1): + with self.assertRaises(ValueError): + next(stream.batches(size, max_events=10)) + + +if __name__ == "__main__": + unittest.main() From 76fc11dc6f647444aaa2cf1112f4c28e8265981d Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:44:59 +0100 Subject: [PATCH 81/88] mg7: gridpack runs without an [amplicol] section in its card --- .../iolibs/template_files/mg7/gridpack.py | 6 +- .../iolibs/test_mg7_gridpack_card.py | 75 +++++++++++++++++++ 2 files changed, 78 insertions(+), 3 deletions(-) create mode 100644 tests/unit_tests/iolibs/test_mg7_gridpack_card.py diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 5d704bec26..1046b928e3 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -77,7 +77,7 @@ def main() -> None: run_args = run_card["run"] gen_args = run_card["generation"] grid_args = run_card.get("gridpack", {}) - amplicol_args = run_card["amplicol"] + amplicol_args = run_card.get("amplicol", {}) madspace_data = load_madspace_data() # parse command line arguments @@ -137,13 +137,13 @@ def main() -> None: parser.add_argument("--freeze_max_weight_after", type=int, default=gen_args["freeze_max_weight_after"]) parser.add_argument( "--first_pass_max_overweight_truncation", type=float, - default=amplicol_args["first_pass_max_overweight_truncation"], + default=amplicol_args.get("first_pass_max_overweight_truncation", 0.01), help="max_overweight_truncation of the leading-colour unweighting of " "[amplicol] processes unweighted to full colour" ) parser.add_argument( "--first_pass_freeze_max_weight_after", type=int, - default=amplicol_args["first_pass_freeze_max_weight_after"], + default=amplicol_args.get("first_pass_freeze_max_weight_after", 10000), help="freeze_max_weight_after of the leading-colour unweighting of " "[amplicol] processes unweighted to full colour" ) diff --git a/tests/unit_tests/iolibs/test_mg7_gridpack_card.py b/tests/unit_tests/iolibs/test_mg7_gridpack_card.py new file mode 100644 index 0000000000..1da76709bd --- /dev/null +++ b/tests/unit_tests/iolibs/test_mg7_gridpack_card.py @@ -0,0 +1,75 @@ +"""Command line defaults of the standalone mg7 gridpack launcher, taken from +its Cards/grid_run_card.toml.""" + +import argparse +import importlib.util +import sys +import types +import unittest +from pathlib import Path +from unittest import mock + +_TEMPLATE = (Path(__file__).resolve().parents[3] / "madgraph" / "iolibs" + / "template_files" / "mg7" / "gridpack.py") + + + +def _run_card(**sections): + """The default run card as plain TOML data, with sections replaced (None + removes one).""" + import madgraph.various.banner as banner_mod + card = banner_mod.RunCardMG7() + data = {key: dict(card[key]) for key in ("run", "generation", "gridpack", "amplicol")} + data["generation"]["max_overweight_truncation"] = 0.001 + for key, value in sections.items(): + if value is None: + del data[key] + else: + data[key] = value + return data + + +class _Parsed(Exception): + pass + + +def _defaults(run_card): + """The parsed arguments of main() without command line options.""" + loader = types.ModuleType("custom_function_loader") + setup = types.ModuleType("gridpack_setup") + setup.__getattr__ = lambda name: None + setup.load_run_card = lambda: run_card + setup.load_madspace_data = dict + spec = importlib.util.spec_from_file_location("mg7_gridpack_card", _TEMPLATE) + module = importlib.util.module_from_spec(spec) + with mock.patch.dict(sys.modules, {"madspace": types.ModuleType("madspace"), + "custom_function_loader": loader, + "gridpack_setup": setup}): + spec.loader.exec_module(module) + + def parse_args(parser): + raise _Parsed(parser.parse_known_args([])[0]) + + with mock.patch.object(argparse.ArgumentParser, "parse_args", parse_args): + try: + module.main() + except _Parsed as parsed: + return parsed.args[0] + raise AssertionError("main() did not parse the arguments") + + +class TestGridpackCard(unittest.TestCase): + def test_card_without_amplicol_section(self): + args = _defaults(_run_card(amplicol=None)) + self.assertEqual(args.first_pass_max_overweight_truncation, 0.01) + self.assertEqual(args.first_pass_freeze_max_weight_after, 10000) + self.assertEqual(args.max_overweight_truncation, 0.001) + + def test_amplicol_section(self): + args = _defaults(_run_card(amplicol={"first_pass_max_overweight_truncation": 0.2})) + self.assertEqual(args.first_pass_max_overweight_truncation, 0.2) + self.assertEqual(args.first_pass_freeze_max_weight_after, 10000) + + +if __name__ == "__main__": + unittest.main() From f14dae8d7bb34ab4df2248f48a9a92bb6e164239 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:46:45 +0100 Subject: [PATCH 82/88] mg7: check the .so [surrogate] path and device, pass it the run-time ME parameters --- docs/source/custom_functions.rst | 21 +++++++ docs/source/run_card.rst | 14 +++-- .../mg7/custom_function_loader.py | 52 +++++++++++++---- madgraph/mg7/madgraph_process.py | 3 +- madgraph/mg7/subprocess_madmatrix.py | 1 + .../interface/test_mg7_weighting_passes.py | 56 +++++++++++++++++++ 6 files changed, 128 insertions(+), 19 deletions(-) diff --git a/docs/source/custom_functions.rst b/docs/source/custom_functions.rst index a8eb54166e..78bd6f7a14 100644 --- a/docs/source/custom_functions.rst +++ b/docs/source/custom_functions.rst @@ -290,6 +290,27 @@ matrix products. If your function relies on it, set ``OMP_NUM_THREADS=1``. The Python API behind these functions is ``madspace.PythonFunction``. It can also be used directly with madspace. +.. _custom-function-surrogate: + +Python surrogates +----------------- + +A Python surrogate of the matrix element (see :ref:`run-card-surrogate`) uses the same +mechanism. Its file defines ``FRAMEWORK`` and ``custom_surrogate(momenta, pdg_ids, +alpha_s)``, which returns :math:`|M|^2` as a float64 array of shape ``(batch,)``. It takes +no parameters. ``momenta`` and ``pdg_ids`` have the shapes and types given above, with the +momenta that the full matrix element gets, and ``alpha_s`` is a float64 array of shape +``(batch,)``. The function is called for all subprocesses, which ``pdg_ids`` identifies: + +.. code-block:: python + + import numpy as np + + FRAMEWORK = "numpy" + + def custom_surrogate(momenta, pdg_ids, alpha_s): + return np.ones(len(momenta)) + Examples -------- diff --git a/docs/source/run_card.rst b/docs/source/run_card.rst index 339dcb8248..af7d8ba1c4 100644 --- a/docs/source/run_card.rst +++ b/docs/source/run_card.rst @@ -385,17 +385,19 @@ The surrogate does not return channel weights, so it needs ``sde_strategy = ``function`` (``""``) The surrogate. An empty string disables it. - - A ``.so`` library implementing the UMAMI interface with the inputs of the madmatrix - libraries (momenta, :math:`\alpha_s`, flavor index and the random numbers). It is - loaded once per subprocess: ``{subprocess}`` and ``{device}`` in the path are replaced - by the subprocess name (e.g. ``P1_gg_ttx``) and the device, like in - ``lib/libmadmatrix_{subprocess}_{device}.so``. + - A ``.so`` library implementing the UMAMI interface. It gets the momenta, + :math:`\alpha_s` and the flavor index of the madmatrix libraries as inputs, and + returns the matrix element. It is loaded once per subprocess: ``{subprocess}`` and + ``{device}`` in the path are replaced by the subprocess name (e.g. ``P1_gg_ttx``) + and the device, like in ``lib/libmadmatrix_{subprocess}_{device}.so``. It has to + be built for the device of the run. If it implements ``umami_set_parameter``, it + gets the run-time parameters of the madmatrix libraries, such as ``bwcutoff``. - A Python file defining ``custom_surrogate(momenta, pdg_ids, alpha_s)``, which returns :math:`|M|^2` for a batch of events. ``momenta`` has the shape ``(batch, n, 4)``, ``pdg_ids`` the shape ``(batch, n)`` and identifies subprocess and flavor, and ``alpha_s`` the shape ``(batch,)``. ``FRAMEWORK = "torch"`` passes torch tensors, which also run on GPUs, otherwise numpy arrays are passed. See - :doc:`custom_functions`. + :ref:`custom-function-surrogate`. ``mode`` (``"unweight"``) ``"unweight"`` unweights the events a second time with respect to the ratio of the full diff --git a/madgraph/iolibs/template_files/mg7/custom_function_loader.py b/madgraph/iolibs/template_files/mg7/custom_function_loader.py index 29b555cd1f..6fcc99c176 100644 --- a/madgraph/iolibs/template_files/mg7/custom_function_loader.py +++ b/madgraph/iolibs/template_files/mg7/custom_function_loader.py @@ -155,25 +155,53 @@ def load_custom_function(ms, contexts, device_types, kind, path, parameters=None ] +def _exports(path, symbol): + """Whether the shared library at ``path`` exports ``symbol``.""" + import ctypes + + return hasattr(ctypes.CDLL(os.path.abspath(path)), symbol) + + def load_surrogate( - ms, contexts, backends, device_types, path, subprocess, pdg_ids, param_card + ms, contexts, backends, device_types, path, subprocess, pdg_ids, param_card, + parameters=None, ): """Load the [surrogate] function for one subprocess into each context, and return the ``MatrixElementApi`` of each context. - A .so path is a UMAMI library with the inputs of the madmatrix libraries, - where ``{subprocess}`` and ``{device}`` are replaced, initialized with - ``param_card``. A .py file defines - ``custom_surrogate(momenta, pdg_ids, alpha_s)``, which returns |M|^2; - ``pdg_ids[flavor_index]`` are the PDG ids of the subprocess for each + A .so path is a UMAMI library with the momenta, alpha_s and flavor index as + inputs, where ``{subprocess}`` and ``{device}`` are replaced, initialized + with ``param_card``. It gets the run-time ``parameters`` ({name: value}) of + the madmatrix libraries if it implements umami_set_parameter. A .py file + defines ``custom_surrogate(momenta, pdg_ids, alpha_s)``, which returns + |M|^2; ``pdg_ids[flavor_index]`` are the PDG ids of the subprocess for each flavor index of its matrix element.""" if path.endswith(".so"): - return [ - context.load_matrix_element( - path.format(subprocess=subprocess, device=backend), param_card - ) - for context, backend in zip(contexts, backends) - ] + apis = [] + for context, backend, device in zip(contexts, backends, device_types): + library = path.format(subprocess=subprocess, device=backend) + if not os.path.isfile(library): + raise RuntimeError(f"[surrogate] function = {path!r}: {library} not found") + library_parameters = {} + if parameters and _exports(library, "umami_set_parameter"): + library_parameters = {str(k): float(v) for k, v in parameters.items()} + try: + api = context.load_matrix_element(library, param_card, library_parameters) + except RuntimeError as error: + if "could not set parameter" not in str(error): + raise + raise RuntimeError( + f"[surrogate] function = {path!r}: {error}. The surrogate gets the " + f"run-time parameters of the madmatrix libraries " + f"({', '.join(library_parameters)}) and has to accept them." + ) from error + if api.device_type() != device: + raise RuntimeError( + f"[surrogate] function = {path!r}: {library} was built for " + f"{api.device_type()}, but the run uses {device}" + ) + apis.append(api) + return apis if not path.endswith(".py"): raise RuntimeError( f"[surrogate] function = {path!r}: expected a .so library or a .py file" diff --git a/madgraph/mg7/madgraph_process.py b/madgraph/mg7/madgraph_process.py index ac27e661aa..857d96df38 100644 --- a/madgraph/mg7/madgraph_process.py +++ b/madgraph/mg7/madgraph_process.py @@ -396,7 +396,8 @@ def init_surrogate(self) -> None: "[surrogate] is not supported with [amplicol] subprocesses: " "set [surrogate] function = \"\" or use the madmatrix backend" ) - if path.endswith(".py") and not os.path.isfile(path): + # a .so path with {subprocess} or {device} is checked when it is loaded + if "{" not in path and not os.path.isfile(path): raise RuntimeError(f"[surrogate] function = {path!r}: file not found") if self.run_card["phasespace"]["sde_strategy"] != "denominators": raise RuntimeError( diff --git a/madgraph/mg7/subprocess_madmatrix.py b/madgraph/mg7/subprocess_madmatrix.py index c28f412cfe..c035e110c2 100644 --- a/madgraph/mg7/subprocess_madmatrix.py +++ b/madgraph/mg7/subprocess_madmatrix.py @@ -431,6 +431,7 @@ def load_surrogate(self, me_path_format: str, member_meta: dict): name, pdg_ids, settings["param_card"], + me_parameters(self.process.run_card), ) if len({api.index() for api in apis}) != 1: raise RuntimeError("surrogate loaded at different indices") diff --git a/tests/unit_tests/interface/test_mg7_weighting_passes.py b/tests/unit_tests/interface/test_mg7_weighting_passes.py index 7eb0296821..29a615ffa2 100644 --- a/tests/unit_tests/interface/test_mg7_weighting_passes.py +++ b/tests/unit_tests/interface/test_mg7_weighting_passes.py @@ -142,5 +142,61 @@ def test_old_and_new_groups(self): self.assertIsNone(groups[2].first_pass_max_overweight_truncation) +class FakeApi: + def __init__(self, device): + self.device = device + + def device_type(self): + return self.device + + +class FakeContext: + def __init__(self, device='cpu'): + self.device = device + self.loaded = [] + + def load_matrix_element(self, path, param_card, parameters): + self.loaded.append((path, param_card, parameters)) + return FakeApi(self.device) + + +class SurrogateLibraryTest(unittest.TestCase): + """A .so [surrogate]: path substitution, missing files, device and the + run-time parameters of the madmatrix libraries.""" + + def setUp(self): + self.loader = _template_module('custom_function_loader') + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.path = os.path.join(self.tmp.name, 'sur_{subprocess}_{device}.so') + open(self.path.format(subprocess='P1_x', device='cpp512y'), 'w').close() + + def load(self, contexts, devices=('cpu',), exports=True, subprocess='P1_x'): + with mock.patch.object(self.loader, '_exports', lambda path, name: exports): + return self.loader.load_surrogate( + None, contexts, ['cpp512y'] * len(contexts), list(devices), + self.path, subprocess, [[1, 2]], 'card.dat', {'bwcutoff': 15}, + ) + + def test_parameters_are_passed(self): + context = FakeContext() + self.load([context]) + library = self.path.format(subprocess='P1_x', device='cpp512y') + self.assertEqual(context.loaded, [(library, 'card.dat', {'bwcutoff': 15.0})]) + + def test_no_parameters_without_set_parameter(self): + context = FakeContext() + self.load([context], exports=False) + self.assertEqual(context.loaded[0][2], {}) + + def test_missing_library(self): + with self.assertRaisesRegex(RuntimeError, 'sur_P2_y_cpp512y.so not found'): + self.load([FakeContext()], subprocess='P2_y') + + def test_wrong_device(self): + with self.assertRaisesRegex(RuntimeError, 'built for cpu, but the run uses cuda'): + self.load([FakeContext('cpu')], devices=('cuda',)) + + if __name__ == '__main__': unittest.main() From e5a0a895f01b3a8243b1208b0f52b485f2f7a149 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:47:14 +0100 Subject: [PATCH 83/88] mg7: run check-pyamplicol before compiling the [amplicol] libraries and report its message --- madgraph/mg7/madgraph_process.py | 17 +++++++ .../interface/test_mg7_weighting_passes.py | 47 +++++++++++++++++++ 2 files changed, 64 insertions(+) diff --git a/madgraph/mg7/madgraph_process.py b/madgraph/mg7/madgraph_process.py index 857d96df38..653fc9678e 100644 --- a/madgraph/mg7/madgraph_process.py +++ b/madgraph/mg7/madgraph_process.py @@ -1468,10 +1468,27 @@ def compile_amplicol(self) -> None: ) }) log_path = os.path.join(self.run_path, "compile_amplicol.log") + for check in sorted({ + os.path.join(os.path.dirname(os.path.normpath(path)), "check-pyamplicol") + for path in paths + }): + # older outputs have no check script + if os.path.isfile(check): + self.check_pyamplicol(check) for path in paths: logger.info(f"Compiling the AmpliCol matrix elements in {path}, see {log_path}") self.make_subprocesses(path, [], 1, log_path) + @staticmethod + def check_pyamplicol(script: str) -> None: + """Run SubProcesses/check-pyamplicol, which fails if the installed + pyamplicol is not the one the output was made with.""" + result = subprocess.run([script], capture_output=True, text=True) + if result.returncode != 0: + message = (result.stderr.strip() or result.stdout.strip() + or f"{script} failed with exit status {result.returncode}") + raise RuntimeError(message) + @staticmethod def make_subprocesses( subproc_path: str, args: list[str], nb_core: int, log_path: str diff --git a/tests/unit_tests/interface/test_mg7_weighting_passes.py b/tests/unit_tests/interface/test_mg7_weighting_passes.py index 29a615ffa2..c6957f6a5b 100644 --- a/tests/unit_tests/interface/test_mg7_weighting_passes.py +++ b/tests/unit_tests/interface/test_mg7_weighting_passes.py @@ -198,5 +198,52 @@ def test_wrong_device(self): self.load([FakeContext('cpu')], devices=('cuda',)) +@unittest.skipUnless(mg7_bootstrap.madspace_is_installed(), + 'madspace is not installed') +class CompileAmplicolTest(unittest.TestCase): + """SubProcesses/check-pyamplicol runs before make, and its message is the + error of the launch.""" + + def setUp(self): + from madgraph.mg7 import madgraph_process + self.module = madgraph_process + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.subproc = os.path.join(self.tmp.name, 'SubProcesses') + os.makedirs(os.path.join(self.subproc, 'P1_gg_ttx')) + self.process = object.__new__(madgraph_process.MadgraphProcess) + self.process.run_path = self.tmp.name + self.process.subprocess_data = [{ + 'backend': 'amplicol', + 'path': os.path.join(self.subproc, 'P1_gg_ttx'), + 'me_path': 'missing.so', 'ratio_path': 'missing.so', + }] + + def write_check(self, body): + path = os.path.join(self.subproc, 'check-pyamplicol') + with open(path, 'w') as f: + f.write('#!/bin/sh\n' + body) + os.chmod(path, 0o755) + + def compile(self): + with mock.patch.object(self.module.MadgraphProcess, 'make_subprocesses') as make: + self.process.compile_amplicol() + return make + + def test_stale_pyamplicol_stops_before_make(self): + self.write_check('echo "reinstall pyamplicol 1.2" >&2\nexit 1\n') + with mock.patch.object(self.module.MadgraphProcess, 'make_subprocesses') as make: + with self.assertRaisesRegex(RuntimeError, 'reinstall pyamplicol 1.2'): + self.process.compile_amplicol() + make.assert_not_called() + + def test_matching_pyamplicol_compiles(self): + self.write_check('exit 0\n') + self.compile().assert_called_once() + + def test_older_output_without_check(self): + self.compile().assert_called_once() + + if __name__ == '__main__': unittest.main() From d84453f3e74d91b44f89a51776821e6c3ed71577 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:47:23 +0100 Subject: [PATCH 84/88] mg7: remove an unused start_time --- madgraph/mg7/madgraph_process.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/madgraph/mg7/madgraph_process.py b/madgraph/mg7/madgraph_process.py index 653fc9678e..3b9107b3fe 100644 --- a/madgraph/mg7/madgraph_process.py +++ b/madgraph/mg7/madgraph_process.py @@ -9,7 +9,7 @@ import tempfile import time -from madgraph.mg7.common import backend_of, fptype_of, get_start_time, interference_helicities_summed, lhe_helicities, lhe_weight_info, logger, me_parameters, mlm_pdf_ladder, ms, resolve_auto_backend, resolve_seed, resolve_verbosity +from madgraph.mg7.common import backend_of, fptype_of, interference_helicities_summed, lhe_helicities, lhe_weight_info, logger, me_parameters, mlm_pdf_ladder, ms, resolve_auto_backend, resolve_seed, resolve_verbosity from madgraph.various import misc from madgraph.various.banner import RunCardMG7 from madgraph.models.check_param_card import ParamCard @@ -1774,7 +1774,6 @@ def generate_events(self) -> None: if self.gridpack_run_mode != "regular": self.finish_without_events() return - start_time = get_start_time() self.event_generator.generate() output_format = self.run_card["run"]["output_format"] systematics = self.build_systematics() From 73a21485f7fc2afe7b6b8b62b1bb1dc698433fc1 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:50:03 +0100 Subject: [PATCH 85/88] madspace: split the flushed weighting-step batches over the idle threads --- .../madspace/driver/event_generator.hpp | 3 +- .../madspace/driver/weighting_pipeline.hpp | 19 ++- madspace/src/driver/event_generator.cpp | 10 +- madspace/src/driver/weighting_pipeline.cpp | 122 +++++++++++++----- 4 files changed, 116 insertions(+), 38 deletions(-) diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index ebed03d570..b0079843ec 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -257,7 +257,8 @@ class EventGenerator { void drain_pipeline(); // unweighting efficiency of each pass, the first one before and after cuts std::tuple> pass_efficiencies() const; - // start the jobs of the steps whose batches are full + // start the jobs of the steps whose batches are full, and the partial ones + // once the generation jobs leave threads idle std::size_t start_pipeline_jobs(); // Layout of the combined events as read from the channel files DataLayout combined_layout() const; diff --git a/madspace/include/madspace/driver/weighting_pipeline.hpp b/madspace/include/madspace/driver/weighting_pipeline.hpp index 0fe3db2d9a..6b3ab9f345 100644 --- a/madspace/include/madspace/driver/weighting_pipeline.hpp +++ b/madspace/include/madspace/driver/weighting_pipeline.hpp @@ -76,9 +76,14 @@ class WeightingPipeline { static bool owns_job(std::size_t job_id) { return job_id & job_id_flag; } /// Take the unweighted events of `job`, in the job order of its channel. void add_events(GeneratorBatchJob& job); - /// Start the jobs of full batches, or of all collected events with - /// `flush`. Returns the number of started jobs. - std::size_t start_jobs(ResultQueue& result_queue, bool flush); + /// Start the jobs of full batches. With `flush`, all collected events are + /// split over the free threads of each context, where `busy_counts` are + /// the other jobs running on it. Returns the number of started jobs. + std::size_t start_jobs( + ResultQueue& result_queue, + bool flush, + const std::vector& busy_counts = {} + ); /// Process the result of the finished job `job_id`. void commit_job(std::size_t job_id); /// Whether events are still collected or in jobs. @@ -143,6 +148,13 @@ class WeightingPipeline { void set_ready(ChunkPtr chunk, std::size_t stage); void commit_stage(ChunkPtr chunk, std::size_t stage); void decide(Chunk& chunk, std::size_t step); + void start_job( + ResultQueue& result_queue, + std::size_t group_index, + std::size_t context_index, + std::size_t stage, + std::vector chunks + ); void run_job(Job& job); void gather_events(Chunk& chunk, ContextRuntimes& runtimes); std::size_t batch_size(std::size_t context_index) const; @@ -158,6 +170,7 @@ class WeightingPipeline { std::vector> _next_sequences; std::vector>> _ready; std::unordered_map _jobs; + std::vector _context_job_counts; std::size_t _job_counter = 0; std::size_t _chunk_count = 0; Sink _sink; diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index 53fa19f543..84e6e60f4b 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -595,7 +595,15 @@ std::size_t EventGenerator::start_jobs() { } std::size_t EventGenerator::start_pipeline_jobs() { - return _pipeline ? _pipeline->start_jobs(_result_queue, false) : 0; + if (!_pipeline) { + return 0; + } + // once no more events are dispatched, partial batches use the idle threads + bool flush = _ready_jobs.empty() && + std::all_of(_context_unweight_queue.begin(), + _context_unweight_queue.end(), + [](auto& queue) { return queue.empty(); }); + return _pipeline->start_jobs(_result_queue, flush, _context_job_counts); } // Batches no longer have to be full: the steps don't depend on the batching. diff --git a/madspace/src/driver/weighting_pipeline.cpp b/madspace/src/driver/weighting_pipeline.cpp index c76f935a62..2fed7bcc6a 100644 --- a/madspace/src/driver/weighting_pipeline.cpp +++ b/madspace/src/driver/weighting_pipeline.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include "madspace/util.hpp" @@ -241,6 +242,7 @@ WeightingPipeline::WeightingPipeline( _channel_sequences(channels.size()), _next_sequences(channels.size()), _ready(channels.size()), + _context_job_counts(contexts.size()), _sink(sink) { if (!_sink) { _sink = [this](StepEvents& step_events) { @@ -560,47 +562,99 @@ void WeightingPipeline::decide(Chunk& chunk, std::size_t step) { } } -std::size_t WeightingPipeline::start_jobs(ResultQueue& result_queue, bool flush) { +std::size_t WeightingPipeline::start_jobs( + ResultQueue& result_queue, bool flush, const std::vector& busy_counts +) { std::size_t started = 0; - for (std::size_t group_index = 0; auto& group : _groups) { - for (std::size_t context_index = 0; context_index < _contexts.size(); - ++context_index) { - auto& buffers = group.buffers.at(context_index); - auto& sizes = group.buffer_sizes.at(context_index); - for (std::size_t stage = 1; stage <= buffers.size(); ++stage) { - auto& buffer = buffers.at(stage - 1); - auto& size = sizes.at(stage - 1); - if (buffer.empty() || (!flush && size < batch_size(context_index))) { - continue; + for (std::size_t context_index = 0; context_index < _contexts.size(); + ++context_index) { + std::vector> stages; + for (std::size_t group_index = 0; auto& group : _groups) { + for (std::size_t stage = 1; stage <= group.settings.size() + 1; ++stage) { + if (!group.buffers.at(context_index).at(stage - 1).empty()) { + stages.emplace_back(group_index, stage); } - std::size_t job_id = job_id_flag | _job_counter++; - auto& job = - _jobs - .emplace( - job_id, - Job{ - .group = group_index, - .stage = stage, - .context_index = context_index, - .chunks = std::move(buffer), - } - ) - .first->second; - buffer.clear(); - size = 0; - result_queue.submit( - _contexts.at(context_index)->thread_pool(), job_id, [this, &job] { - run_job(job); - } - ); - ++started; } + ++group_index; + } + std::size_t free = 0, share = 1; + if (flush && !stages.empty()) { + // with flush, the buffers share the free threads + std::size_t thread_count = + _contexts.at(context_index)->thread_pool().thread_count(); + std::size_t busy = _context_job_counts.at(context_index) + + (busy_counts.empty() ? 0 : busy_counts.at(context_index)); + free = thread_count > busy ? thread_count - busy : 0; + share = std::max(1, free / stages.size()); + } + for (auto [group_index, stage] : stages) { + auto& group = _groups.at(group_index); + auto& buffer = group.buffers.at(context_index).at(stage - 1); + auto& size = group.buffer_sizes.at(context_index).at(stage - 1); + bool full = size >= batch_size(context_index); + if (!full && free == 0) { + continue; + } + std::size_t job_count = + std::clamp(std::min(share, free), 1, buffer.size()); + free -= std::min(free, job_count); + // split at the chunk boundaries into jobs of similar size + std::vector chunks; + std::size_t job_events = 0, done_events = 0; + for (std::size_t job_index = 0; auto& chunk : buffer) { + chunks.push_back(chunk); + job_events += chunk->size(); + if (job_events * (job_count - job_index) >= size - done_events || + &chunk == &buffer.back()) { + done_events += job_events; + job_events = 0; + start_job( + result_queue, + group_index, + context_index, + stage, + std::move(chunks) + ); + chunks.clear(); + ++job_index; + ++started; + } + } + buffer.clear(); + size = 0; } - ++group_index; } return started; } +void WeightingPipeline::start_job( + ResultQueue& result_queue, + std::size_t group_index, + std::size_t context_index, + std::size_t stage, + std::vector chunks +) { + std::size_t job_id = job_id_flag | _job_counter++; + auto& job = + _jobs + .emplace( + job_id, + Job{ + .group = group_index, + .stage = stage, + .context_index = context_index, + .chunks = std::move(chunks), + } + ) + .first->second; + ++_context_job_counts.at(context_index); + result_queue.submit( + _contexts.at(context_index)->thread_pool(), job_id, [this, &job] { + run_job(job); + } + ); +} + void WeightingPipeline::gather_events(Chunk& chunk, ContextRuntimes& runtimes) { auto& source = *chunk.source; std::size_t count = chunk.size(); @@ -703,6 +757,7 @@ void WeightingPipeline::run_job(Job& job) { void WeightingPipeline::commit_job(std::size_t job_id) { auto node = _jobs.extract(job_id); auto& job = node.mapped(); + --_context_job_counts.at(job.context_index); for (auto& chunk : job.chunks) { set_ready(chunk, job.stage); } @@ -710,6 +765,7 @@ void WeightingPipeline::commit_job(std::size_t job_id) { void WeightingPipeline::reset() { _jobs.clear(); + std::fill(_context_job_counts.begin(), _context_job_counts.end(), 0); for (auto& group : _groups) { for (auto& buffers : group.buffers) { for (auto& buffer : buffers) { From 6d451dc59a990bba358a53b990f4e7a1d2f0f2fd Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 20:54:54 +0100 Subject: [PATCH 86/88] madspace: one gather per source and one concatenation per weighting-step job --- .../madspace/driver/weighting_pipeline.hpp | 16 ++- madspace/src/driver/weighting_pipeline.cpp | 132 ++++++++++++------ 2 files changed, 107 insertions(+), 41 deletions(-) diff --git a/madspace/include/madspace/driver/weighting_pipeline.hpp b/madspace/include/madspace/driver/weighting_pipeline.hpp index 6b3ab9f345..da476a08f3 100644 --- a/madspace/include/madspace/driver/weighting_pipeline.hpp +++ b/madspace/include/madspace/driver/weighting_pipeline.hpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -125,10 +126,12 @@ class WeightingPipeline { struct ContextRuntimes { RuntimePtr gather; - RuntimePtr concat; + // by the number of concatenated parts, built when first needed + std::map concats; std::vector steps; }; struct Group { + NamedVector types; std::vector settings; std::vector runtimes; // [context][stage - 1], stage step_count + 1 copies the events to the CPU @@ -157,6 +160,16 @@ class WeightingPipeline { ); void run_job(Job& job); void gather_events(Chunk& chunk, ContextRuntimes& runtimes); + // the rows of `chunks` from their shared `source`, concatenated + TensorVec gather_rows( + const TensorVec& source, + const std::vector& chunks, + std::size_t context_index, + ContextRuntimes& runtimes + ); + RuntimePtr& concat_runtime( + std::size_t group_index, std::size_t context_index, std::size_t part_count + ); std::size_t batch_size(std::size_t context_index) const; std::vector _contexts; @@ -171,6 +184,7 @@ class WeightingPipeline { std::vector>> _ready; std::unordered_map _jobs; std::vector _context_job_counts; + std::mutex _concat_mutex; std::size_t _job_counter = 0; std::size_t _chunk_count = 0; Sink _sink; diff --git a/madspace/src/driver/weighting_pipeline.cpp b/madspace/src/driver/weighting_pipeline.cpp index 2fed7bcc6a..cd82c5994b 100644 --- a/madspace/src/driver/weighting_pipeline.cpp +++ b/madspace/src/driver/weighting_pipeline.cpp @@ -47,9 +47,9 @@ Function gather_function(const NamedVector& types) { return fb.function(); } -Function concat_function(const NamedVector& types) { +Function concat_function(const NamedVector& types, std::size_t part_count) { NamedVector arg_types; - for (std::size_t part = 1; part <= 2; ++part) { + for (std::size_t part = 1; part <= part_count; ++part) { for (auto [key, type] : zip(types.keys(), types)) { arg_types.push_back(std::format("arg{}_{}", part, key), type); } @@ -57,7 +57,11 @@ Function concat_function(const NamedVector& types) { FunctionBuilder fb(arg_types, types); int count = static_cast(types.size()); for (int i = 0; i < count; ++i) { - auto [cat, sizes] = fb.batch_cat({fb.input(i), fb.input(i + count)}); + ValueVec parts; + for (std::size_t part = 0; part < part_count; ++part) { + parts.push_back(fb.input(i + static_cast(part) * count)); + } + auto [cat, sizes] = fb.batch_cat(parts); fb.output(i, cat); } return fb.function(); @@ -359,17 +363,16 @@ WeightingPipeline::WeightingPipeline( } Function gather = gather_function(types); - Function concat = concat_function(types); for (auto& context : contexts) { ContextRuntimes runtimes{ .gather = build_runtime(gather, context, false), - .concat = build_runtime(concat, context, false), }; for (auto& function : step_functions) { runtimes.steps.push_back(build_runtime(function, context, false)); } group.runtimes.push_back(std::move(runtimes)); } + group.types = types; group.buffers.assign( contexts.size(), std::vector>(step_count + 1) ); @@ -661,37 +664,64 @@ void WeightingPipeline::gather_events(Chunk& chunk, ContextRuntimes& runtimes) { if (!chunk.keep && chunk.offset == 0 && source.at(0).size(0) == count) { chunk.events = source; } else { - Tensor indices(DataType::dt_int, {count}); - auto view = indices.view(); - for (std::size_t i = 0; i < count; ++i) { - view[i] = static_cast( - chunk.offset + (chunk.keep ? chunk.keep->at(i) : i) - ); - } - auto device = _contexts.at(chunk.context_index)->device(); - if (device != cpu_device()) { - Tensor device_indices(DataType::dt_int, {count}, device); - device_indices.copy_from(indices); - indices = device_indices; - } - TensorVec args{indices}; - args.insert(args.end(), source.begin(), source.end()); - chunk.events = runtimes.gather->run(args); + chunk.events = gather_rows(source, {&chunk}, chunk.context_index, runtimes); } chunk.source.reset(); chunk.keep.reset(); chunk.offset = 0; } +TensorVec WeightingPipeline::gather_rows( + const TensorVec& source, + const std::vector& chunks, + std::size_t context_index, + ContextRuntimes& runtimes +) { + std::size_t count = 0; + for (auto chunk : chunks) { + count += chunk->size(); + } + Tensor indices(DataType::dt_int, {count}); + auto view = indices.view(); + for (std::size_t row = 0; auto chunk : chunks) { + for (std::size_t i = 0; i < chunk->size(); ++i, ++row) { + view[row] = static_cast( + chunk->offset + (chunk->keep ? chunk->keep->at(i) : i) + ); + } + } + auto device = _contexts.at(context_index)->device(); + if (device != cpu_device()) { + Tensor device_indices(DataType::dt_int, {count}, device); + device_indices.copy_from(indices); + indices = device_indices; + } + TensorVec args{indices}; + args.insert(args.end(), source.begin(), source.end()); + return runtimes.gather->run(args); +} + +RuntimePtr& WeightingPipeline::concat_runtime( + std::size_t group_index, std::size_t context_index, std::size_t part_count +) { + std::lock_guard lock(_concat_mutex); + auto& group = _groups.at(group_index); + auto& runtime = group.runtimes.at(context_index).concats[part_count]; + if (!runtime) { + runtime = build_runtime( + concat_function(group.types, part_count), _contexts.at(context_index), false + ); + } + return runtime; +} + void WeightingPipeline::run_job(Job& job) { auto& group = _groups.at(job.group); auto& runtimes = group.runtimes.at(job.context_index); - for (auto& chunk : job.chunks) { - gather_events(*chunk, runtimes); - } if (job.stage > group.settings.size()) { // after the last step, the events go to the CPU to be written for (auto& chunk : job.chunks) { + gather_events(*chunk, runtimes); for (auto& tensor : chunk->events) { chunk->cpu_events.push_back(tensor.cpu()); } @@ -700,26 +730,48 @@ void WeightingPipeline::run_job(Job& job) { return; } - std::vector parts; + // one gather per source, then one concatenation of the sources + std::vector> source_chunks; for (auto& chunk : job.chunks) { - parts.push_back(std::move(chunk->events)); - chunk->events.clear(); - } - // concatenate pairwise, so every event is copied log(chunks) times - while (parts.size() > 1) { - std::vector next_parts; - for (std::size_t i = 0; i < parts.size(); i += 2) { - if (i + 1 == parts.size()) { - next_parts.push_back(std::move(parts.at(i))); - } else { - TensorVec args = parts.at(i); - args.insert(args.end(), parts.at(i + 1).begin(), parts.at(i + 1).end()); - next_parts.push_back(runtimes.concat->run(args)); + auto found = + std::find_if(source_chunks.begin(), source_chunks.end(), [&](auto& chunks) { + return chunks.at(0)->source == chunk->source; + }); + if (found == source_chunks.end()) { + source_chunks.push_back({chunk}); + } else { + found->push_back(chunk); + } + } + job.chunks.clear(); + TensorVec events, concat_args; + for (auto& chunks : source_chunks) { + auto& first = *chunks.at(0); + auto& source = *first.source; + TensorVec part; + if (chunks.size() == 1 && !first.keep && first.offset == 0 && + source.at(0).size(0) == first.size()) { + part = source; + } else { + std::vector chunk_ptrs; + for (auto& chunk : chunks) { + chunk_ptrs.push_back(chunk.get()); } + part = gather_rows(source, chunk_ptrs, job.context_index, runtimes); + } + concat_args.insert(concat_args.end(), part.begin(), part.end()); + for (auto& chunk : chunks) { + chunk->source.reset(); + chunk->keep.reset(); + chunk->offset = 0; + // the chunks in the order of the events + job.chunks.push_back(chunk); } - parts = std::move(next_parts); } - TensorVec events = std::move(parts.at(0)); + events = source_chunks.size() == 1 + ? std::move(concat_args) + : concat_runtime(job.group, job.context_index, source_chunks.size()) + ->run(concat_args); TensorVec step_args{ events.at(group.momenta), events.at(group.alpha_s), events.at(group.flavor) From f3c73bcfbad0ae11ca158e18d40b03318d747e74 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 21:03:42 +0100 Subject: [PATCH 87/88] madspace: skip the weighting steps of survey iterations that are cleared --- madspace/src/driver/event_generator.cpp | 48 +++++++++++++++++++++---- madspace/tests/test_weighting_steps.py | 21 +++++++++++ 2 files changed, 62 insertions(+), 7 deletions(-) diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index 84e6e60f4b..2ce20a971f 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -291,15 +291,27 @@ void EventGenerator::survey(std::size_t survey_pass) { } print_survey_init(); + auto iterates = [&](auto& channel, std::size_t iter) { + return channel->status().iterations <= iter && + !(iter >= min_iters && + channel->abs_cross_section().rel_error() < target_precision); + }; + // the events for the steps, until it is known whether the channel iterates + // again and clears them + std::vector> deferred_jobs(_channels.size()); + auto release_deferred = [&](std::size_t channel_index, bool keep) { + auto& jobs = deferred_jobs.at(channel_index); + if (keep) { + for (auto& job : jobs) { + _pipeline->add_events(job); + } + } + jobs.clear(); + }; std::size_t iter = 0; for (; !done && iter < max_iters; ++iter) { for (std::size_t i = 0; auto channel : _channels) { - if (channel->status().iterations > iter) { - ++i; - continue; - } - if (iter >= min_iters && - channel->abs_cross_section().rel_error() < target_precision) { + if (!iterates(channel, iter)) { ++i; continue; } @@ -378,8 +390,12 @@ void EventGenerator::survey(std::size_t survey_pass) { order.pop_front(); auto& uw_job = _running_jobs.at(commit_id); auto& channel = _channels.at(uw_job.channel_index); + bool defer = false; if (_pipeline && _pipeline->contains(uw_job.channel_index)) { - _pipeline->add_events(uw_job); + defer = iter + 1 < max_iters; + if (!defer) { + _pipeline->add_events(uw_job); + } } else { channel->write_events( uw_job.unweighted_events, uw_job.max_weight @@ -398,7 +414,16 @@ void EventGenerator::survey(std::size_t survey_pass) { channel->optimize_vegas(uw_job); done_event_count += uw_job.batch_event_count; } + std::size_t channel_index = uw_job.channel_index; + if (defer) { + deferred_jobs.at(channel_index).push_back(std::move(uw_job)); + } _running_jobs.erase(commit_id); + // otherwise it depends on the other channels, see below + if (defer && channel_job_count == 0 && + (!done || !iterates(channel, iter + 1))) { + release_deferred(channel_index, !iterates(channel, iter + 1)); + } } } std::size_t job_id_refill = _job_id; @@ -407,6 +432,15 @@ void EventGenerator::survey(std::size_t survey_pass) { start_pipeline_jobs(); print_survey_update(false, done_event_count, total_event_count, iter); } + for (std::size_t channel_index = 0; channel_index < _channels.size(); + ++channel_index) { + release_deferred( + channel_index, + done || iter + 1 >= max_iters || + !iterates(_channels.at(channel_index), iter + 1) + ); + } + update_counts(); drain_pipeline(); } print_survey_update(true, done_event_count, total_event_count, iter - 1); diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py index d242c0467b..ddf7a1a44f 100644 --- a/madspace/tests/test_weighting_steps.py +++ b/madspace/tests/test_weighting_steps.py @@ -247,6 +247,27 @@ def test_reproducible_step_freeze(tmp_path): assert status_a.mean == status_b.mean +def test_survey_skips_cleared_iterations(tmp_path): + """Only the middle channel iterates again in the survey, so the steps of its + earlier iteration are skipped, independently of the threads.""" + results = [ + generate( + str(tmp_path / str(threads)), + [surrogate_me, heavy_tailed_me, surrogate_me], + ["reweight", "unweight"], + threads=threads, + pt_mins=(20.0, 40.0, 60.0), + grouped=(0, 1, 2), + config_kwargs={"survey_max_iters": 6, "survey_target_precision": 0.02}, + ) + for threads in (1, 4) + ] + (status_a, events_a, channels_a), (status_b, events_b, _) = results + assert [c.status().iterations for c in channels_a] == [3, 4, 3] + assert np.array_equal(events_a, events_b) + assert status_a.mean == status_b.mean + + def test_step_budget_of_smaller_target(tmp_path): """A smaller target applies its truncation budget to the step maxima too, and rescales the unweighted count accordingly. EventGenerator does this at From 90701244fa0d96592d9408a1e561f224dfdc1e63 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 29 Oct 2026 21:21:56 +0100 Subject: [PATCH 88/88] mg7: no automatic MLM merging for outputs with [amplicol] processes --- madgraph/iolibs/export_cpp.py | 6 ++++-- madgraph/various/banner.py | 8 +++++--- tests/unit_tests/various/test_banner.py | 9 +++++++++ 3 files changed, 18 insertions(+), 5 deletions(-) diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index 6862f48397..dad9e27728 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3629,12 +3629,14 @@ def create_run_card(self, matrix_elements, history): for me in matrix_elements['matrix_elements']] except (KeyError, TypeError): processes = None - if self.amplicol_processes(): + amplicol = bool(self.amplicol_processes()) + if amplicol: processes = (processes or []) + [self.amplicol_processes()] if processes: run_card.create_default_for_process(self.proc_characteristic, - history, processes) + history, processes, + amplicol=amplicol) # persist the model so the runtime can reload it: to compute the # widths set to 'auto' in the param_card (and recompute them at # each scan point) and to reset the parameters the model derives diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index b5738fd238..192df15a2c 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -8340,9 +8340,11 @@ def write_gridpack_card(self, output_file): # multiplicities, madevent's default_mlm_xqcut = 30.0 - def create_default_for_process(self, proc_characteristic, history, proc_def): + def create_default_for_process(self, proc_characteristic, history, proc_def, + amplicol=False): """Adjust default values depending on the process, then read the - site/user default file (``input/default_run_card_mg7.toml``).""" + site/user default file (``input/default_run_card_mg7.toml``). + ``amplicol``: the output has [amplicol] processes, which MLM does not support.""" # collect the initial-state (beam) particle ids beam_id = set() @@ -8390,7 +8392,7 @@ def create_default_for_process(self, proc_characteristic, history, proc_def): # drjl = 0). The shower then does the jet matching (is_mlm_merged), and # auto_ptj_mjj ties the jet cuts to xqcut. if beam_id & hadronic and not self['beam']['leptonic'] and \ - mlm_matching_needed(proc_def): + not amplicol and mlm_matching_needed(proc_def): self['beam']['dynamical_scale_choice'] = 'mlm' self['phasespace']['xqcut'] = self.default_mlm_xqcut cuts = self.dynamic_sections['cuts'] diff --git a/tests/unit_tests/various/test_banner.py b/tests/unit_tests/various/test_banner.py index d03e2e7ef7..b242304205 100755 --- a/tests/unit_tests/various/test_banner.py +++ b/tests/unit_tests/various/test_banner.py @@ -2578,6 +2578,15 @@ def test_extra_jets_turn_mlm_on(self): self.assertNotIn('jet-lepton-delta_r', rc.dynamic_sections['cuts']) self.assertIn('lepton-delta_r', rc.dynamic_sections['cuts']) + def test_amplicol_output_is_no_merging(self): + rc = bannermod.RunCardMG7() + rc.create_default_for_process( + self.PC(), '', [[mg7_proc([2, -1], [-11, 12])], + [mg7_proc([2, -1], [-11, 12, 21])]], amplicol=True) + self.assertFalse(rc.is_mlm_merged()) + self.assertEqual(rc['phasespace']['xqcut'], 0.0) + self.assertIn('jet-delta_r', rc.dynamic_sections['cuts']) + def test_an_extra_photon_is_no_merging(self): rc = self.build([[mg7_proc([2, -1], [-11, 12])], [mg7_proc([2, -1], [-11, 12, 22])]])