From 304a2bd6df317f2d297832da3a4946e0dae21120 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 9 Oct 2026 11:03:04 +0200 Subject: [PATCH 01/37] add option to create tar.gz for gridpacks, false by default for now --- madgraph/iolibs/template_files/mg7/launch.py | 3 +++ madgraph/iolibs/template_files/mg7/run_card.toml | 1 + madgraph/various/banner.py | 1 + tests/input_files/mg7_run_card.toml | 1 + 4 files changed, 6 insertions(+) diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index 2aa49beede..4ed46a9845 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -1952,6 +1952,9 @@ def save_gridpack(self) -> None: 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 diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 32024933d9..eb7cd6271a 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -57,6 +57,7 @@ save_gridpack = %(gridpack.save_gridpack)s include_source = %(gridpack.include_source)s include_madspace = %(gridpack.include_madspace)s include_madspace_source = %(gridpack.include_madspace_source)s +compress = %(gridpack.compress)s [beam] e_cm = %(beam.e_cm)s diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 24694a672b..b9c058c03c 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6758,6 +6758,7 @@ def default_setup(self): self.add_toml_param('gridpack', 'include_source', False) self.add_toml_param('gridpack', 'include_madspace', True) self.add_toml_param('gridpack', 'include_madspace_source', False) + self.add_toml_param('gridpack', 'compress', False) # ----------------------------- [beam] ------------------------- self.add_toml_param('beam', 'e_cm', 13000.0) diff --git a/tests/input_files/mg7_run_card.toml b/tests/input_files/mg7_run_card.toml index 9e4290477c..d1c543d0f3 100644 --- a/tests/input_files/mg7_run_card.toml +++ b/tests/input_files/mg7_run_card.toml @@ -23,6 +23,7 @@ save_gridpack = false include_source = false include_madspace = true include_madspace_source = false +compress = false [beam] e_cm = 1000.0 From 0e5eeffc75bedb4c83b644e6fb5e9e93974af0a8 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 9 Oct 2026 11:05:58 +0200 Subject: [PATCH 02/37] ignore include folder when copying madspace build into gridpack --- madgraph/iolibs/template_files/mg7/launch.py | 1 + 1 file changed, 1 insertion(+) diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index 4ed46a9845..39cb2b363f 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -1869,6 +1869,7 @@ def save_gridpack(self) -> None: shutil.copytree( _INSTALL_DIR / "madspace", os.path.join(gridpack_path, "madspace", "install", "madspace"), + ignore=shutil.ignore_patterns("include"), ) matrix_elements = [] From 8bf2d3cfaf4a1a59001f42e330e3cda425fec06c Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 9 Oct 2026 11:15:05 +0200 Subject: [PATCH 03/37] options to write temporary and final files to different target directories --- .../iolibs/template_files/mg7/gridpack.py | 71 ++++++++++++++----- madgraph/iolibs/template_files/mg7/launch.py | 55 +++++++++----- .../iolibs/template_files/mg7/run_card.toml | 5 ++ madgraph/various/banner.py | 3 + tests/input_files/mg7_run_card.toml | 2 + 5 files changed, 99 insertions(+), 37 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 07ae921be6..67016f1540 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -86,6 +86,10 @@ def write_lhe_header(run_path: str, meta) -> None: del writer # closes the file (writes the closing tag) +# directory the gridpack was called from; relative command line paths refer to it +_INVOCATION_DIR = os.getcwd() + + 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 @@ -99,6 +103,7 @@ def main() -> None: run_card = tomllib.load(f) 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) @@ -120,6 +125,17 @@ def main() -> None: "seed each run instead of fixing one here (still recorded in the " "run's info.json)" ) + parser.add_argument( + "--output_dir", type=str, default=None, + help="directory for the final output files instead of a new Events/ " + "folder (default: the run card value; relative paths refer to the " + "current directory)" + ) + parser.add_argument( + "--temp_output_dir", type=str, default=None, + help="directory for the temporary npy files (default: the run card " + "value, else the output directory)" + ) parser.add_argument("--device", type=str, nargs="*") parser.add_argument( "--cpu_thread_pool_size", type=int, default=run_args["cpu_thread_pool_size"] @@ -147,23 +163,40 @@ def main() -> None: args = parser.parse_args() seed = resolve_seed(args.seed) - # initialize event directory - run_name = args.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: - run_path = f"{run_dir_prefix}{run_index:02d}" - os.mkdir(run_path) - break - except FileExistsError: - run_index += 1 + # initialize output directories; command line paths are relative to the + # invocation directory, run card paths to the gridpack + def resolve_dir(cli_value, card_value): + if cli_value: + return os.path.abspath( + os.path.join(_INVOCATION_DIR, os.path.expanduser(cli_value))) + return os.path.abspath(os.path.expanduser(card_value)) if card_value else None + + output_dir = resolve_dir(args.output_dir, grid_args.get("output_dir")) + temp_dir = resolve_dir(args.temp_output_dir, grid_args.get("temp_output_dir")) + if output_dir is not None: + run_path = output_dir + os.makedirs(run_path, exist_ok=True) + else: + run_name = args.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: + run_path = f"{run_dir_prefix}{run_index:02d}" + os.mkdir(run_path) + break + except FileExistsError: + run_index += 1 + if temp_dir is None: + temp_dir = run_path + else: + os.makedirs(temp_dir, exist_ok=True) # initialize context device_names = args.device if args.device else run_args["device"] @@ -220,8 +253,8 @@ def main() -> None: ms.ChannelEventGenerator.load( os.path.join("data", "channels", file), contexts, - event_file=os.path.join(run_path, f"events.{name}.npy"), - weight_file=os.path.join(run_path, f"weights.{name}.npy"), + 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, file in madspace_data["channels"].items() diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index 39cb2b363f..2025e11c5a 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -409,22 +409,34 @@ def init_backend(self) -> None: ms.set_simd_vector_size(self.run_card["run"]["simd_vector_size"]) def init_event_dir(self) -> None: - 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 + 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 @@ -1241,8 +1253,8 @@ def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenera channel.event_generator = ms.ChannelEventGenerator( contexts=self.contexts, integrand=integrand, - event_file=os.path.join(self.run_path, f"events.{i}.{channel.name}.npy"), - weight_file=os.path.join(self.run_path, f"weights.{i}.{channel.name}.npy"), + 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}", @@ -1913,12 +1925,19 @@ def save_gridpack(self) -> None: # (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) diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index eb7cd6271a..1d8df94d54 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -58,6 +58,11 @@ include_source = %(gridpack.include_source)s include_madspace = %(gridpack.include_madspace)s include_madspace_source = %(gridpack.include_madspace_source)s compress = %(gridpack.compress)s +# gridpack runs only: directories for the temporary npy files and for the final +# output files. "" uses the run folder Events/_; if output_dir +# is set, no run folder is created. +temp_output_dir = %(gridpack.temp_output_dir)s +output_dir = %(gridpack.output_dir)s [beam] e_cm = %(beam.e_cm)s diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index b9c058c03c..05ec1e09db 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6759,6 +6759,9 @@ def default_setup(self): self.add_toml_param('gridpack', 'include_madspace', True) self.add_toml_param('gridpack', 'include_madspace_source', False) self.add_toml_param('gridpack', 'compress', False) + # used when running the gridpack; "" keeps the default Events/ folder + self.add_toml_param('gridpack', 'temp_output_dir', "", gridpack=True) + self.add_toml_param('gridpack', 'output_dir', "", gridpack=True) # ----------------------------- [beam] ------------------------- self.add_toml_param('beam', 'e_cm', 13000.0) diff --git a/tests/input_files/mg7_run_card.toml b/tests/input_files/mg7_run_card.toml index d1c543d0f3..9bcaec12fe 100644 --- a/tests/input_files/mg7_run_card.toml +++ b/tests/input_files/mg7_run_card.toml @@ -24,6 +24,8 @@ include_source = false include_madspace = true include_madspace_source = false compress = false +temp_output_dir = "" +output_dir = "" [beam] e_cm = 1000.0 From 51a9a8346344466df85ff5c4d926dd42a1d6a04e Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 9 Oct 2026 11:37:14 +0200 Subject: [PATCH 04/37] store globals inside tar file --- .../iolibs/template_files/mg7/gridpack.py | 4 +- madgraph/iolibs/template_files/mg7/launch.py | 2 +- madspace/CMakeLists.txt | 2 + madspace/include/madspace/driver.hpp | 3 +- madspace/include/madspace/driver/context.hpp | 9 +- madspace/include/madspace/driver/io.hpp | 13 +- madspace/include/madspace/driver/tar.hpp | 66 +++++ madspace/include/madspace/driver/tensor.hpp | 3 +- madspace/src/driver/context.cpp | 42 +-- madspace/src/driver/io.cpp | 151 ++++++++--- madspace/src/driver/tar.cpp | 253 ++++++++++++++++++ madspace/tests/test_tar.py | 105 ++++++++ 12 files changed, 579 insertions(+), 74 deletions(-) create mode 100644 madspace/include/madspace/driver/tar.hpp create mode 100644 madspace/src/driver/tar.cpp create mode 100644 madspace/tests/test_tar.py diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 67016f1540..46e4d2e2aa 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -240,9 +240,9 @@ def resolve_dir(cli_value, card_value): # 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", {}) - global_dir = os.path.join("data", "globals") + globals_path = os.path.join("data", "globals.tar") for context, backend in zip(contexts, backends): - context.load_globals(global_dir) + 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 diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index 2025e11c5a..7400dd049b 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -1851,7 +1851,7 @@ def save_gridpack(self) -> None: os.mkdir(gridpack_path) os.mkdir(data_path) os.mkdir(events_path) - self.contexts[0].save_globals(os.path.join(data_path, "globals")) + self.contexts[0].save_globals(os.path.join(data_path, "globals.tar")) channel_path = os.path.join(data_path, "channels") os.mkdir(channel_path) diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index 316453d193..c7166ea774 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -213,6 +213,7 @@ add_library( src/driver/status_file.cpp src/driver/thread_pool.cpp src/driver/io.cpp + src/driver/tar.cpp src/driver/vegas_optimizer.cpp src/driver/discrete_optimizer.cpp src/driver/backend.cpp @@ -267,6 +268,7 @@ add_library( include/madspace/driver/channel_generator.hpp include/madspace/driver/generator_data.hpp include/madspace/driver/io.hpp + include/madspace/driver/tar.hpp include/madspace/driver/vegas_optimizer.hpp include/madspace/driver/discrete_optimizer.hpp include/madspace/driver/backend.hpp diff --git a/madspace/include/madspace/driver.hpp b/madspace/include/madspace/driver.hpp index 26afff2a88..948427276c 100644 --- a/madspace/include/madspace/driver.hpp +++ b/madspace/include/madspace/driver.hpp @@ -5,6 +5,7 @@ #include "driver/context.hpp" #include "driver/discrete_optimizer.hpp" #include "driver/event_generator.hpp" +#include "driver/event_histograms.hpp" #include "driver/format.hpp" #include "driver/io.hpp" #include "driver/lhe_output.hpp" @@ -12,7 +13,7 @@ #include "driver/madnis_training.hpp" #include "driver/status_file.hpp" #include "driver/systematics.hpp" -#include "driver/event_histograms.hpp" +#include "driver/tar.hpp" #include "driver/tensor.hpp" #include "driver/thread_pool.hpp" #include "driver/vegas_optimizer.hpp" diff --git a/madspace/include/madspace/driver/context.hpp b/madspace/include/madspace/driver/context.hpp index 521cfb4493..8980847b5a 100644 --- a/madspace/include/madspace/driver/context.hpp +++ b/madspace/include/madspace/driver/context.hpp @@ -244,10 +244,11 @@ class Context { /// The matrix element previously loaded at position `index`; see @ref /// load_matrix_element. const MatrixElementApi& matrix_element(std::size_t index) const; - /// Write every global to `dir`, one `.npy` file per name. - void save_globals(const std::string& dir) const; - /// Define a global from every `.npy` file in `dir`, named after the file. - void load_globals(const std::string& dir); + /// Write every global to the tar archive `file`, as one `.npy` file per name. + void save_globals(const std::string& file) const; + /// Define a global from every `.npy` file in the tar archive `file`, named + /// after the file. + void load_globals(const std::string& file); /// The device every global and matrix element of this context lives on. DevicePtr device() { return _device; } /// The thread pool backing this context's device-local resources. diff --git a/madspace/include/madspace/driver/io.hpp b/madspace/include/madspace/driver/io.hpp index 5e631ad78d..7d7f189a9c 100644 --- a/madspace/include/madspace/driver/io.hpp +++ b/madspace/include/madspace/driver/io.hpp @@ -6,12 +6,19 @@ #include #include "madspace/driver/lhe_output.hpp" +#include "madspace/driver/tar.hpp" #include "madspace/driver/tensor.hpp" namespace madspace { +/// Read a tensor from the npy file `file`. Tensor load_tensor(const std::string& file); +/// Write `tensor` to the npy file `file`. void save_tensor(const std::string& file, Tensor tensor); +/// Read a tensor from the npy file `name` inside the tar archive `tar`. +Tensor load_tensor(TarReader& tar, const std::string& name); +/// Append `tensor` as an npy file called `name` to the tar archive `tar`. +void save_tensor(TarWriter& tar, const std::string& name, Tensor tensor); struct FieldLayout { static constexpr const char* i32 = "). - static std::vector layout( - int fields, - const std::vector& syst_weight_ids = {} - ) { + static std::vector + layout(int fields, const std::vector& syst_weight_ids = {}) { std::vector ret; if (fields & f_weight) { ret.push_back({"weight", FieldLayout::f64, 0}); diff --git a/madspace/include/madspace/driver/tar.hpp b/madspace/include/madspace/driver/tar.hpp new file mode 100644 index 0000000000..1d4ffa7453 --- /dev/null +++ b/madspace/include/madspace/driver/tar.hpp @@ -0,0 +1,66 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace madspace { + +/// Writes a POSIX ustar archive, readable by any standard tar tool. Entries are +/// regular files, added one at a time. Names longer than 100 characters are +/// split over the ustar `prefix` and `name` fields; names that do not fit even +/// then are rejected. The archive is finished by @ref close or the destructor. +class TarWriter { +public: + /// Create (or truncate) the archive at `file_name`. + explicit TarWriter(const std::string& file_name); + TarWriter(const TarWriter&) = delete; + TarWriter& operator=(const TarWriter&) = delete; + ~TarWriter(); + /// Append a file called `name` with the given contents. + void add(const std::string& name, const char* data, std::size_t size); + /// Write the end-of-archive marker and close the file. + void close(); + +private: + std::string _file_name; + std::ofstream _stream; + bool _closed; +}; + +/// Reads the regular files of a tar archive (ustar or GNU, as written by +/// standard tools). The archive is indexed when opened; entries are then read +/// by name. Entries other than regular files are skipped, and long names +/// stored in GNU or pax extension headers are not interpreted. +class TarReader { +public: + /// Open and index the archive at `file_name`. + explicit TarReader(const std::string& file_name); + /// The names of all files in the archive, in sorted order. + std::vector names() const; + /// Whether the archive contains a file called `name`. + bool contains(const std::string& name) const; + /// The size in bytes of the file called `name`. + std::size_t size(const std::string& name) const; + /// Read the file called `name` into `data`, which must hold `size(name)` + /// bytes. + void read(const std::string& name, char* data); + /// The contents of the file called `name`. + std::vector read(const std::string& name); + +private: + struct Entry { + std::uint64_t offset; + std::uint64_t size; + }; + const Entry& entry(const std::string& name) const; + + std::string _file_name; + std::ifstream _stream; + std::map _entries; +}; + +} // namespace madspace diff --git a/madspace/include/madspace/driver/tensor.hpp b/madspace/include/madspace/driver/tensor.hpp index f71daf0b35..0aacb1a3ef 100644 --- a/madspace/include/madspace/driver/tensor.hpp +++ b/madspace/include/madspace/driver/tensor.hpp @@ -733,7 +733,8 @@ class Tensor { return copy(*impl->device, hint); } - /// Whether the elements are stored without gaps in row-major order. + /// Whether the elements are stored without gaps in column-major (Fortran) + /// order, i.e. the first index varies fastest. bool is_contiguous() const { return impl->contiguous_dims == impl->shape.size(); } /// Number of leading dimensions that are contiguous. diff --git a/madspace/src/driver/context.cpp b/madspace/src/driver/context.cpp index f988b8f731..869c345fd0 100644 --- a/madspace/src/driver/context.cpp +++ b/madspace/src/driver/context.cpp @@ -1,5 +1,6 @@ #include "madspace/driver/context.hpp" +#include #include #include #include @@ -12,7 +13,9 @@ using namespace madspace; using json = nlohmann::json; namespace { -UmamiStatus umami_key_query_not_implemented(bool const**, int*) { return UMAMI_ERROR_NOT_IMPLEMENTED; } +UmamiStatus umami_key_query_not_implemented(bool const**, int*) { + return UMAMI_ERROR_NOT_IMPLEMENTED; +} thread_local std::optional current_caller_stream; } // namespace @@ -160,8 +163,7 @@ void MatrixElementApi::throw_error(const std::string& message) const { ); } -const MatrixElementApi& -Context::load_matrix_element( +const MatrixElementApi& Context::load_matrix_element( const std::string& file, const std::string& param_card, const std::unordered_map& parameters @@ -310,27 +312,31 @@ const MatrixElementApi& Context::matrix_element(std::size_t index) const { return *_matrix_elements.at(index).get(); } -void Context::save_globals(const std::string& dir) const { - namespace fs = std::filesystem; - fs::path dir_path(dir); - fs::create_directory(dir_path); - +void Context::save_globals(const std::string& file) const { + // sorted, so that the archive does not depend on the hash map order + std::vector names; + names.reserve(_globals.size()); for (auto& [name, tensor_and_grad] : _globals) { - auto& [tensor, requires_grad] = tensor_and_grad; - fs::path tensor_file = dir_path / name; - tensor_file += ".npy"; - save_tensor(tensor_file, tensor); + names.push_back(name); + } + std::sort(names.begin(), names.end()); + + TarWriter tar(file); + for (auto& name : names) { + save_tensor(tar, name + ".npy", _globals.at(name).first); } + tar.close(); } -void Context::load_globals(const std::string& dir) { - namespace fs = std::filesystem; - for (auto& file : fs::directory_iterator(dir)) { - if (file.path().extension() != ".npy") { +void Context::load_globals(const std::string& file) { + constexpr std::string_view extension = ".npy"; + TarReader tar(file); + for (auto& file_name : tar.names()) { + if (!file_name.ends_with(extension)) { continue; } - std::string name = file.path().stem(); - Tensor tensor = load_tensor(file.path()); + std::string name = file_name.substr(0, file_name.size() - extension.size()); + Tensor tensor = load_tensor(tar, file_name); Tensor global_tensor = define_global( name, tensor.dtype(), diff --git a/madspace/src/driver/io.cpp b/madspace/src/driver/io.cpp index 3adafb2db9..2c733d9280 100644 --- a/madspace/src/driver/io.cpp +++ b/madspace/src/driver/io.cpp @@ -36,19 +36,20 @@ DataType str_to_dtype(std::string dtype) { } } -json parse_header(std::fstream& file_stream) { - char magic_num[8]; - file_stream.read(magic_num, 8); - if (std::memcmp(magic_num, "\x93NUMPY\x01\x00", 8) != 0) { +constexpr std::size_t npy_preamble_size = 10; + +/// Check the magic number and version of the first 10 bytes of an npy file and +/// return the length of the header that follows. +std::size_t parse_preamble(const char* bytes) { + if (std::memcmp(bytes, "\x93NUMPY\x01\x00", 8) != 0) { throw std::runtime_error("invalid header of npy file"); } - union { - uint16_t size; - char chars[2]; - } header_size; - file_stream.read(&header_size.chars[0], 2); - std::string header(header_size.size, '\0'); - file_stream.read(header.data(), header_size.size); + // little-endian 16 bit length + return static_cast(static_cast(bytes[8])) | + (static_cast(static_cast(bytes[9])) << 8); +} + +json parse_header_text(std::string header) { header.erase( std::remove_if( header.begin(), header.end(), [](char x) { return std::isspace(x); } @@ -80,6 +81,62 @@ json parse_header(std::fstream& file_stream) { return json::parse(header); } +json parse_header(std::fstream& file_stream) { + char preamble[npy_preamble_size]; + file_stream.read(preamble, npy_preamble_size); + if (file_stream.fail()) { + throw std::runtime_error("invalid header of npy file"); + } + std::string header(parse_preamble(preamble), '\0'); + file_stream.read(header.data(), header.size()); + return parse_header_text(std::move(header)); +} + +/// The dtype and shape described by the parsed header of a tensor npy file. +std::pair parse_tensor_header(json header) { + if (!header.is_object()) { + throw std::runtime_error("Invalid header"); + } + json descr = header.at("descr"); + json fortran_order = header.at("fortran_order"); + json header_shape = header.at("shape"); + if (!descr.is_string() || !fortran_order.is_boolean() || + !fortran_order.get() || !header_shape.is_array()) { + throw std::runtime_error("Invalid file header"); + } + DataType dtype = str_to_dtype(descr); + SizeVec shape; + for (auto& size : header_shape) { + if (!size.is_number_unsigned()) { + throw std::runtime_error("Invalid file header"); + } + shape.push_back(size.get()); + } + return {dtype, shape}; +} + +/// The npy header of `tensor`: 128 bytes, padded with spaces. +std::string tensor_header(const Tensor& tensor) { + using namespace std::string_literals; + // magic, version 1.0, little-endian header length of 118 + std::string header = "\x93NUMPY\x01\x00"s; + header.push_back(static_cast(118)); + header.push_back('\0'); + header += std::format( + "{{'descr':'{}','fortran_order':True,'shape':(", dtype_to_str(tensor.dtype()) + ); + for (std::size_t size : tensor.shape()) { + header += std::format("{},", size); + } + header += ")}"; + if (header.size() > 127) { + throw std::runtime_error("Tensor has too many dimensions to save"); + } + header.resize(127, ' '); + header.push_back('\n'); + return header; +} + std::vector> full_descr(std::size_t particle_count, const DataLayout& layout) { std::vector> field_descr; @@ -154,8 +211,13 @@ std::tuple write_event_header( } file_stream.put('\n'); file_stream.seekp(8); - uint16_t header_size_short = header_size - 10; - file_stream.write(reinterpret_cast(&header_size_short), 2); + // little-endian 16 bit length + std::size_t header_length = header_size - npy_preamble_size; + char length_bytes[2] = { + static_cast(header_length & 0xff), + static_cast((header_length >> 8) & 0xff) + }; + file_stream.write(length_bytes, 2); file_stream.seekp(header_size); return {header_size, shape_pos}; } @@ -167,25 +229,7 @@ Tensor madspace::load_tensor(const std::string& file) { if (file_stream.fail()) { throw std::runtime_error(std::format("Could not open file '{}'", file)); } - json header = parse_header(file_stream); - if (!header.is_object()) { - throw std::runtime_error("Invalid header"); - } - json descr = header.at("descr"); - json fortran_order = header.at("fortran_order"); - json header_shape = header.at("shape"); - if (!descr.is_string() || !fortran_order.is_boolean() || - !fortran_order.get() || !header_shape.is_array()) { - throw std::runtime_error("Invalid file header"); - } - DataType dtype = str_to_dtype(descr); - SizeVec shape; - for (auto& size : header_shape) { - if (!size.is_number_unsigned()) { - throw std::runtime_error("Invalid file header"); - } - shape.push_back(size.get()); - } + auto [dtype, shape] = parse_tensor_header(parse_header(file_stream)); Tensor tensor(dtype, shape); file_stream.read(static_cast(tensor.data()), tensor.byte_size()); if (file_stream.fail()) { @@ -195,24 +239,45 @@ Tensor madspace::load_tensor(const std::string& file) { } void madspace::save_tensor(const std::string& file, Tensor tensor) { - using namespace std::string_literals; Tensor cpu_tensor = tensor.cpu().contiguous(); std::ofstream file_stream(file, std::ios::binary); - file_stream << "\x93NUMPY\x01\x00\x76\x00{'descr':'"s - << dtype_to_str(tensor.dtype()) << "','fortran_order':True,'shape':("; - for (std::size_t size : tensor.shape()) { - file_stream << size << ","; - } - file_stream << ")}"; - for (int i = file_stream.tellp(); i < 127; ++i) { - file_stream.put(' '); - } - file_stream.put('\n'); + std::string header = tensor_header(tensor); + file_stream.write(header.data(), header.size()); file_stream.write( static_cast(cpu_tensor.data()), cpu_tensor.byte_size() ); } +Tensor madspace::load_tensor(TarReader& tar, const std::string& name) { + std::vector data = tar.read(name); + if (data.size() < npy_preamble_size) { + throw std::runtime_error(std::format("Invalid npy file '{}'", name)); + } + std::size_t header_length = parse_preamble(data.data()); + std::size_t data_offset = npy_preamble_size + header_length; + if (data.size() < data_offset) { + throw std::runtime_error(std::format("Invalid npy file '{}'", name)); + } + auto [dtype, shape] = parse_tensor_header( + parse_header_text(std::string(data.data() + npy_preamble_size, header_length)) + ); + Tensor tensor(dtype, shape); + if (data.size() - data_offset != tensor.byte_size()) { + throw std::runtime_error(std::format("Wrong data size in npy file '{}'", name)); + } + std::memcpy(tensor.data(), data.data() + data_offset, tensor.byte_size()); + return tensor; +} + +void madspace::save_tensor(TarWriter& tar, const std::string& name, Tensor tensor) { + Tensor cpu_tensor = tensor.cpu().contiguous(); + std::string header = tensor_header(tensor); + std::vector data(header.size() + cpu_tensor.byte_size()); + std::memcpy(data.data(), header.data(), header.size()); + std::memcpy(data.data() + header.size(), cpu_tensor.data(), cpu_tensor.byte_size()); + tar.add(name, data.data(), data.size()); +} + EventFile::EventFile( const std::string& file_name, const DataLayout& layout, diff --git a/madspace/src/driver/tar.cpp b/madspace/src/driver/tar.cpp new file mode 100644 index 0000000000..d64ce7198b --- /dev/null +++ b/madspace/src/driver/tar.cpp @@ -0,0 +1,253 @@ +#include "madspace/driver/tar.hpp" + +#include +#include +#include +#include +#include + +using namespace madspace; + +namespace { + +constexpr std::size_t block_size = 512; +using Block = std::array; + +// ustar header layout +constexpr std::size_t name_pos = 0, name_len = 100; +constexpr std::size_t mode_pos = 100, mode_len = 8; +constexpr std::size_t uid_pos = 108, uid_len = 8; +constexpr std::size_t gid_pos = 116, gid_len = 8; +constexpr std::size_t size_pos = 124, size_len = 12; +constexpr std::size_t mtime_pos = 136, mtime_len = 12; +constexpr std::size_t chksum_pos = 148, chksum_len = 8; +constexpr std::size_t typeflag_pos = 156; +constexpr std::size_t magic_pos = 257, magic_len = 6; +constexpr std::size_t version_pos = 263; +constexpr std::size_t prefix_pos = 345, prefix_len = 155; + +/// Write `value` as zero-padded octal into a field of `len` bytes, ending +/// with a NUL. +void put_octal(Block& block, std::size_t pos, std::size_t len, std::uint64_t value) { + for (std::size_t i = len - 1; i-- > 0;) { + block[pos + i] = static_cast('0' + (value & 7)); + value >>= 3; + } + if (value != 0) { + throw std::runtime_error("Value too large for tar header field"); + } + block[pos + len - 1] = '\0'; +} + +/// Read an octal field, ignoring leading spaces and stopping at the first NUL +/// or space after the digits. +std::uint64_t get_octal(const Block& block, std::size_t pos, std::size_t len) { + if (static_cast(block[pos]) & 0x80) { + throw std::runtime_error("Unsupported binary number in tar header"); + } + std::uint64_t value = 0; + bool digits = false; + for (std::size_t i = 0; i < len; ++i) { + char c = block[pos + i]; + if (c == ' ' && !digits) { + continue; + } + if (c < '0' || c > '7') { + break; + } + value = (value << 3) | static_cast(c - '0'); + digits = true; + } + return value; +} + +std::uint64_t checksum(const Block& block) { + std::uint64_t sum = 0; + for (std::size_t i = 0; i < block_size; ++i) { + bool in_chksum = i >= chksum_pos && i < chksum_pos + chksum_len; + sum += in_chksum ? 32 : static_cast(block[i]); + } + return sum; +} + +bool is_zero_block(const Block& block) { + return std::all_of(block.begin(), block.end(), [](char c) { return c == '\0'; }); +} + +std::string get_string(const Block& block, std::size_t pos, std::size_t len) { + std::size_t end = 0; + while (end < len && block[pos + end] != '\0') { + ++end; + } + return std::string(&block[pos], end); +} + +/// Split `name` into the ustar prefix and name fields. +std::pair split_name(const std::string& name) { + if (name.empty()) { + throw std::invalid_argument("Empty tar entry name"); + } + if (name.size() <= name_len) { + return {"", name}; + } + // the split has to be at a '/', with the last part fitting in the name field + for (std::size_t pos = name.find('/'); pos != std::string::npos; + pos = name.find('/', pos + 1)) { + if (pos <= prefix_len && name.size() - pos - 1 <= name_len && + name.size() - pos - 1 > 0) { + return {name.substr(0, pos), name.substr(pos + 1)}; + } + } + throw std::invalid_argument(std::format("Tar entry name too long: {}", name)); +} + +} // namespace + +TarWriter::TarWriter(const std::string& file_name) : + _file_name(file_name), + _stream(file_name, std::ios::binary | std::ios::trunc), + _closed(false) { + if (!_stream) { + throw std::runtime_error(std::format("Could not open file '{}'", file_name)); + } +} + +TarWriter::~TarWriter() { + try { + close(); + } catch (...) { + } +} + +void TarWriter::add(const std::string& name, const char* data, std::size_t size) { + if (_closed) { + throw std::runtime_error("Tar archive is already closed"); + } + auto [prefix, short_name] = split_name(name); + Block header{}; + std::memcpy(&header[name_pos], short_name.data(), short_name.size()); + std::memcpy(&header[prefix_pos], prefix.data(), prefix.size()); + put_octal(header, mode_pos, mode_len, 0644); + put_octal(header, uid_pos, uid_len, 0); + put_octal(header, gid_pos, gid_len, 0); + put_octal(header, size_pos, size_len, size); + put_octal(header, mtime_pos, mtime_len, 0); + header[typeflag_pos] = '0'; + std::memcpy(&header[magic_pos], "ustar", magic_len); // includes the NUL + header[version_pos] = '0'; + header[version_pos + 1] = '0'; + // six octal digits, NUL, space + put_octal(header, chksum_pos, 7, checksum(header)); + header[chksum_pos + 7] = ' '; + + _stream.write(header.data(), block_size); + _stream.write(data, static_cast(size)); + std::size_t padding = (block_size - size % block_size) % block_size; + static const Block zeros{}; + _stream.write(zeros.data(), static_cast(padding)); + if (!_stream) { + throw std::runtime_error( + std::format("Failed to write to file '{}'", _file_name) + ); + } +} + +void TarWriter::close() { + if (_closed) { + return; + } + _closed = true; + static const Block zeros{}; + _stream.write(zeros.data(), block_size); + _stream.write(zeros.data(), block_size); + _stream.close(); + if (_stream.fail()) { + throw std::runtime_error( + std::format("Failed to write to file '{}'", _file_name) + ); + } +} + +TarReader::TarReader(const std::string& file_name) : + _file_name(file_name), _stream(file_name, std::ios::binary) { + if (!_stream) { + throw std::runtime_error(std::format("Could not open file '{}'", file_name)); + } + Block header; + std::uint64_t pos = 0; + while (true) { + _stream.read(header.data(), block_size); + if (_stream.gcount() == 0 || is_zero_block(header)) { + break; + } + if (_stream.gcount() != static_cast(block_size)) { + throw std::runtime_error("Truncated tar header"); + } + if (get_octal(header, chksum_pos, chksum_len) != checksum(header)) { + throw std::runtime_error( + std::format("Invalid tar header checksum in '{}'", file_name) + ); + } + std::uint64_t size = get_octal(header, size_pos, size_len); + char type = header[typeflag_pos]; + pos += block_size; + if (type == '0' || type == '\0') { + std::string name = get_string(header, name_pos, name_len); + // the prefix field is only defined for ustar archives + if (std::memcmp(&header[magic_pos], "ustar", 5) == 0) { + std::string prefix = get_string(header, prefix_pos, prefix_len); + if (!prefix.empty()) { + name = prefix + "/" + name; + } + } + _entries[name] = {pos, size}; + } + std::uint64_t padded_size = (size + block_size - 1) / block_size * block_size; + pos += padded_size; + _stream.seekg(static_cast(pos)); + } + _stream.clear(); +} + +std::vector TarReader::names() const { + std::vector ret; + ret.reserve(_entries.size()); + for (auto& [name, entry] : _entries) { + ret.push_back(name); + } + return ret; +} + +bool TarReader::contains(const std::string& name) const { + return _entries.contains(name); +} + +const TarReader::Entry& TarReader::entry(const std::string& name) const { + auto search = _entries.find(name); + if (search == _entries.end()) { + throw std::runtime_error( + std::format("File '{}' not found in '{}'", name, _file_name) + ); + } + return search->second; +} + +std::size_t TarReader::size(const std::string& name) const { return entry(name).size; } + +void TarReader::read(const std::string& name, char* data) { + const Entry& e = entry(name); + _stream.clear(); + _stream.seekg(static_cast(e.offset)); + _stream.read(data, static_cast(e.size)); + if (!_stream) { + throw std::runtime_error( + std::format("Failed to read '{}' from '{}'", name, _file_name) + ); + } +} + +std::vector TarReader::read(const std::string& name) { + std::vector data(size(name)); + read(name, data.data()); + return data; +} diff --git a/madspace/tests/test_tar.py b/madspace/tests/test_tar.py new file mode 100644 index 0000000000..55cd9091bc --- /dev/null +++ b/madspace/tests/test_tar.py @@ -0,0 +1,105 @@ +import subprocess +import tarfile + +import numpy as np +import pytest + +import madspace as ms + +torch = pytest.importorskip("torch") + + +def make_context(tmp_path): + ctx = ms.Context(device=ms.cpu_device(), thread_count=1) + rng = np.random.default_rng(1) + values = { + "a.float": rng.normal(size=(1, 4, 50)), + "b.float": rng.normal(size=(1, 3, 5)), + "c.int": rng.integers(-100, 100, size=(1, 2, 7)).astype(np.int32), + "d.scalar": rng.normal(size=(1, 1)), + } + for name, value in values.items(): + dtype = ms.DataType.int if value.dtype == np.int32 else ms.DataType.float + glob = ctx.define_global(name, dtype, list(value.shape[1:])) + glob.torch().copy_(torch.from_numpy(value)) + return ctx, values + + +def test_roundtrip(tmp_path): + ctx, values = make_context(tmp_path) + path = str(tmp_path / "globals.tar") + ctx.save_globals(path) + + ctx2 = ms.Context(device=ms.cpu_device(), thread_count=1) + ctx2.load_globals(path) + assert sorted(ctx2.global_names()) == sorted(values) + for name, value in values.items(): + np.testing.assert_array_equal( + ctx2.get_global(name).numpy(), ctx.get_global(name).numpy() + ) + + +def test_standard_tar(tmp_path): + ctx, values = make_context(tmp_path) + path = tmp_path / "globals.tar" + ctx.save_globals(str(path)) + + # system tar and the Python standard library can read it + listing = subprocess.run( + ["tar", "tf", str(path)], capture_output=True, text=True, check=True + ).stdout.split() + assert sorted(listing) == sorted(name + ".npy" for name in values) + with tarfile.open(path) as tar: + assert tar.getnames() == sorted(name + ".npy" for name in values) + for name in values: + tar.extract(name + ".npy", tmp_path / "out", filter="data") + + # and so can numpy, with the same values as the context holds + for name in values: + expected = ctx.get_global(name).numpy() + loaded = np.load(tmp_path / "out" / (name + ".npy")) + np.testing.assert_array_equal(loaded, expected) + + +def test_extract_with_system_tar(tmp_path): + ctx, values = make_context(tmp_path) + path = tmp_path / "globals.tar" + ctx.save_globals(str(path)) + out = tmp_path / "out" + out.mkdir() + subprocess.run(["tar", "xf", str(path), "-C", str(out)], check=True) + assert sorted(p.name for p in out.iterdir()) == sorted( + name + ".npy" for name in values + ) + + +def test_read_tar_from_other_tool(tmp_path): + ctx, values = make_context(tmp_path) + first = tmp_path / "first.tar" + ctx.save_globals(str(first)) + src = tmp_path / "src" + src.mkdir() + with tarfile.open(first) as tar: + tar.extractall(src, filter="data") + + # an archive written by Python's tarfile (GNU format) and by system tar + for fmt, name in [ + (tarfile.GNU_FORMAT, "gnu.tar"), + (tarfile.USTAR_FORMAT, "ustar.tar"), + ]: + other = tmp_path / name + with tarfile.open(other, "w", format=fmt) as tar: + for p in sorted(src.iterdir()): + tar.add(p, arcname=p.name) + ctx2 = ms.Context(device=ms.cpu_device(), thread_count=1) + ctx2.load_globals(str(other)) + for gname in values: + np.testing.assert_array_equal( + ctx2.get_global(gname).numpy(), ctx.get_global(gname).numpy() + ) + + +def test_missing_file(tmp_path): + ctx = ms.Context(device=ms.cpu_device(), thread_count=1) + with pytest.raises(RuntimeError): + ctx.load_globals(str(tmp_path / "nonexistent.tar")) From 23973a82d2de20be46993d3207954e154f67a058 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 9 Oct 2026 11:59:13 +0200 Subject: [PATCH 05/37] one json file for all channels --- .../iolibs/template_files/mg7/gridpack.py | 8 ++++--- madgraph/iolibs/template_files/mg7/launch.py | 17 +++++++------ .../madspace/driver/channel_generator.hpp | 11 +++++++++ madspace/src/driver/channel_generator.cpp | 24 +++++++++++++++++-- madspace/src/python/madspace.cpp | 19 ++++++++++++++- 5 files changed, 64 insertions(+), 15 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 46e4d2e2aa..86705bce92 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -249,15 +249,17 @@ def resolve_dir(cli_value, card_value): ) # set up generators + with open(os.path.join("data", "channels.json")) as f: + channels = json.load(f) channel_generators = [ - ms.ChannelEventGenerator.load( - os.path.join("data", "channels", file), + 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, file in madspace_data["channels"].items() + for name, channel in channels.items() ] event_generator = ms.EventGenerator( contexts=contexts, diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index 7400dd049b..cd4671e7a8 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -1853,14 +1853,14 @@ def save_gridpack(self) -> None: os.mkdir(events_path) self.contexts[0].save_globals(os.path.join(data_path, "globals.tar")) - channel_path = os.path.join(data_path, "channels") - os.mkdir(channel_path) - channel_files = {} - for channel in self.event_generator.channels(): - name = channel.status().name - file = f"channel{name}.json" - channel_files[name] = file - channel.save(os.path.join(channel_path, file)) + # 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()}" + 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"]: @@ -1948,7 +1948,6 @@ def save_gridpack(self) -> None: os.chmod(gen_events_file, 0o755) data = { - "channels": channel_files, "matrix_elements": matrix_elements, "me_parameters": me_parameters(self.run_card), "source_hash": ms.SOURCE_HASH, diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 1830817719..cc0324739c 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -39,6 +39,15 @@ class ChannelEventGenerator { const std::string& weight_file, const GeneratorConfig& config ); + /// Like @ref load, from the JSON text produced by @ref to_json instead of + /// a file. + static ChannelEventGenerator load_json( + const std::string& channel_json, + const std::vector& contexts, + const std::string& event_file, + const std::string& weight_file, + const GeneratorConfig& config + ); /** * @param contexts One context per device to run on. @@ -143,6 +152,8 @@ class ChannelEventGenerator { void write_events(const TensorVec& unweighted_events, double job_max_weight); /// 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. + std::string to_json() const; private: void apply_truncation_budget(); diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 86f471c5a2..6581e60abe 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -809,9 +809,13 @@ void ChannelEventGenerator::write_events( void ChannelEventGenerator::save(const std::string& file_name) const { std::ofstream f(file_name); + f << to_json(); +} + +std::string ChannelEventGenerator::to_json() const { json j; j = *this; - f << j.dump(); + return j.dump(); } ChannelEventGenerator ChannelEventGenerator::load( @@ -822,7 +826,23 @@ ChannelEventGenerator ChannelEventGenerator::load( const GeneratorConfig& config ) { std::ifstream f(channel_file); - json channel = json::parse(f); + if (!f) { + throw std::runtime_error(std::format("Could not open file '{}'", channel_file)); + } + std::string text( + (std::istreambuf_iterator(f)), std::istreambuf_iterator() + ); + return load_json(text, contexts, event_file, weight_file, config); +} + +ChannelEventGenerator ChannelEventGenerator::load_json( + const std::string& channel_json, + const std::vector& contexts, + const std::string& event_file, + const std::string& weight_file, + const GeneratorConfig& config +) { + json channel = json::parse(channel_json); std::optional hist_function; std::vector histograms; if (!channel.at("histogram_function").is_null()) { diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index b1fede50dc..6589823058 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1313,7 +1313,9 @@ PYBIND11_MODULE(_madspace_py, m) { &PhaseSpaceMapping::channel_count, pydoc::doc("PhaseSpaceMapping::channel_count") ) - .def("empty", &PhaseSpaceMapping::empty, pydoc::doc("PhaseSpaceMapping::empty")); + .def( + "empty", &PhaseSpaceMapping::empty, pydoc::doc("PhaseSpaceMapping::empty") + ); py::classh( m, "MultiChannelFunction", pydoc::doc("MultiChannelFunction") @@ -3079,6 +3081,16 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("config"), pydoc::doc("ChannelEventGenerator::load") ) + .def_static( + "load_json", + &ChannelEventGenerator::load_json, + py::arg("channel_json"), + py::arg("contexts"), + py::arg("event_file"), + py::arg("weight_file"), + py::arg("config"), + pydoc::doc("ChannelEventGenerator::load_json") + ) .def( py::init< const std::vector&, @@ -3109,6 +3121,11 @@ PYBIND11_MODULE(_madspace_py, m) { &ChannelEventGenerator::save, py::arg("save"), pydoc::doc("ChannelEventGenerator::save") + ) + .def( + "to_json", + &ChannelEventGenerator::to_json, + pydoc::doc("ChannelEventGenerator::to_json") ); py::classh(m, "PdfMemberSpec", pydoc::doc("PdfMemberSpec")) From 73b9b1c01ed35e9de39c2df7bb3ddf459824b820 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 9 Oct 2026 12:05:15 +0200 Subject: [PATCH 06/37] gridpack: make a madspace source hash mismatch an error, add --ignore_source_hash --- .../iolibs/template_files/mg7/gridpack.py | 27 +++++++++++++------ 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 86705bce92..848799a7d6 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -107,14 +107,6 @@ def main() -> None: 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) - if madspace_data["source_hash"] != ms.SOURCE_HASH: - print() - print( - "\033[1m\033[31mWARNING\033[39m: The madspace version is not identical " - "to the one used to generate the gridpack. This can lead to errors or " - "incorrect results\033[0m" - ) - print() # parse command line arguments parser = argparse.ArgumentParser() @@ -136,6 +128,11 @@ def main() -> None: help="directory for the temporary npy files (default: the run card " "value, else the output directory)" ) + parser.add_argument( + "--ignore_source_hash", action="store_true", + help="run even if the madspace version differs from the one used to " + "generate the gridpack (can lead to errors or incorrect results)" + ) parser.add_argument("--device", type=str, nargs="*") parser.add_argument( "--cpu_thread_pool_size", type=int, default=run_args["cpu_thread_pool_size"] @@ -161,6 +158,20 @@ def main() -> None: 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() + + 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 args.ignore_source_hash: + sys.exit( + f"\033[1m\033[31mERROR\033[39m: {message}.\n" + "Use --ignore_source_hash to run anyway.\033[0m" + ) + print() + print(f"\033[1m\033[31mWARNING\033[39m: {message}\033[0m") + print() seed = resolve_seed(args.seed) # initialize output directories; command line paths are relative to the From 5eff3d45a4a18c7ead29b6bba07ea711608163eb Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 9 Oct 2026 12:37:48 +0200 Subject: [PATCH 07/37] gridpack run_mode: minimal/fix_max_weight modes without events, fixed channel max weights and integral priors --- madgraph/iolibs/template_files/mg7/launch.py | 30 ++++++++- .../iolibs/template_files/mg7/run_card.toml | 10 +++ madgraph/various/banner.py | 2 + .../madspace/driver/channel_generator.hpp | 23 ++++++- .../madspace/driver/event_generator.hpp | 14 +++++ madspace/src/driver/channel_generator.cpp | 63 +++++++++++++++++-- madspace/src/driver/event_generator.cpp | 61 +++++++++++++++--- madspace/src/python/madspace.cpp | 17 +++++ tests/input_files/mg7_run_card.toml | 1 + 9 files changed, 206 insertions(+), 15 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index cd4671e7a8..63da39cf77 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -237,6 +237,8 @@ class HistItem: class MadgraphProcess: def __init__(self): + self.madnis_trained = False + self.lhe_completer = None self.load_cards() self.init_backend() self.init_event_dir() @@ -1461,6 +1463,7 @@ def train_madnis(self) -> None: 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 @@ -1544,7 +1547,30 @@ def update_madnis_status_multi( 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"] @@ -1856,7 +1882,7 @@ def save_gridpack(self) -> None: # 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()}" + 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: @@ -4200,6 +4226,8 @@ def run_single(switch=None) -> "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) diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 1d8df94d54..9b4d30f003 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -58,6 +58,16 @@ include_source = %(gridpack.include_source)s include_madspace = %(gridpack.include_madspace)s include_madspace_source = %(gridpack.include_madspace_source)s compress = %(gridpack.compress)s +# how the run that creates the gridpack ends (only used if save_gridpack is set): +# regular a normal run with events, then the gridpack is stored +# minimal no events: the gridpack is stored right after the MadNIS training, +# or, with VEGAS, once the grid of every channel has converged +# fix_max_weight no events: every channel runs until freeze_max_weight_after +# 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) +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 # is set, no run folder is created. diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 05ec1e09db..3be0cdb0a9 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6759,6 +6759,8 @@ def default_setup(self): self.add_toml_param('gridpack', 'include_madspace', True) self.add_toml_param('gridpack', 'include_madspace_source', False) self.add_toml_param('gridpack', 'compress', False) + self.add_toml_param('gridpack', 'run_mode', "minimal", + allowed=['regular', 'minimal', 'fix_max_weight']) # used when running the gridpack; "" keeps the default Events/ folder self.add_toml_param('gridpack', 'temp_output_dir', "", gridpack=True) self.add_toml_param('gridpack', 'output_dir', "", gridpack=True) diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index cc0324739c..884d72bd2f 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -92,6 +92,18 @@ class ChannelEventGenerator { EventFile& weight_file() { return _weight_file; } /// Current maximum-weight estimate, used for unweighting. double max_weight() const { return _max_weight; } + /// 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); + /// Set the integral of the absolute weights known from an earlier run, with + /// its standard error. Both must be positive. + void set_abs_integral_prior(double mean, double error); + /// Whether @ref set_abs_integral_prior was called. + 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; /// Current generation batch size. std::size_t batch_size() const { return _batch_size; } /// Whether a VEGAS grid or channel-weight optimization pass is still @@ -152,8 +164,12 @@ class ChannelEventGenerator { void write_events(const TensorVec& unweighted_events, double job_max_weight); /// 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. - std::string to_json() const; + /// The JSON text that @ref save writes; see @ref load_json. With + /// `include_estimates`, the current maximum weight and the integral of the + /// absolute weights are stored as well. @ref load_json then keeps the + /// maximum weight fixed and uses the integral as a prior, see + /// @ref abs_integral_estimate. + std::string to_json(bool include_estimates = false) const; private: void apply_truncation_budget(); @@ -216,6 +232,9 @@ class ChannelEventGenerator { RunningIntegral _cross_section; RunningIntegral _abs_cross_section; double _max_weight = 0.; + bool _max_weight_fixed = false; + // mean and standard error + std::optional> _abs_integral_prior; std::size_t _count_requested_opt = 0; // Monotonic per-job counters keying each job's deterministic random stream; // never reset. Separate for survey/generate so neither depends on the other. diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 7ad0aae2bb..9e3294c60f 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -73,6 +73,15 @@ class EventGenerator { /// Generate unweighted events until every channel reaches its target /// count. void generate(); + /// Run only the VEGAS and channel-weight optimization of every channel + /// that still needs it, until all are optimized. No events are kept. + void optimize(); + /// Like @ref generate, but every channel with a nonzero cross section runs + /// until it has `freeze_max_weight_after` unweighted events, regardless of + /// its share of the cross section, so that its maximum weight is final + /// (see @ref ChannelEventGenerator::max_weight). The events are not meant + /// to be combined. + void fix_max_weights(); /** * Combine every channel's generated events into a single unweighted * `.npy` file. @@ -184,6 +193,11 @@ class EventGenerator { // Salted by this counter so repeated calls don't replay the same stream. std::size_t _unweight_call_index = 0; + // What a run of the generate loop does; see run_generate(). + enum class GenerateMode { events, optimize_only, fix_max_weights }; + void run_generate(GenerateMode mode); + GenerateMode _generate_mode = GenerateMode::events; + // Scheduling context for the running survey()/generate() call, read by // start_jobs() to derive job seeds. bool _survey_job = false; diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 6581e60abe..ea1d4f6595 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -638,7 +638,9 @@ std::size_t ChannelEventGenerator::next_vegas_batch_size() { void ChannelEventGenerator::clear_events() { _status.count_unweighted = 0; - _max_weight = 0; + if (!_max_weight_fixed) { + _max_weight = 0; + } _unweighted_count = 0; _unweighted_accept_count = 0; _status.count_opt = 0; @@ -656,7 +658,8 @@ void ChannelEventGenerator::clear_events() { } void ChannelEventGenerator::update_max_weight(Tensor weights) { - if (_status.count_unweighted > _config.freeze_max_weight_after) { + if (_max_weight_fixed || + _status.count_unweighted > _config.freeze_max_weight_after) { return; } @@ -686,7 +689,7 @@ 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. - if (target_count < old_target && _max_weight != 0 && + if (target_count < old_target && _max_weight != 0 && !_max_weight_fixed && _status.count_unweighted <= _config.freeze_max_weight_after) { apply_truncation_budget(); } @@ -812,9 +815,16 @@ void ChannelEventGenerator::save(const std::string& file_name) const { f << to_json(); } -std::string ChannelEventGenerator::to_json() const { +std::string ChannelEventGenerator::to_json(bool include_estimates) const { json j; 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}}; + } + } return j.dump(); } @@ -858,7 +868,7 @@ ChannelEventGenerator ChannelEventGenerator::load_json( } } - return ChannelEventGenerator( + ChannelEventGenerator generator( contexts, channel.at("particle_count").get(), channel.at("integrand_channel_function").get(), @@ -873,6 +883,49 @@ ChannelEventGenerator ChannelEventGenerator::load_json( config, histograms ); + if (channel.contains("max_weight") && channel.at("max_weight").get() > 0) { + generator.set_fixed_max_weight(channel.at("max_weight").get()); + } + if (channel.contains("abs_integral")) { + generator.set_abs_integral_prior( + channel.at("abs_integral").at("mean").get(), + channel.at("abs_integral").at("error").get() + ); + } + return generator; +} + +void ChannelEventGenerator::set_fixed_max_weight(double max_weight) { + if (!(max_weight > 0)) { + throw std::invalid_argument("The fixed maximum weight must be positive"); + } + _max_weight = max_weight; + _max_weight_fixed = true; + _unweighted_count = 0; + _unweighted_accept_count = 0; +} + +void ChannelEventGenerator::set_abs_integral_prior(double mean, double error) { + if (!(mean > 0) || !(error > 0) || !std::isfinite(mean) || !std::isfinite(error)) { + throw std::invalid_argument("The integral prior must be positive and finite"); + } + _abs_integral_prior = {mean, error}; +} + +double ChannelEventGenerator::abs_integral_estimate() const { + double mean = _abs_cross_section.mean(); + if (!_abs_integral_prior) { + return mean; + } + auto [prior_mean, prior_error] = *_abs_integral_prior; + double error = _abs_cross_section.error(); + // too few samples for a variance + if (_abs_cross_section.count() < 2 || !std::isfinite(error) || !(error > 0)) { + return prior_mean; + } + double prior_weight = 1. / (prior_error * prior_error); + double weight = 1. / (error * error); + return (prior_mean * prior_weight + mean * weight) / (prior_weight + weight); } void madspace::to_json(nlohmann::json& j, const ChannelEventGenerator& channel) { diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index c6e6a61bcf..fae44fdae9 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -106,7 +106,14 @@ void EventGenerator::finish_channel_job(const GeneratorBatchJob& job) { // Jobs run on the pool, but commit strictly in per-channel dispatch order, so // results are a pure function of the seed. -void EventGenerator::generate() { +void EventGenerator::generate() { run_generate(GenerateMode::events); } + +void EventGenerator::optimize() { run_generate(GenerateMode::optimize_only); } + +void EventGenerator::fix_max_weights() { run_generate(GenerateMode::fix_max_weights); } + +void EventGenerator::run_generate(GenerateMode mode) { + _generate_mode = mode; _survey_job = false; reset_start_time(); print_gen_init(); @@ -116,11 +123,22 @@ void EventGenerator::generate() { // so round one truncates at all; _max_weight only rises, and set_target_count() // tightens it when the real (smaller) target arrives. for (auto& channel : _channels) { - if (channel->status().count_opt == 0) { - channel->set_target_count(_config.target_count); + 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 + : _config.target_count + ); } } + // channels loaded with a known integral can share the events from the start + if (std::any_of(_channels.begin(), _channels.end(), [](auto& channel) { + return channel->has_abs_integral_prior(); + })) { + update_integral_fractions(); + } + while (true) { _abort_check_function(); @@ -128,8 +146,13 @@ void EventGenerator::generate() { ++channel_index) { auto& channel = _channels.at(channel_index); double integral_frac = _channel_integral_fractions.at(channel_index); - if (integral_frac > 0 && - channel->status().count_unweighted >= channel->status().count_target) { + if (mode == GenerateMode::optimize_only) { + if (!channel->needs_optimization()) { + continue; + } + } else if (integral_frac > 0 && + channel->status().count_unweighted >= + channel->status().count_target) { continue; } if (channel->needs_optimization()) { @@ -191,6 +214,15 @@ void EventGenerator::generate() { round_in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched; } + if (mode == GenerateMode::optimize_only) { + if (std::none_of(_channels.begin(), _channels.end(), [](auto& channel) { + return channel->needs_optimization(); + })) { + break; + } + continue; + } + update_integral_fractions(); update_counts(); @@ -572,9 +604,20 @@ void EventGenerator::update_integral_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(); + } + } for (auto [channel, integral_fraction] : zip(_channels, _channel_integral_fractions)) { - integral_fraction = channel->abs_cross_section().mean() / _status.mean_abs; + integral_fraction = channel->abs_integral_estimate() / total_estimate; } // Distribute events between channels, ensure sum is exactly target_count @@ -601,7 +644,11 @@ void EventGenerator::update_integral_fractions() { for (std::size_t index : order | std::views::take(leftover)) { ++counts.at(index); } - for (auto [channel, count] : zip(_channels, counts)) { + 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; + } channel->set_target_count(count); } } diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 6589823058..6e64487fd0 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -3125,7 +3125,14 @@ PYBIND11_MODULE(_madspace_py, m) { .def( "to_json", &ChannelEventGenerator::to_json, + py::arg("include_estimates") = false, pydoc::doc("ChannelEventGenerator::to_json") + ) + .def( + "set_fixed_max_weight", + &ChannelEventGenerator::set_fixed_max_weight, + py::arg("max_weight"), + pydoc::doc("ChannelEventGenerator::set_fixed_max_weight") ); py::classh(m, "PdfMemberSpec", pydoc::doc("PdfMemberSpec")) @@ -3633,6 +3640,16 @@ PYBIND11_MODULE(_madspace_py, m) { &EventGenerator::generate, pydoc::doc("EventGenerator::generate") ) + .def( + "optimize", + &EventGenerator::optimize, + pydoc::doc("EventGenerator::optimize") + ) + .def( + "fix_max_weights", + &EventGenerator::fix_max_weights, + pydoc::doc("EventGenerator::fix_max_weights") + ) .def( "combine_to_compact_npy", &EventGenerator::combine_to_compact_npy, diff --git a/tests/input_files/mg7_run_card.toml b/tests/input_files/mg7_run_card.toml index 9bcaec12fe..8a450a4fdb 100644 --- a/tests/input_files/mg7_run_card.toml +++ b/tests/input_files/mg7_run_card.toml @@ -24,6 +24,7 @@ include_source = false include_madspace = true include_madspace_source = false compress = false +run_mode = "minimal" temp_output_dir = "" output_dir = "" From 1182c38d9d220d240ad034aaf42135bfdb9aaee9 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 9 Oct 2026 13:22:53 +0200 Subject: [PATCH 08/37] logging improvements --- .../madspace/driver/event_generator.hpp | 2 + .../madspace/driver/generator_data.hpp | 4 + madspace/src/driver/channel_generator.cpp | 1 + madspace/src/driver/event_generator.cpp | 212 ++++++++++++++---- madspace/src/driver/generator_data.cpp | 1 + madspace/src/python/madspace.cpp | 5 + 6 files changed, 176 insertions(+), 49 deletions(-) diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 9e3294c60f..53f4e7f10e 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -196,6 +196,8 @@ class EventGenerator { // What a run of the generate loop does; see run_generate(). enum class GenerateMode { events, optimize_only, fix_max_weights }; void run_generate(GenerateMode mode); + // Channels with an optimizer, and how many of them are optimized. + std::pair optimized_channel_count() const; GenerateMode _generate_mode = GenerateMode::events; // Scheduling context for the running survey()/generate() call, read by diff --git a/madspace/include/madspace/driver/generator_data.hpp b/madspace/include/madspace/driver/generator_data.hpp index ee2c4e87e9..5f734d0f7d 100644 --- a/madspace/include/madspace/driver/generator_data.hpp +++ b/madspace/include/madspace/driver/generator_data.hpp @@ -192,6 +192,10 @@ struct GeneratorStatus { std::size_t count_target; /// Number of survey iterations run. std::size_t iterations; + /// Number of optimization iterations in a row without enough improvement; + /// the channel counts as optimized when this reaches + /// @ref GeneratorConfig::optimization_patience. + std::size_t iters_without_improvement; /// Whether the survey/optimization phase has finished. bool optimized; /// Whether generation of this subprocess is complete. diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index ea1d4f6595..b8bb79ee77 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -414,6 +414,7 @@ void ChannelEventGenerator::optimize_vegas(const GeneratorBatchJob& job) { } } _best_rsd = std::min(rsd, _best_rsd); + _status.iters_without_improvement = _iters_without_improvement; ++_status.iterations; } diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index fae44fdae9..de3a2a83c5 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -667,6 +667,13 @@ void EventGenerator::update_counts() { } _status.count_unweighted = total_eff_count; _status.done = done; + if (_generate_mode == GenerateMode::fix_max_weights) { + // every channel runs to its own target, whatever the config says + _status.count_target = 0; + for (auto& channel : _channels) { + _status.count_target += channel->status().count_target; + } + } } void EventGenerator::combine_to_compact_npy( @@ -1419,11 +1426,26 @@ void EventGenerator::print_survey_update_log( } } +std::pair EventGenerator::optimized_channel_count() const { + std::size_t optimizable = 0, optimized = 0; + for (auto& channel : _channels) { + // a channel without an optimizer is never marked optimized + if (channel->status().optimized) { + ++optimizable; + ++optimized; + } else if (channel->needs_optimization()) { + ++optimizable; + } + } + return {optimized, optimizable}; +} + void EventGenerator::print_gen_init() { init_status("generate"); _last_print_time = std::chrono::steady_clock::now(); + bool optimize_only = _generate_mode == GenerateMode::optimize_only; if (_config.verbosity != Verbosity::pretty) { - Logger::info("generating started"); + Logger::info(optimize_only ? "optimizing started" : "generating started"); return; } @@ -1432,27 +1454,55 @@ void EventGenerator::print_gen_init() { _pretty_box_lower = PrettyBox( "Individual channels", _channels.size() < 21 ? _channels.size() + 1 : 22, - {6, 16, 9, 9, 7, 6, 0} + optimize_only ? std::vector{6, 16, 9, 7, 6, 0} + : std::vector{6, 16, 9, 9, 7, 6, 0} ); _pretty_box_lower.set_row( - 0, {"#", "integral ↓", "RSD", "uweff", "N", "opt", "unweighted"} + 0, + optimize_only + ? std::vector< + std::string>{"#", "integral ↓", "RSD", "N", "iter", "patience"} + : std::vector{ + "#", "integral ↓", "RSD", "uweff", "N", "opt", "unweighted" + } ); if (_channels.size() > 20) { _pretty_box_lower.set_cell(21, 0, ".."); } offset = _pretty_box_lower.line_count(); } - _pretty_box_upper = PrettyBox("Integration and unweighting", 7, {19, 0}, offset); - _pretty_box_upper.set_column( - 0, - {"Result:", - "Rel. error:", - "Rel. stddev:", - "Number of events:", - "Unweighting eff.:", - "Unweighted events:", - "Run time:"} - ); + if (optimize_only) { + _pretty_box_upper = + PrettyBox("Integration and optimization", 6, {19, 0}, offset); + _pretty_box_upper.set_column( + 0, + {"Result:", + "Rel. error:", + "Rel. stddev:", + "Number of events:", + "Optimized channels:", + "Run time:"} + ); + } else { + _pretty_box_upper = PrettyBox( + _generate_mode == GenerateMode::fix_max_weights + ? "Integration and max. weight estimation" + : "Integration and unweighting", + 7, + {19, 0}, + offset + ); + _pretty_box_upper.set_column( + 0, + {"Result:", + "Rel. error:", + "Rel. stddev:", + "Number of events:", + "Unweighting eff.:", + "Unweighted events:", + "Run time:"} + ); + } _pretty_box_upper.print_first(); if (_channels.size() > 1) { _pretty_box_lower.print_first(); @@ -1478,6 +1528,7 @@ void EventGenerator::print_gen_update_pretty(bool done) { return; } _last_print_time = now; + bool optimize_only = _generate_mode == GenerateMode::optimize_only; std::string int_str, rel_str, rsd_str, uweff_str, count_str; count_str = std::format( @@ -1498,28 +1549,43 @@ void EventGenerator::print_gen_update_pretty(bool done) { _status.count_unweighted / _status.count_after_cuts_opt ); } - std::string unw_str = std::format( - "{} / {}", - format_si_prefix(_status.count_unweighted), - format_si_prefix(_status.count_target) - ); + // the progress of the run: unweighted events, or optimized channels + std::string progress_str; + double progress_fraction; + if (optimize_only) { + auto [optimized, optimizable] = optimized_channel_count(); + progress_str = std::format("{} / {}", optimized, optimizable); + progress_fraction = + optimizable > 0 ? static_cast(optimized) / optimizable : 1.; + } else { + progress_str = std::format( + "{} / {}", + format_si_prefix(_status.count_unweighted), + format_si_prefix(_status.count_target) + ); + progress_fraction = _status.count_unweighted / _status.count_target; + } std::string time_str; if (done) { auto [wall_time_sec, cpu_time_sec] = _timing_data.at("generate"); time_str = format_run_time(wall_time_sec, cpu_time_sec); } else { - unw_str = std::format( - "{:<15} {}", - unw_str, - format_progress(_status.count_unweighted / _status.count_target, 52) + progress_str = std::format( + "{:<15} {}", progress_str, format_progress(progress_fraction, 52) ); time_str = std::format( "{:%H:%M:%S}", std::chrono::round(now - _start_time) ); } - _pretty_box_upper.set_column( - 1, {int_str, rel_str, rsd_str, count_str, uweff_str, unw_str, time_str} - ); + if (optimize_only) { + _pretty_box_upper.set_column( + 1, {int_str, rel_str, rsd_str, count_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} + ); + } _pretty_box_upper.print_update(); if (_channels.size() > 1) { @@ -1530,7 +1596,7 @@ void EventGenerator::print_gen_update_pretty(bool done) { for (std::size_t row = 1; auto& channel : channels | std::views::take(20)) { std::string index_str = std::format("{}", channel.name); - std::string int_str, rsd_str, count_str, unw_str, opt_str; + std::string int_str, rsd_str, count_str, unw_str, opt_str, patience_str; if (!std::isnan(channel.error)) { int_str = format_with_error(channel.mean, channel.error); rsd_str = std::format("{:.3f}", channel.rel_std_dev); @@ -1555,11 +1621,37 @@ void EventGenerator::print_gen_update_pretty(bool done) { ); } unw_str = std::format("{:<14} {:<19}", unw_count_str, progress); + std::size_t patience = _config.optimization_patience; + // an optimized channel stays so, even if its counter was reset later + std::size_t waited = + channel.optimized ? patience : channel.iters_without_improvement; + patience_str = std::format( + "{:>3} / {:<3} {}", + waited, + patience, + format_progress( + static_cast(waited) / + std::max(patience, 1), + 19 + ) + ); + } + if (optimize_only) { + _pretty_box_lower.set_row( + row, {index_str, int_str, rsd_str, count_str, opt_str, patience_str} + ); + } else { + _pretty_box_lower.set_row( + row, + {index_str, + int_str, + rsd_str, + uweff_str, + count_str, + opt_str, + unw_str} + ); } - _pretty_box_lower.set_row( - row, - {index_str, int_str, rsd_str, uweff_str, count_str, opt_str, unw_str} - ); ++row; } _pretty_box_lower.print_update(); @@ -1573,33 +1665,55 @@ void EventGenerator::print_gen_update_log(bool done) { return; } _last_print_time = now; + bool optimize_only = _generate_mode == GenerateMode::optimize_only; std::string rel_str = std::abs(_status.error) < std::abs(_status.mean) ? std::format("{:.4f} %", _status.error / _status.mean * 100) : ""; - 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}", - format_si_prefix(_status.count_unweighted), - format_si_prefix(_status.count_target), - format_with_error(_status.mean, _status.error), - rel_str, - _status.rel_std_dev, - format_si_prefix(_status.count), - format_si_prefix(_status.count_after_cuts), - _status.count_unweighted / _status.count_opt, - _status.count_unweighted / _status.count_after_cuts_opt, - std::chrono::round(now - _start_time) - ) - ); + if (optimize_only) { + auto [optimized, optimizable] = optimized_channel_count(); + Logger::info( + std::format( + "optimizing, optimized channels: {} / {}, integral: {}, " + "rel. error: {}, RSD: {:.3f}, samps: {}, samps. after cuts: {}, " + "time: {:%H:%M:%S}", + optimized, + optimizable, + format_with_error(_status.mean, _status.error), + rel_str, + _status.rel_std_dev, + format_si_prefix(_status.count), + format_si_prefix(_status.count_after_cuts), + std::chrono::round(now - _start_time) + ) + ); + } else { + 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}", + format_si_prefix(_status.count_unweighted), + format_si_prefix(_status.count_target), + format_with_error(_status.mean, _status.error), + rel_str, + _status.rel_std_dev, + format_si_prefix(_status.count), + format_si_prefix(_status.count_after_cuts), + _status.count_unweighted / _status.count_opt, + _status.count_unweighted / _status.count_after_cuts_opt, + std::chrono::round(now - _start_time) + ) + ); + } if (done) { auto [wall_time_sec, cpu_time_sec] = _timing_data.at("generate"); Logger::info( std::format( - "generating done, {}", format_run_time(wall_time_sec, cpu_time_sec) + "{} done, {}", + optimize_only ? "optimizing" : "generating", + format_run_time(wall_time_sec, cpu_time_sec) ) ); } diff --git a/madspace/src/driver/generator_data.cpp b/madspace/src/driver/generator_data.cpp index 579bfd4f9a..3c1a0da788 100644 --- a/madspace/src/driver/generator_data.cpp +++ b/madspace/src/driver/generator_data.cpp @@ -73,6 +73,7 @@ void madspace::to_json(nlohmann::json& j, const GeneratorStatus& status) { {"count_unweighted", status.count_unweighted}, {"count_target", status.count_target}, {"iterations", status.iterations}, + {"iters_without_improvement", status.iters_without_improvement}, {"optimized", status.optimized}, {"done", status.done}, }; diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 6e64487fd0..89af1a5b34 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -2638,6 +2638,11 @@ PYBIND11_MODULE(_madspace_py, m) { &GeneratorStatus::iterations, pydoc::doc("GeneratorStatus::iterations") ) + .def_readwrite( + "iters_without_improvement", + &GeneratorStatus::iters_without_improvement, + pydoc::doc("GeneratorStatus::iters_without_improvement") + ) .def_readwrite( "optimized", &GeneratorStatus::optimized, From aac711480bbb02e3d40f029a917c024d136bcc17 Mon Sep 17 00:00:00 2001 From: Daniele Massaro Date: Fri, 9 Oct 2026 18:09:40 +0200 Subject: [PATCH 09/37] Copy backend folder when full gridpack source is asked After splitting the backends code, the subfolder "backend" was added to the process folder. This folder should be copied when the full source gridpack is asked. --- madgraph/iolibs/template_files/mg7/launch.py | 1 + 1 file changed, 1 insertion(+) diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index 63da39cf77..4acb1af30b 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -1893,6 +1893,7 @@ def save_gridpack(self) -> None: 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) From 55924ce408cf0369ed2e0e943ae8a03ecd46322a Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 9 Oct 2026 18:34:41 +0200 Subject: [PATCH 10/37] 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 16716253cf4fd6b4711638251a18ded5cd9f9a0c Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sat, 10 Oct 2026 12:42:55 +0200 Subject: [PATCH 11/37] update docs --- docs/source/gridpacks.rst | 56 ++++++++++++++++--- docs/source/madspace/cpp-api.rst | 2 + docs/source/run_card.rst | 24 ++++++++ .../iolibs/template_files/mg7/run_card.toml | 15 +---- madgraph/various/banner.py | 6 +- 5 files changed, 80 insertions(+), 23 deletions(-) diff --git a/docs/source/gridpacks.rst b/docs/source/gridpacks.rst index 84b11fb660..07b0e8b030 100644 --- a/docs/source/gridpacks.rst +++ b/docs/source/gridpacks.rst @@ -16,8 +16,13 @@ start a run: set save_gridpack true launch -The gridpack is written to ``Events/_NN/gridpack`` after the run. It is a plain -directory. Pack it with ``tar`` to move it to another machine. +The gridpack is written to ``Events/_NN/gridpack``, or to ``gridpack`` in the +``output_dir`` of the ``[gridpack]`` section if it is set. It is a plain directory. With +``compress = true``, it is packed into ``gridpack.tar.gz`` instead, ready to be moved to +another machine. + +By default, the run stops once the gridpack is ready and does not generate events. See +:ref:`gridpack-run-mode` for the other options. A gridpack contains: @@ -33,7 +38,7 @@ A gridpack contains: The trained phase-space channels and the information needed to write LHE files and the scale and PDF weights. -``lib/`` or ``src/`` and ``SubProcesses/`` +``lib/`` or ``src/``, ``SubProcesses/`` and ``backend/`` The matrix elements, see :ref:`gridpack-matrix-elements`. ``madspace/`` @@ -42,6 +47,33 @@ A gridpack contains: ``Events/`` Empty. The output of each gridpack run goes here. +.. _gridpack-run-mode: + +Ending the run +^^^^^^^^^^^^^^ + +``run_mode`` in the ``[gridpack]`` section sets how the run that creates the gridpack +ends: + +``minimal`` (default) + No events are generated. The gridpack is stored right after the MadNIS training or, + with VEGAS, once the grid of every channel has converged. This is the fastest way + to create a gridpack. Each gridpack run then estimates the maximum weights and the + channel integrals on its own. + +``fix_max_weight`` + No events are generated. Each channel runs until ``freeze_max_weight_after`` + unweighted events are reached. Its maximum weight is stored in the gridpack and kept + fixed in every gridpack run, so all runs unweight against the same maximum. The + channel integrals are stored as well and combined with the ones of the gridpack run + to share the events between the channels from the start. Creating the gridpack takes + longer, especially for channels with a low unweighting efficiency, but the gridpack + runs do not have to estimate the maximum weights first. This helps if each run only + generates a few events. + +``regular`` + The events of the run are generated as usual, then the gridpack is stored. + Generating events ----------------- @@ -51,7 +83,10 @@ Run the script from anywhere: ./bin/generate_events --events 10000 --seed 42 -Each call creates a new directory ``Events/_NN``. A gridpack generates LHE files +Each call creates a new directory ``Events/_NN``. Use ``--output_dir`` to write +the output to a given directory instead, and ``--temp_output_dir`` to put the large +temporary npy files of the channels somewhere else, for example on a fast local disk of a +cluster node. Relative paths on the command line refer to the current directory. A gridpack generates LHE files by default, since it usually feeds a shower or detector simulation. The file is compressed to ``events.lhe.gz``. Use ``--output_format`` to get one of the npy formats instead. @@ -108,6 +143,9 @@ gridpack is created. * - - ``cut_efficiency_threshold``, ``max_cut_repetitions`` - Card only. + * - ``[gridpack]`` + - ``output_dir``, ``temp_output_dir`` + - ``--output_dir``, ``--temp_output_dir`` * - ``[systematics]`` - ``enable``, ``mur``, ``muf``, ``together``, ``dynamical_scale``, ``pdf``, ``write_inputs`` @@ -115,6 +153,8 @@ gridpack is created. An option on the command line overrides the card. The settings have the same meaning as in the main run card, see :doc:`run_card`. ``./bin/generate_events --help`` lists all options. +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. 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 @@ -143,7 +183,8 @@ Matrix elements that you used. A gridpack can only run on devices for which a library exists. ``include_source = true`` - The source code is copied to ``src/`` and ``SubProcesses/`` and ``lib/`` stays empty. + The source code is copied to ``src/``, ``SubProcesses/`` and ``backend/`` and ``lib/`` + stays empty. This is useful if the gridpack has to run on a different platform. The gridpack does not compile anything by itself, so build the libraries once before the first run: @@ -179,8 +220,9 @@ If both options are on, the compiled copy is used and the source is only a fallb rebuilding it. The compiled copy has to match the platform. The source code makes the gridpack portable. -The script warns if the ``madspace`` it uses differs from the one that created the gridpack. -In this case, the results can be wrong. +The script stops with an error if the ``madspace`` it uses differs from the one that +created the gridpack, since the results can be wrong in this case. Use +``--ignore_source_hash`` to run anyway, which only prints a warning. The options can be combined. The smallest gridpack sets ``include_madspace = false`` and relies on an installed ``madspace``. A fully portable gridpack sets ``include_source = true`` diff --git a/docs/source/madspace/cpp-api.rst b/docs/source/madspace/cpp-api.rst index 64384364ef..8aad6b8723 100644 --- a/docs/source/madspace/cpp-api.rst +++ b/docs/source/madspace/cpp-api.rst @@ -159,6 +159,8 @@ The high-level driver plus the tensor and device layer it runs on: integration, cpp/SubprocessSystArgs cpp/SystematicsCalculator cpp/SystematicsConfig + cpp/TarReader + cpp/TarWriter cpp/Tensor cpp/TensorView cpp/ThreadPool diff --git a/docs/source/run_card.rst b/docs/source/run_card.rst index 9893e64731..20ef869dce 100644 --- a/docs/source/run_card.rst +++ b/docs/source/run_card.rst @@ -112,6 +112,30 @@ See :doc:`gridpacks` for details. ``include_madspace_source`` (``false``) Include the ``madspace`` source code. +``compress`` (``false``) + Pack the gridpack into ``gridpack.tar.gz`` instead of leaving it as a directory. + +``run_mode`` (``"minimal"``) + How the run that creates the gridpack ends. Only used if ``save_gridpack`` is set. + + - ``regular``: generate the events of the run, then store the gridpack. + - ``minimal``: no events. The gridpack is stored right after the MadNIS training or, + with VEGAS, once the grid of every channel has converged. + - ``fix_max_weight``: no events. Each channel runs until ``freeze_max_weight_after`` + unweighted events are reached. Its maximum weight is then fixed in the gridpack, + and its integral is used to share the events between the channels. This can be + slow for channels with a low unweighting efficiency. + + See :ref:`gridpack-run-mode`. + +``temp_output_dir`` (``""``) + Directory for the temporary npy files of the channels. ``""`` uses the output + directory. + +``output_dir`` (``""``) + Directory for the output files. ``""`` creates a new run folder + ``Events/_NN``. If set, no run folder is created. + [beam] ------ diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 173b938cdd..d94c80f747 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -44,20 +44,11 @@ include_source = %(gridpack.include_source)s include_madspace = %(gridpack.include_madspace)s include_madspace_source = %(gridpack.include_madspace_source)s compress = %(gridpack.compress)s -# how the run that creates the gridpack ends (only used if save_gridpack is set): -# regular a normal run with events, then the gridpack is stored -# minimal no events: the gridpack is stored right after the MadNIS training, -# or, with VEGAS, once the grid of every channel has converged -# fix_max_weight no events: every channel runs until freeze_max_weight_after -# 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) +# end of a run creating a gridpack; options: regular, minimal (no events), fix_max_weight (no events) 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 -# is set, no run folder is created. +# temporary npy files; "" uses the output directory temp_output_dir = %(gridpack.temp_output_dir)s +# output files; "" creates a new run folder in Events output_dir = %(gridpack.output_dir)s [beam] diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 6014597e2b..86539475b9 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6755,11 +6755,9 @@ def default_setup(self): comment="how the run creating the gridpack ends: regular (with events), " "minimal or fix_max_weight (no events)") self.add_toml_param('gridpack', 'temp_output_dir', "", gridpack=True, - comment="gridpack runs only: directory for the temporary npy files; " - "\"\" uses the run folder") + comment="directory for the temporary npy files; \"\" uses the output directory") self.add_toml_param('gridpack', 'output_dir', "", gridpack=True, - comment="gridpack runs only: directory for the output files; " - "\"\" uses the run folder") + comment="directory for the output files; \"\" creates a new run folder in Events") # ----------------------------- [beam] ------------------------- self.add_toml_param('beam', 'e_cm', 13000.0, From 02181940b5f186c207678d9c123ab9fb675fd1a5 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sun, 11 Oct 2026 14:39:56 +0200 Subject: [PATCH 12/37] test_mg7_launch: add gridpack section to fake run card --- tests/unit_tests/interface/test_mg7_launch.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 fa1aaaf1b2083ebb30a43096472d78ac7d69bae4 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 12 Oct 2026 11:19:24 +0200 Subject: [PATCH 13/37] 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 14/37] 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 f4c78a86ec13fa3eebf7b3e4dd80feefd707a221 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 12 Oct 2026 13:18:04 +0200 Subject: [PATCH 15/37] madspace: UMAMI keys and generators for custom cut, scale, PDF and bias functions Add UMAMI input keys for particle count, PDG ids, momentum fraction, scale, PDG id and beam index, and output keys CUT, SCALE, PDF and BIAS. A new CustomFunction generator calls a loaded UMAMI library on the event momenta. Cuts, EnergyScale, PartonDensity (via CustomPdf) and Integrand accept these custom functions. --- madspace/CMakeLists.txt | 1 + madspace/include/madspace/phasespace.hpp | 1 + .../madspace/phasespace/custom_function.hpp | 103 +++++++++++++++++ madspace/include/madspace/phasespace/cuts.hpp | 22 +++- .../include/madspace/phasespace/integrand.hpp | 18 ++- .../madspace/phasespace/matrix_element.hpp | 30 +++-- madspace/include/madspace/phasespace/pdf.hpp | 20 +++- .../include/madspace/phasespace/scale.hpp | 9 +- madspace/include/madspace/umami.h | 28 ++++- madspace/src/compgraphs/instruction.cpp | 18 +++ madspace/src/phasespace/custom_function.cpp | 104 +++++++++++++++++ madspace/src/phasespace/cuts.cpp | 40 +++++-- madspace/src/phasespace/integrand.cpp | 51 ++++++++- madspace/src/phasespace/matrix_element.cpp | 85 +++++++++++--- madspace/src/phasespace/pdf.cpp | 51 +++++++++ madspace/src/phasespace/scale.cpp | 48 +++++--- madspace/src/python/madspace.cpp | 107 ++++++++++++++++-- 17 files changed, 662 insertions(+), 74 deletions(-) create mode 100644 madspace/include/madspace/phasespace/custom_function.hpp create mode 100644 madspace/src/phasespace/custom_function.cpp diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index 316453d193..d96db0a477 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -240,6 +240,7 @@ add_library( src/phasespace/madnis.cpp src/phasespace/rambo.cpp src/phasespace/cuts.cpp + src/phasespace/custom_function.cpp src/phasespace/observable.cpp src/phasespace/histograms.cpp src/phasespace/chili.cpp diff --git a/madspace/include/madspace/phasespace.hpp b/madspace/include/madspace/phasespace.hpp index ca6cb3cf8d..d4e3990962 100644 --- a/madspace/include/madspace/phasespace.hpp +++ b/madspace/include/madspace/phasespace.hpp @@ -6,6 +6,7 @@ #include "phasespace/channel_weights.hpp" #include "phasespace/color_ordered_mapping.hpp" #include "phasespace/cross_section.hpp" +#include "phasespace/custom_function.hpp" #include "phasespace/cuts.hpp" #include "phasespace/discrete_flow.hpp" #include "phasespace/discrete_sampler.hpp" diff --git a/madspace/include/madspace/phasespace/custom_function.hpp b/madspace/include/madspace/phasespace/custom_function.hpp new file mode 100644 index 0000000000..8b04abee39 --- /dev/null +++ b/madspace/include/madspace/phasespace/custom_function.hpp @@ -0,0 +1,103 @@ +#pragma once + +#include "madspace/driver/context.hpp" +#include "madspace/phasespace/base.hpp" + +namespace madspace { + +/** + * A user-supplied function of the event momenta, loaded as a UMAMI library. + * + * Custom cuts, scales and biases use the UMAMI interface as a generic batched + * function. The library is loaded like a matrix element, with + * Context::load_matrix_element. It receives the momenta + * (`UMAMI_IN_MOMENTA`), the particle count (`UMAMI_IN_PARTICLE_COUNT`) and the + * PDG ids of the external particles (`UMAMI_IN_PDG_IDS`). It returns one value + * per event, under `UMAMI_OUT_CUT`, `UMAMI_OUT_SCALE` or `UMAMI_OUT_BIAS`. + * + * With a single PDG-id option the ids are constant. With several options the + * ids are selected per event by the `flavor_index` argument. + * + * `batch` is the leading batch dimension. + * + * **Arguments** + * - `momenta` – `float`, shape `(batch, particle_count, 4)` – the event momenta. + * - `flavor_index` – `int`, shape `(batch,)` – index into the PDG-id options. + * Present only with more than one option. + * + * **Returns** + * - `cut`, `scale` or `bias` – `float`, shape `(batch,)` – the function value. + */ +class CustomFunction : public FunctionGenerator { +public: + /// What the custom function computes. + enum CustomFunctionType { + cut, ///< Acceptance factor, 0 rejects the event. + scale, ///< Energy scale. + bias ///< Event bias. + }; + + /** + * @param function_index Index of the loaded library in the context. + * @param type What the function computes. + * @param particle_count Number of external particles. + * @param pdg_id_options PDG ids of the external particles, one list per + * flavor option. Empty passes zeros. + */ + CustomFunction( + std::size_t function_index, + CustomFunctionType type, + std::size_t particle_count, + const nested_vector2& pdg_id_options = {} + ); + /** + * Take the index from an already-loaded library. + * @param function_api The loaded library. + * @param type What the function computes. + * @param particle_count Number of external particles. + * @param pdg_id_options See the other constructor. + */ + CustomFunction( + const MatrixElementApi& function_api, + CustomFunctionType type, + std::size_t particle_count, + const nested_vector2& pdg_id_options = {} + ) : + CustomFunction(function_api.index(), type, particle_count, pdg_id_options) {} + + /// Index of the loaded library in the context. + std::size_t function_index() const { return _function_index; } + /// What the function computes. + CustomFunctionType type() const { return _type; } + /// Whether the PDG ids are selected per event by a `flavor_index` argument. + bool dynamic_pdg_ids() const { return _pdg_id_options.size() > 1; } + +private: + NamedVector build_function_impl( + FunctionBuilder& fb, const NamedVector& args + ) const override; + + std::size_t _function_index; + CustomFunctionType _type; + std::size_t _particle_count; + nested_vector2 _pdg_id_options; +}; + +/** + * A user-supplied parton density, loaded as a UMAMI library. + * + * Passed in place of a @ref PdfGrid. The library receives the momentum + * fraction, the factorization scale, the parton PDG id and the beam index. It + * returns @f$x f(x, q)@f$ under `UMAMI_OUT_PDF`. + */ +struct CustomPdf { + /// @param function_index Index of the loaded library in the context. + CustomPdf(std::size_t function_index) : function_index(function_index) {} + /// @param function_api The loaded library. + CustomPdf(const MatrixElementApi& function_api) : + function_index(function_api.index()) {} + /// Index of the loaded library in the context. + std::size_t function_index; +}; + +} // namespace madspace diff --git a/madspace/include/madspace/phasespace/cuts.hpp b/madspace/include/madspace/phasespace/cuts.hpp index f3ff571ac4..1933855ed9 100644 --- a/madspace/include/madspace/phasespace/cuts.hpp +++ b/madspace/include/madspace/phasespace/cuts.hpp @@ -2,6 +2,7 @@ #include "madspace/compgraphs.hpp" #include "madspace/phasespace/base.hpp" +#include "madspace/phasespace/custom_function.hpp" #include "madspace/phasespace/observable.hpp" #include @@ -17,7 +18,9 @@ namespace madspace { * a list of @ref Cuts::CutItem "cut items" [1]. Each item evaluates one * @ref Observable for its particle selection and requires the result to lie in * `[min, max]`. When the selection yields several values, @ref CutMode decides - * whether all or any of them must pass. See Sec. 3.2.4 of [1]. + * whether all or any of them must pass. See Sec. 3.2.4 of [1]. An optional + * @ref CustomFunction of type CustomFunction::cut multiplies its value into the + * mask. * * The accessors re-expose the active lower bounds in the layout the * phase-space mappings consume. `pt_min()` and `eta_max()` are indexed by @@ -57,11 +60,19 @@ class Cuts : public FunctionGenerator { CutMode mode = CutMode::all; }; - /// @param cut_data The cut items to apply. - Cuts(const std::vector& cut_data); - /// Build a pass-through mask with no cuts. + /// @param cut_data The cut items to apply. + /// @param custom_cut Optional user-supplied cut function. + Cuts( + const std::vector& cut_data, + const std::optional& custom_cut = std::nullopt + ); + /// Build a mask with no cut items. /// @param particle_count Number of external particles. - Cuts(std::size_t particle_count); + /// @param custom_cut Optional user-supplied cut function. + Cuts( + std::size_t particle_count, + const std::optional& custom_cut = std::nullopt + ); /// Largest required partonic center-of-mass energy, or 0 if unconstrained. double sqrt_s_min() const; /// Per-outgoing-particle maximum pseudorapidity (infinity where inactive). @@ -96,6 +107,7 @@ class Cuts : public FunctionGenerator { ) const; std::vector _cut_data; + std::optional _custom_cut; }; } // namespace madspace diff --git a/madspace/include/madspace/phasespace/integrand.hpp b/madspace/include/madspace/phasespace/integrand.hpp index 0ca15b8cac..7cb088358a 100644 --- a/madspace/include/madspace/phasespace/integrand.hpp +++ b/madspace/include/madspace/phasespace/integrand.hpp @@ -3,6 +3,7 @@ #include "madspace/phasespace/channel_weight_network.hpp" #include "madspace/phasespace/channel_weights.hpp" #include "madspace/phasespace/cross_section.hpp" +#include "madspace/phasespace/custom_function.hpp" #include "madspace/phasespace/discrete_flow.hpp" #include "madspace/phasespace/discrete_sampler.hpp" #include "madspace/phasespace/flow.hpp" @@ -45,6 +46,8 @@ namespace madspace { * `float`, shape `(batch,)` – scales and momentum fractions. * - `random` – `float`, shape `(batch, mapping.random_dim())` – the underlying * random numbers. + * - `bias` – `float`, shape `(batch,)` – the event bias, present only with a + * custom bias. `weight` then includes the bias factor. * * **References** * - [1] T. Heimel, O. Mattelaer, R. Winterhalder, "MadSpace", @@ -57,6 +60,8 @@ class Integrand : public FunctionGenerator { /// The optional adaptive sampler for a set of discrete choices. using AdaptiveDiscrete = std::variant; + /// The optional parton density: an LHAPDF grid or a user-supplied library. + using PdfSource = std::variant; /// The @ref MatrixElement inputs an `Integrand` supplies. inline static const std::vector @@ -91,7 +96,8 @@ class Integrand : public FunctionGenerator { * the flavour choice. * @param pid_options Allowed initial-parton flavour * combinations. - * @param pdf_grid Optional shared @ref PdfGrid. + * @param pdf_grid Optional shared @ref PdfGrid or + * @ref CustomPdf. * @param running_coupling Optional shared * @ref RunningCoupling. * @param energy_scale Optional shared @ref EnergyScale. @@ -126,6 +132,10 @@ class Integrand : public FunctionGenerator { * @param flavor_per_subproc_remap Per-sub-process flavour remap. * @param compressed_channel_weight_count Channel-weight entries kept per * event. + * @param custom_bias Optional user-supplied event + * bias of type CustomFunction::bias. + * Its PDG-id options must match + * @p pid_options. */ Integrand( const PhaseSpaceMapping& mapping, @@ -134,7 +144,7 @@ class Integrand : public FunctionGenerator { const AdaptiveDiscrete& discrete_sym = std::monostate{}, const AdaptiveDiscrete& discrete_flavor = std::monostate{}, const nested_vector2& pid_options = {}, - const std::optional& pdf_grid = std::nullopt, + const PdfSource& pdf_grid = std::monostate{}, const std::optional& running_coupling = std::nullopt, const std::optional& energy_scale = std::nullopt, const std::optional& prop_chan_weights = std::nullopt, @@ -155,7 +165,8 @@ class Integrand : public FunctionGenerator { const std::vector& flavor_diff_xs_indices = {}, const std::vector& flavor_subproc_indices = {}, const std::vector& flavor_per_subproc_remap = {}, - std::size_t compressed_channel_weight_count = 50 + std::size_t compressed_channel_weight_count = 50, + const std::optional& custom_bias = std::nullopt ); /// Total number of particles, incoming and outgoing. std::size_t particle_count() const { return _mapping.particle_count(); } @@ -259,6 +270,7 @@ class Integrand : public FunctionGenerator { std::vector _flavor_diff_xs_indices; std::vector _flavor_subproc_indices; std::vector _flavor_per_subproc_remap; + std::optional _custom_bias; friend class IntegrandProbability; friend class IntegrandChannelPart; diff --git a/madspace/include/madspace/phasespace/matrix_element.hpp b/madspace/include/madspace/phasespace/matrix_element.hpp index 70d2c7afc8..f677765d78 100644 --- a/madspace/include/madspace/phasespace/matrix_element.hpp +++ b/madspace/include/madspace/phasespace/matrix_element.hpp @@ -40,15 +40,21 @@ class MatrixElement : public FunctionGenerator { public: /// A quantity the matrix element can be asked to consume. enum MatrixElementInput { - momenta_in, ///< Final-state momenta. - alpha_s_in, ///< Strong coupling. - flavor_in, ///< Flavour assignment. - random_color_in, ///< Random number for the color draw. - random_helicity_in, ///< Random number for the helicity draw. - random_diagram_in, ///< Random number for the diagram draw. - helicity_in, ///< Fixed helicity. - channel_in, ///< Multi-channel index. - diagram_in ///< Fixed diagram. + momenta_in, ///< Final-state momenta. + alpha_s_in, ///< Strong coupling. + flavor_in, ///< Flavour assignment. + random_color_in, ///< Random number for the color draw. + random_helicity_in, ///< Random number for the helicity draw. + random_diagram_in, ///< Random number for the diagram draw. + helicity_in, ///< Fixed helicity. + channel_in, ///< Multi-channel index. + diagram_in, ///< Fixed diagram. + particle_count_in, ///< Number of external particles (custom functions). + pdg_ids_in, ///< PDG ids of the external particles. + momentum_fraction_in, ///< Momentum fraction of an incoming parton. + scale_in, ///< Energy scale, e.g. the factorization scale. + pdg_id_in, ///< PDG id of a single parton. + beam_index_in ///< Index of the beam (0 or 1). }; /// A quantity the matrix element can be asked to produce. @@ -57,7 +63,11 @@ class MatrixElement : public FunctionGenerator { diagram_amp2_out, ///< Per-diagram squared amplitudes. color_index_out, ///< Sampled color. helicity_index_out, ///< Sampled helicity. - diagram_index_out ///< Sampled diagram. + diagram_index_out, ///< Sampled diagram. + 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. }; /** diff --git a/madspace/include/madspace/phasespace/pdf.hpp b/madspace/include/madspace/phasespace/pdf.hpp index 916fe357b1..599b9e87c1 100644 --- a/madspace/include/madspace/phasespace/pdf.hpp +++ b/madspace/include/madspace/phasespace/pdf.hpp @@ -2,6 +2,7 @@ #include "madspace/driver/context.hpp" #include "madspace/phasespace/base.hpp" +#include "madspace/phasespace/custom_function.hpp" namespace madspace { @@ -57,7 +58,8 @@ struct PdfGrid { * * Interpolates a @ref PdfGrid at the sampled @f$(x, q)@f$, returning * @f$x f(x, q)@f$ for the requested flavours; the numerical output agrees with - * LHAPDF to double precision [1] (Sec. 3.2.3 of [2]). + * LHAPDF to double precision [1] (Sec. 3.2.3 of [2]). With a @ref CustomPdf, + * the user-supplied library is called once per requested flavour instead. * * `batch` is the leading batch dimension. * @@ -92,6 +94,19 @@ class PartonDensity : public FunctionGenerator { bool dynamic_pid = false, const std::string& prefix = "" ); + /** + * @param custom_pdf The user-supplied parton density. + * @param pids PDG ids of the flavours to return. + * @param beam_index Index of the beam (0 or 1), passed to the library. + * @param dynamic_pid If true, a single flavour selected per event by the + * `flavor_index` argument is returned. + */ + PartonDensity( + const CustomPdf& custom_pdf, + const std::vector& pids, + int beam_index, + bool dynamic_pid = false + ); private: NamedVector build_function_impl( @@ -99,6 +114,9 @@ class PartonDensity : public FunctionGenerator { ) const override; std::vector _pid_indices; + std::vector _pids; + std::optional _custom_pdf; + me_int_t _beam_index = 0; bool _dynamic_pid; std::string _prefix; std::vector _logx_shape; diff --git a/madspace/include/madspace/phasespace/scale.hpp b/madspace/include/madspace/phasespace/scale.hpp index d25fe2279b..8f8abab21c 100644 --- a/madspace/include/madspace/phasespace/scale.hpp +++ b/madspace/include/madspace/phasespace/scale.hpp @@ -1,6 +1,7 @@ #pragma once #include "madspace/phasespace/base.hpp" +#include "madspace/phasespace/custom_function.hpp" namespace madspace { @@ -24,6 +25,8 @@ namespace madspace { * @f$H_\mathrm{T}/2@f$ (`half_transverse_mass`) or @f$\sqrt{\hat s}@f$ * (`partonic_energy`). The dynamical scale is multiplied by `scale_factor`, * e.g. @f$H_\mathrm{T}/4@f$ is `half_transverse_mass` with `scale_factor = 0.5`. + * A @ref CustomFunction of type CustomFunction::scale replaces the dynamical + * choice, like `user_dynamical_scale` in MadGraph5. * * `batch` is the leading batch dimension. * @@ -85,6 +88,8 @@ class EnergyScale : public FunctionGenerator { * @param fact_scale2 Fixed factorization scale for the second beam. * @param scale_factor Factor on the dynamical scale; fixed scales are * left alone. + * @param custom_scale Optional user-supplied function that replaces + * the dynamical scale choice. */ EnergyScale( std::size_t particle_count, @@ -94,7 +99,8 @@ class EnergyScale : public FunctionGenerator { double ren_scale, double fact_scale1, double fact_scale2, - double scale_factor = 1. + double scale_factor = 1., + const std::optional& custom_scale = std::nullopt ); private: @@ -109,6 +115,7 @@ class EnergyScale : public FunctionGenerator { double _fact_scale1; double _fact_scale2; double _scale_factor; + std::optional _custom_scale; }; } // namespace madspace diff --git a/madspace/include/madspace/umami.h b/madspace/include/madspace/umami.h index 57d079d1c1..a30959b2a2 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 0 +#define UMAMI_MINOR_VERSION 1 typedef enum { /** operation was executed successfully */ @@ -93,15 +93,27 @@ typedef enum { 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_CHANNEL_INDEX + 1) +#define UMAMI_INPUT_KEY_COUNT (UMAMI_IN_BEAM_INDEX + 1) typedef enum { /** value of the matrix element, type: `double`, shape: `()` */ UMAMI_OUT_MATRIX_ELEMENT, - /** selected color index, type: `double`, shape: `(diagram count)` */ + /** squared amplitudes of the diagrams, type: `double`, shape: `(diagram count)` */ UMAMI_OUT_DIAGRAM_AMP2, /** selected color index, type: `int`, shape: `()` */ UMAMI_OUT_COLOR_INDEX, @@ -112,10 +124,18 @@ typedef enum { /** 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, } UmamiOutputKey; /** Number of values in `UmamiOutputKey` */ -#define UMAMI_OUTPUT_KEY_COUNT (UMAMI_OUT_GPU_STREAM + 1) +#define UMAMI_OUTPUT_KEY_COUNT (UMAMI_OUT_BIAS + 1) /** Implementation-defined pointer to a matrix element instance */ typedef void* UmamiHandle; diff --git a/madspace/src/compgraphs/instruction.cpp b/madspace/src/compgraphs/instruction.cpp index b05d0486ae..85bc85fe3b 100644 --- a/madspace/src/compgraphs/instruction.cpp +++ b/madspace/src/compgraphs/instruction.cpp @@ -939,6 +939,8 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const { case UMAMI_IN_RANDOM_COLOR: case UMAMI_IN_RANDOM_HELICITY: case UMAMI_IN_RANDOM_DIAGRAM: + case UMAMI_IN_MOMENTUM_FRACTION: + case UMAMI_IN_SCALE: if (input_type.dtype != DataType::dt_float || input_type.shape.size() != 0) { throw std::invalid_argument( @@ -952,6 +954,9 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const { case UMAMI_IN_HELICITY_INDEX: case UMAMI_IN_DIAGRAM_INDEX: case UMAMI_IN_CHANNEL_INDEX: + case UMAMI_IN_PARTICLE_COUNT: + case UMAMI_IN_PDG_ID: + case UMAMI_IN_BEAM_INDEX: if (input_type.dtype != DataType::dt_int || input_type.shape.size() != 0) { throw std::invalid_argument( std::format( @@ -960,6 +965,15 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const { ); } break; + case UMAMI_IN_PDG_IDS: + if (input_type.dtype != DataType::dt_int || input_type.shape.size() != 1) { + throw std::invalid_argument( + std::format( + "matrix_element, argument {}: expected array of integers", i + 1 + ) + ); + } + break; default: throw std::invalid_argument( std::format( @@ -979,6 +993,10 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const { std::vector shape; switch (output_key) { case UMAMI_OUT_MATRIX_ELEMENT: + case UMAMI_OUT_CUT: + case UMAMI_OUT_SCALE: + case UMAMI_OUT_PDF: + case UMAMI_OUT_BIAS: dtype = DataType::dt_float; shape = {}; break; diff --git a/madspace/src/phasespace/custom_function.cpp b/madspace/src/phasespace/custom_function.cpp new file mode 100644 index 0000000000..44c4f5ce01 --- /dev/null +++ b/madspace/src/phasespace/custom_function.cpp @@ -0,0 +1,104 @@ +#include "madspace/phasespace/custom_function.hpp" + +#include "madspace/phasespace/matrix_element.hpp" + +#include + +using namespace madspace; + +namespace { + +const char* output_name(CustomFunction::CustomFunctionType type) { + switch (type) { + case CustomFunction::cut: + return "cut"; + case CustomFunction::scale: + return "scale"; + case CustomFunction::bias: + return "bias"; + } + throw std::invalid_argument("unknown custom function type"); +} + +MatrixElement::MatrixElementOutput output_key(CustomFunction::CustomFunctionType type) { + switch (type) { + case CustomFunction::cut: + return MatrixElement::cut_out; + case CustomFunction::scale: + return MatrixElement::scale_out; + case CustomFunction::bias: + return MatrixElement::bias_out; + } + throw std::invalid_argument("unknown custom function type"); +} + +} // namespace + +CustomFunction::CustomFunction( + std::size_t function_index, + CustomFunctionType type, + std::size_t particle_count, + const nested_vector2& pdg_id_options +) : + FunctionGenerator( + "CustomFunction", + [&] { + NamedVector arg_types{ + {"momenta", batch_four_vec_array(particle_count)} + }; + if (pdg_id_options.size() > 1) { + arg_types.push_back("flavor_index", batch_int); + } + return arg_types; + }(), + {{output_name(type), batch_float}} + ), + _function_index(function_index), + _type(type), + _particle_count(particle_count), + _pdg_id_options(pdg_id_options) { + for (auto& option : pdg_id_options) { + if (option.size() != particle_count) { + throw std::invalid_argument( + std::format( + "custom function: expected {} PDG ids per option, got {}", + particle_count, + option.size() + ) + ); + } + } +} + +NamedVector CustomFunction::build_function_impl( + FunctionBuilder& fb, const NamedVector& args +) const { + auto momenta = args.at(0); + auto batch_size = fb.batch_size({momenta}); + ValueVec pdg_ids; + for (std::size_t i = 0; i < _particle_count; ++i) { + if (_pdg_id_options.empty()) { + pdg_ids.push_back(fb.full({static_cast(0), batch_size})); + } else if (_pdg_id_options.size() == 1) { + pdg_ids.push_back(fb.full({_pdg_id_options.at(0).at(i), batch_size})); + } else { + std::vector column; + for (auto& option : _pdg_id_options) { + column.push_back(option.at(i)); + } + pdg_ids.push_back(fb.gather_int(args.at(1), column)); + } + } + MatrixElement function( + _function_index, + _particle_count, + {MatrixElement::momenta_in, + MatrixElement::particle_count_in, + MatrixElement::pdg_ids_in}, + {output_key(_type)} + ); + auto particle_count = fb.full({static_cast(_particle_count), batch_size}); + auto value = + function.build_function(fb, {momenta, particle_count, fb.stack(pdg_ids)}).at(0); + return {{output_name(_type), value}}; +} diff --git a/madspace/src/phasespace/cuts.cpp b/madspace/src/phasespace/cuts.cpp index 69795e5814..64a408f4ce 100644 --- a/madspace/src/phasespace/cuts.cpp +++ b/madspace/src/phasespace/cuts.cpp @@ -7,18 +7,42 @@ using namespace madspace; -Cuts::Cuts(const std::vector& cut_data) : +namespace { + +void check_custom_cut(const std::optional& custom_cut) { + if (custom_cut && + (custom_cut->type() != CustomFunction::cut || custom_cut->dynamic_pdg_ids())) { + throw std::invalid_argument( + "custom cut must be of type cut, with a single PDG-id option" + ); + } +} + +} // namespace + +Cuts::Cuts( + const std::vector& cut_data, + const std::optional& custom_cut +) : FunctionGenerator( "Cuts", cut_data.at(0).observable.arg_types(), {{"mask", batch_float}} ), - _cut_data(cut_data) {} + _cut_data(cut_data), + _custom_cut(custom_cut) { + check_custom_cut(custom_cut); +} -Cuts::Cuts(std::size_t particle_count) : +Cuts::Cuts( + std::size_t particle_count, const std::optional& custom_cut +) : FunctionGenerator( "Cuts", {{"momenta", batch_four_vec_array(particle_count)}}, {{"mask", batch_float}} - ) {} + ), + _custom_cut(custom_cut) { + check_custom_cut(custom_cut); +} NamedVector Cuts::build_function_impl(FunctionBuilder& fb, const NamedVector& args) const { @@ -36,6 +60,9 @@ Cuts::build_function_impl(FunctionBuilder& fb, const NamedVector& args) c weights.push_back(fb.cut_any(obs, item.min, item.max)); } } + if (_custom_cut) { + weights.push_back(_custom_cut->build_function(fb, {args.at(0)}).at(0)); + } return {{"mask", fb.product(weights)}}; } @@ -168,9 +195,8 @@ PairList delta_r_pairs(const Observable& o) { } // namespace -std::vector> Cuts::pairwise_min( - const std::function& pairs -) const { +std::vector> +Cuts::pairwise_min(const std::function& pairs) const { std::size_t n = arg_types().at(0).shape.at(0) - 2; std::vector> out(n, std::vector(n, 0.)); for (auto& item : _cut_data) { diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index 379c3bbdae..b8f5fde48e 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -41,7 +41,7 @@ Integrand::Integrand( const AdaptiveDiscrete& discrete_sym, const AdaptiveDiscrete& discrete_flavor, const nested_vector2& pid_options, - const std::optional& pdf_grid, + const PdfSource& pdf_grid, const std::optional& running_coupling, const std::optional& energy_scale, const std::optional& prop_chan_weights, @@ -62,7 +62,8 @@ Integrand::Integrand( const std::vector& flavor_diff_xs_indices, const std::vector& flavor_subproc_indices, const std::vector& flavor_per_subproc_remap, - std::size_t compressed_channel_weight_count + std::size_t compressed_channel_weight_count, + const std::optional& custom_bias ) : FunctionGenerator( "Integrand", @@ -117,7 +118,8 @@ Integrand::Integrand( if (flav_count > 1 && !std::holds_alternative(discrete_flavor)) { ret_types.push_back("discrete_flavor_index", batch_int); - if (pdf_grid && energy_scale) { + if (!std::holds_alternative(pdf_grid) && + energy_scale) { ret_types.push_back("pdf_prior", batch_float_array(flav_count)); } } @@ -149,6 +151,9 @@ Integrand::Integrand( } } ret_types.push_back("random", batch_float_array(mapping.random_dim())); + if (custom_bias) { + ret_types.push_back("bias", batch_float); + } if (mapping.channel_count() > 1 && !std::holds_alternative(discrete_sym)) { ret_types.push_back("channel_index_in_group", batch_int); @@ -198,8 +203,9 @@ Integrand::Integrand( ), _flavor_per_subproc_remap( flavor_per_subproc_remap.begin(), flavor_per_subproc_remap.end() - ) { - if (pdf_grid) { + ), + _custom_bias(custom_bias) { + if (!std::holds_alternative(pdf_grid)) { for (std::size_t i = 0; i < 2; ++i) { std::set pids; for (auto& option : pid_options) { @@ -210,7 +216,24 @@ Integrand::Integrand( std::distance(pids.begin(), pids.find(option.at(i))) ); } - _pdfs.at(i) = PartonDensity(pdf_grid.value(), {pids.begin(), pids.end()}); + std::vector pid_vec(pids.begin(), pids.end()); + if (auto grid = std::get_if(&pdf_grid)) { + _pdfs.at(i) = PartonDensity(*grid, pid_vec); + } else { + _pdfs.at(i) = PartonDensity( + std::get(pdf_grid), pid_vec, static_cast(i) + ); + } + } + } + if (custom_bias) { + if (custom_bias->type() != CustomFunction::bias) { + throw std::invalid_argument("custom bias must be of type bias"); + } + if (custom_bias->dynamic_pdg_ids() != (pid_options.size() > 1)) { + throw std::invalid_argument( + "the PDG-id options of the custom bias must match pid_options" + ); } } @@ -780,6 +803,17 @@ NamedVector Integrand::build_common_part( diff_xs_acc = fb.mul(diff_xs_acc, fb.gather(flavor_id, _flavor_factors)); } ValueVec weights_after_cuts{args.at("weight_after_cuts"), diff_xs_acc}; + // the bias enters the weight seen by the adaptation and the unweighting, + // the event generator divides it out again + Value bias_acc; + if (_custom_bias) { + ValueVec bias_args{momenta_me}; + if (_custom_bias->dynamic_pdg_ids()) { + bias_args.push_back(flavor_id); + } + bias_acc = _custom_bias->build_function(fb, bias_args).at(0); + weights_after_cuts.push_back(bias_acc); + } ValueVec extra_weights_after_cuts; if (args.index_map().contains("extra_weight_after_cuts")) { extra_weights_after_cuts.push_back(args.at("extra_weight_after_cuts")); @@ -971,6 +1005,11 @@ NamedVector Integrand::build_common_part( } } outputs.push_back("random", optional_cut(args.at("r"))); + if (bias_acc) { + outputs.push_back( + "bias", scatter_or_drop(fb.full({1., batch_size_val}), bias_acc) + ); + } if (has_permutations && !std::holds_alternative(_discrete_sym)) { outputs.push_back( diff --git a/madspace/src/phasespace/matrix_element.cpp b/madspace/src/phasespace/matrix_element.cpp index 851bfb8c09..dabc80972b 100644 --- a/madspace/src/phasespace/matrix_element.cpp +++ b/madspace/src/phasespace/matrix_element.cpp @@ -10,19 +10,20 @@ namespace { /// A per-particle 0/1 selector, 1 for every one-based particle number in /// `indices`. `what` names the setting the numbers came from, for the error. std::vector particle_mask( - std::size_t particle_count, - const std::vector& indices, - const char* what + std::size_t particle_count, const std::vector& indices, const char* what ) { std::vector mask(particle_count, 0); for (me_int_t index : indices) { if (index < 1 || std::cmp_greater(index, particle_count)) { - throw std::invalid_argument(std::format( - "{} particle {} out of range, this process has {} external particles", - what, - index, - particle_count - )); + throw std::invalid_argument( + std::format( + "{} particle {} out of range, this process has {} external " + "particles", + what, + index, + particle_count + ) + ); } mask.at(index - 1) = 1; } @@ -45,14 +46,12 @@ MatrixElement::MatrixElement( "MatrixElement", [&] { NamedVector arg_types; - bool found_momenta = false; for (auto input : inputs) { switch (input) { case momenta_in: arg_types.push_back( "momenta", batch_four_vec_array(particle_count) ); - found_momenta = true; break; case alpha_s_in: arg_types.push_back("alpha_s", batch_float); @@ -84,12 +83,30 @@ MatrixElement::MatrixElement( case channel_in: arg_types.push_back("channel", batch_int); break; + case particle_count_in: + arg_types.push_back("particle_count", batch_int); + break; + case pdg_ids_in: + arg_types.push_back("pdg_ids", batch_int_array(particle_count)); + break; + case momentum_fraction_in: + arg_types.push_back("momentum_fraction", batch_float); + break; + case scale_in: + arg_types.push_back("scale", batch_float); + break; + case pdg_id_in: + arg_types.push_back("pdg_id", batch_int); + break; + case beam_index_in: + arg_types.push_back("beam_index", batch_int); + break; default: throw std::invalid_argument("unknown input type"); } } - if (!found_momenta) { - throw std::invalid_argument("missing momentum input"); + if (arg_types.size() == 0) { + throw std::invalid_argument("matrix element needs at least one input"); } return arg_types; }(), @@ -114,6 +131,18 @@ MatrixElement::MatrixElement( case diagram_index_out: ret_types.push_back("diagram_index", batch_int); break; + case cut_out: + ret_types.push_back("cut", batch_float); + break; + case scale_out: + ret_types.push_back("scale", batch_float); + break; + case pdf_out: + ret_types.push_back("pdf", batch_float); + break; + case bias_out: + ret_types.push_back("bias", batch_float); + break; default: throw std::invalid_argument("unknown output type"); } @@ -207,6 +236,24 @@ NamedVector MatrixElement::build_function_impl( case channel_in: input_key = UMAMI_IN_CHANNEL_INDEX; break; + case particle_count_in: + input_key = UMAMI_IN_PARTICLE_COUNT; + break; + case pdg_ids_in: + input_key = UMAMI_IN_PDG_IDS; + break; + case momentum_fraction_in: + input_key = UMAMI_IN_MOMENTUM_FRACTION; + break; + case scale_in: + input_key = UMAMI_IN_SCALE; + break; + case pdg_id_in: + input_key = UMAMI_IN_PDG_ID; + break; + case beam_index_in: + input_key = UMAMI_IN_BEAM_INDEX; + break; } matrix_args.push_back(static_cast(input_key)); if (_sample_random_inputs && @@ -261,6 +308,18 @@ NamedVector MatrixElement::build_function_impl( case diagram_index_out: output_key = UMAMI_OUT_DIAGRAM_INDEX; break; + case cut_out: + output_key = UMAMI_OUT_CUT; + break; + case scale_out: + output_key = UMAMI_OUT_SCALE; + break; + case pdf_out: + output_key = UMAMI_OUT_PDF; + break; + case bias_out: + output_key = UMAMI_OUT_BIAS; + break; } matrix_args.push_back(static_cast(output_key)); matrix_args.push_back(size_arg); diff --git a/madspace/src/phasespace/pdf.cpp b/madspace/src/phasespace/pdf.cpp index 0a437f341e..fc3a28025c 100644 --- a/madspace/src/phasespace/pdf.cpp +++ b/madspace/src/phasespace/pdf.cpp @@ -1,5 +1,7 @@ #include "madspace/phasespace/pdf.hpp" +#include "madspace/phasespace/matrix_element.hpp" + #include #include #include @@ -384,11 +386,60 @@ PartonDensity::PartonDensity( } } +PartonDensity::PartonDensity( + const CustomPdf& custom_pdf, + const std::vector& pids, + int beam_index, + bool dynamic_pid +) : + FunctionGenerator( + "PartonDensity", + [&] { + NamedVector arg_types{{"x", batch_float}, {"q", batch_float}}; + if (dynamic_pid) { + arg_types.push_back("flavor_index", batch_int); + } + return arg_types; + }(), + {{"pdf", dynamic_pid ? batch_float : batch_float_array(pids.size())}} + ), + _pids(pids.begin(), pids.end()), + _custom_pdf(custom_pdf), + _beam_index(beam_index), + _dynamic_pid(dynamic_pid) {} + NamedVector PartonDensity::build_function_impl( FunctionBuilder& fb, const NamedVector& args ) const { auto x = args.at(0); auto q = args.at(1); + if (_custom_pdf) { + MatrixElement function( + _custom_pdf->function_index, + 0, + {MatrixElement::momentum_fraction_in, + MatrixElement::scale_in, + MatrixElement::pdg_id_in, + MatrixElement::beam_index_in}, + {MatrixElement::pdf_out} + ); + auto batch_size = fb.batch_size({x}); + auto beam_index = fb.full({_beam_index, batch_size}); + if (_dynamic_pid) { + auto pid = fb.gather_int(args.at(2), _pids); + return { + {"pdf", function.build_function(fb, {x, q, pid, beam_index}).at(0)} + }; + } + ValueVec pdfs; + for (me_int_t pid : _pids) { + auto pid_value = fb.full({pid, batch_size}); + pdfs.push_back( + function.build_function(fb, {x, q, pid_value, beam_index}).at(0) + ); + } + return {{"pdf", fb.stack(pdfs)}}; + } auto grid_logx = fb.global( prefixed_name(_prefix, "pdf_logx"), DataType::dt_float, diff --git a/madspace/src/phasespace/scale.cpp b/madspace/src/phasespace/scale.cpp index 46715ca972..ddcc82524d 100644 --- a/madspace/src/phasespace/scale.cpp +++ b/madspace/src/phasespace/scale.cpp @@ -10,7 +10,8 @@ EnergyScale::EnergyScale( double ren_scale, double fact_scale1, double fact_scale2, - double scale_factor + double scale_factor, + const std::optional& custom_scale ) : FunctionGenerator( "EnergyScale", @@ -25,7 +26,16 @@ EnergyScale::EnergyScale( _ren_scale(ren_scale), _fact_scale1(fact_scale1), _fact_scale2(fact_scale2), - _scale_factor(scale_factor) {} + _scale_factor(scale_factor), + _custom_scale(custom_scale) { + if (custom_scale && + (custom_scale->type() != CustomFunction::scale || + custom_scale->dynamic_pdg_ids())) { + throw std::invalid_argument( + "custom scale must be of type scale, with a single PDG-id option" + ); + } +} NamedVector EnergyScale::build_function_impl( FunctionBuilder& fb, const NamedVector& args @@ -40,21 +50,25 @@ NamedVector EnergyScale::build_function_impl( }; } Value scale; - switch (_dynamical_scale_type) { - case transverse_energy: - scale = fb.scale_transverse_energy(momenta); - break; - case transverse_mass: - scale = fb.scale_transverse_mass(momenta); - break; - case half_transverse_mass: - scale = fb.scale_half_transverse_mass(momenta); - break; - case partonic_energy: - scale = fb.scale_partonic_energy(momenta); - break; - default: - throw std::runtime_error("invalid dynamical scale type"); + if (_custom_scale) { + scale = _custom_scale->build_function(fb, {momenta}).at(0); + } else { + switch (_dynamical_scale_type) { + case transverse_energy: + scale = fb.scale_transverse_energy(momenta); + break; + case transverse_mass: + scale = fb.scale_transverse_mass(momenta); + break; + case half_transverse_mass: + scale = fb.scale_half_transverse_mass(momenta); + break; + case partonic_energy: + scale = fb.scale_partonic_energy(momenta); + break; + default: + throw std::runtime_error("invalid dynamical scale type"); + } } // dynamical scale only; a fixed scale is an absolute value if (_scale_factor != 1.) { diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 5cae93612f..9b267ec9d0 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -955,6 +955,71 @@ PYBIND11_MODULE(_madspace_py, m) { .def_readonly_static("missing_pids", &Observable::missing_pids) .def_readonly_static("photon_pids", &Observable::photon_pids); + py::classh custom_function( + m, "CustomFunction", pydoc::doc("CustomFunction") + ); + add_enum( + custom_function, + "CustomFunctionType", + { + {"cut", CustomFunction::cut}, + {"scale", CustomFunction::scale}, + {"bias", CustomFunction::bias}, + } + ); + custom_function + .def( + py::init< + std::size_t, + CustomFunction::CustomFunctionType, + std::size_t, + const nested_vector2&>(), + py::arg("function_index"), + py::arg("type"), + py::arg("particle_count"), + py::arg("pdg_id_options") = nested_vector2{}, + pydoc::doc("CustomFunction::CustomFunction") + ) + .def( + py::init< + const MatrixElementApi&, + CustomFunction::CustomFunctionType, + std::size_t, + const nested_vector2&>(), + py::arg("function_api"), + py::arg("type"), + py::arg("particle_count"), + py::arg("pdg_id_options") = nested_vector2{}, + pydoc::doc("CustomFunction::CustomFunction#2") + ) + .def( + "function_index", + &CustomFunction::function_index, + pydoc::doc("CustomFunction::function_index") + ) + .def("type", &CustomFunction::type, pydoc::doc("CustomFunction::type")) + .def( + "dynamic_pdg_ids", + &CustomFunction::dynamic_pdg_ids, + pydoc::doc("CustomFunction::dynamic_pdg_ids") + ); + py::classh(m, "CustomPdf", pydoc::doc("CustomPdf")) + .def( + py::init(), + py::arg("function_index"), + pydoc::doc("CustomPdf::CustomPdf") + ) + .def( + py::init(), + py::arg("function_api"), + pydoc::doc("CustomPdf::CustomPdf#2") + ) + .def_readonly( + "function_index", + &CustomPdf::function_index, + pydoc::doc("CustomPdf::function_index") + ); + auto cuts = py::classh(m, "Cuts", pydoc::doc("Cuts")); add_enum( cuts, @@ -991,13 +1056,17 @@ PYBIND11_MODULE(_madspace_py, m) { "min", &Cuts::PairMassAny::min, pydoc::doc("Cuts::PairMassAny::min") ); cuts.def( - py::init&>(), + py::init< + const std::vector&, + const std::optional&>(), py::arg("cut_data"), + py::arg("custom_cut") = std::nullopt, pydoc::doc("Cuts::Cuts") ) .def( - py::init(), + py::init&>(), py::arg("particle_count"), + py::arg("custom_cut") = std::nullopt, pydoc::doc("Cuts::Cuts#2") ) .def("sqrt_s_min", &Cuts::sqrt_s_min, pydoc::doc("Cuts::sqrt_s_min")) @@ -1313,7 +1382,9 @@ PYBIND11_MODULE(_madspace_py, m) { &PhaseSpaceMapping::channel_count, pydoc::doc("PhaseSpaceMapping::channel_count") ) - .def("empty", &PhaseSpaceMapping::empty, pydoc::doc("PhaseSpaceMapping::empty")); + .def( + "empty", &PhaseSpaceMapping::empty, pydoc::doc("PhaseSpaceMapping::empty") + ); py::classh( m, "MultiChannelFunction", pydoc::doc("MultiChannelFunction") @@ -1341,6 +1412,12 @@ PYBIND11_MODULE(_madspace_py, m) { {"helicity_in", MatrixElement::helicity_in}, {"diagram_in", MatrixElement::diagram_in}, {"channel_in", MatrixElement::channel_in}, + {"particle_count_in", MatrixElement::particle_count_in}, + {"pdg_ids_in", MatrixElement::pdg_ids_in}, + {"momentum_fraction_in", MatrixElement::momentum_fraction_in}, + {"scale_in", MatrixElement::scale_in}, + {"pdg_id_in", MatrixElement::pdg_id_in}, + {"beam_index_in", MatrixElement::beam_index_in}, } ); add_enum( @@ -1352,6 +1429,10 @@ PYBIND11_MODULE(_madspace_py, m) { {"color_index_out", MatrixElement::color_index_out}, {"helicity_index_out", MatrixElement::helicity_index_out}, {"diagram_index_out", MatrixElement::diagram_index_out}, + {"cut_out", MatrixElement::cut_out}, + {"scale_out", MatrixElement::scale_out}, + {"pdf_out", MatrixElement::pdf_out}, + {"bias_out", MatrixElement::bias_out}, } ); matrix_element @@ -1862,6 +1943,14 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("dynamic_pid") = false, py::arg("prefix") = "", pydoc::doc("PartonDensity::PartonDensity") + ) + .def( + py::init&, int, bool>(), + py::arg("custom_pdf"), + py::arg("pids"), + py::arg("beam_index"), + py::arg("dynamic_pid") = false, + pydoc::doc("PartonDensity::PartonDensity#2") ); py::classh(m, "AlphaSGrid", pydoc::doc("AlphaSGrid")) @@ -1951,7 +2040,8 @@ PYBIND11_MODULE(_madspace_py, m) { double, double, double, - double>(), + double, + const std::optional&>(), py::arg("particle_count"), py::arg("dynamical_scale_type"), py::arg("ren_scale_fixed"), @@ -1960,6 +2050,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("fact_scale1"), py::arg("fact_scale2"), py::arg("scale_factor") = 1., + py::arg("custom_scale") = std::nullopt, pydoc::doc("EnergyScale::EnergyScale#4") ); @@ -2051,7 +2142,7 @@ PYBIND11_MODULE(_madspace_py, m) { const Integrand::AdaptiveDiscrete&, const Integrand::AdaptiveDiscrete&, const nested_vector2&, - const std::optional&, + const Integrand::PdfSource&, const std::optional&, const std::optional&, const std::optional&, @@ -2072,14 +2163,15 @@ PYBIND11_MODULE(_madspace_py, m) { const std::vector&, const std::vector&, const std::vector&, - std::size_t>(), + std::size_t, + const std::optional&>(), py::arg("mapping"), py::arg("diff_xs"), py::arg("adaptive_map") = std::monostate{}, py::arg("discrete_sym") = std::monostate{}, py::arg("discrete_flavor") = std::monostate{}, py::arg("pid_options") = nested_vector2{}, - py::arg("pdf_grid") = std::nullopt, + py::arg("pdf_grid") = std::monostate{}, py::arg("running_coupling") = std::nullopt, py::arg("energy_scale") = std::nullopt, py::arg("prop_chan_weights") = std::nullopt, @@ -2101,6 +2193,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("flavor_subproc_indices") = std::vector{}, py::arg("flavor_per_subproc_remap") = std::vector{}, py::arg("compressed_channel_weight_count") = 50, + py::arg("custom_bias") = std::nullopt, pydoc::doc("Integrand::Integrand") ) .def( From a0422fdd7fa498fcbb0a24b0cbc360f58b4cfefe Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 12 Oct 2026 13:22:53 +0200 Subject: [PATCH 16/37] mg7: bin/init_custom_func and [custom_functions] run card section bin/init_custom_func writes custom_/ with a UMAMI wrapper, a function body for one phase-space point (C, C++, Fortran, CUDA or HIP), the UMAMI header and a Makefile that builds lib/custom_.so. Options are taken from the command line or asked interactively. The run card gets a [custom_functions] section with the library paths. --- madgraph/iolibs/export_cpp.py | 12 + .../mg7/custom_functions/Makefile | 93 +++++++ .../mg7/custom_functions/bias.c | 29 +++ .../mg7/custom_functions/bias.cpp | 33 +++ .../mg7/custom_functions/bias.cu | 28 +++ .../mg7/custom_functions/bias.f90 | 28 +++ .../template_files/mg7/custom_functions/cut.c | 27 +++ .../mg7/custom_functions/cut.cpp | 31 +++ .../mg7/custom_functions/cut.cu | 26 ++ .../mg7/custom_functions/cut.f90 | 30 +++ .../template_files/mg7/custom_functions/pdf.c | 27 +++ .../mg7/custom_functions/pdf.cpp | 36 +++ .../mg7/custom_functions/pdf.cu | 26 ++ .../mg7/custom_functions/pdf.f90 | 32 +++ .../mg7/custom_functions/scale.c | 36 +++ .../mg7/custom_functions/scale.cpp | 40 ++++ .../mg7/custom_functions/scale.cu | 35 +++ .../mg7/custom_functions/scale.f90 | 35 +++ .../mg7/custom_functions/wrapper_cpu.c | 195 +++++++++++++++ .../mg7/custom_functions/wrapper_gpu.cu | 226 ++++++++++++++++++ .../template_files/mg7/init_custom_func.py | 222 +++++++++++++++++ .../iolibs/template_files/mg7/run_card.toml | 8 + madgraph/various/banner.py | 11 + 23 files changed, 1266 insertions(+) create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/Makefile create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/bias.c create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/bias.cpp create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/bias.cu create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/bias.f90 create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/cut.c create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/cut.cpp create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/cut.cu create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/cut.f90 create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/pdf.c create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/pdf.cpp create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/pdf.cu create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/pdf.f90 create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/scale.c create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/scale.cpp create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/scale.cu create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/scale.f90 create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/wrapper_cpu.c create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/wrapper_gpu.cu create mode 100644 madgraph/iolibs/template_files/mg7/init_custom_func.py diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index 14c0756790..e70cbfe98c 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3325,6 +3325,18 @@ def copy_template(self, model): ) os.chmod(npy_bin, 0o755) + custom_bin = os.path.join("bin", "init_custom_func") + with open(custom_bin, "w") as f: + f.write( + "#! /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" + "if __name__ == '__main__':\n" + " sys.exit(main(me_dir=os.path.dirname(os.path.dirname(os.path.realpath(__file__)))))\n" + ) + os.chmod(custom_bin, 0o755) + # Recorded in Cards/me5_configuration.txt: the tools a run needs but cannot # rediscover on its own. LHAPDF above all -- bin/generate_events may be # driven from a shell that never sourced anything MadGraph-related. diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/Makefile b/madgraph/iolibs/template_files/mg7/custom_functions/Makefile new file mode 100644 index 0000000000..0435e1e525 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/Makefile @@ -0,0 +1,93 @@ +# Makefile for the custom %(kind)s function, generated by bin/init_custom_func. +# Builds $(LIB) from %(body_file)s and the UMAMI wrapper. +# Add extra compiler flags to CFLAGS and libraries to LIBS, e.g. for LHAPDF: +# make CFLAGS="$$(lhapdf-config --cflags)" LIBS="$$(lhapdf-config --libs)" + +LANGUAGE = %(language)s +LIB = ../lib/custom_%(kind)s.so +BODY = %(body_file)s + +OPTFLAGS ?= -O3 +CFLAGS ?= +LIBS ?= + +# make predefines CC, CXX and FC; only replace its f77 default +ifeq ($(origin FC),default) + FC = gfortran +endif + +UNAME_S := $(shell uname -s) +ifeq ($(UNAME_S),Darwin) + SONAME_FLAGS = -Wl,-install_name,@rpath/$(notdir $(LIB)) +else + SONAME_FLAGS = -Wl,-soname,$(notdir $(LIB)) +endif + +all: $(LIB) + +ifeq ($(LANGUAGE),cuda) + + NVCC ?= nvcc + # One may pass MADGRAPH_CUDA_ARCHITECTURE (comma-separated list), as for the matrix elements + comma := , + DETECTED_CC := $(shell nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | grep -E '^[0-9]+\.[0-9]+$$' | tr -d '.' | sort -un) + MADGRAPH_CUDA_ARCHITECTURE ?= $(foreach arch,$(DETECTED_CC),$(arch)$(comma)) + CUDA_ARCH_LIST := $(subst $(comma), ,$(MADGRAPH_CUDA_ARCHITECTURE)) + ifeq ($(strip $(CUDA_ARCH_LIST)),) + CUDA_ARCH_LIST := 70 + endif + GPUARCHFLAGS := $(foreach arch,$(CUDA_ARCH_LIST),-gencode arch=compute_$(arch),code=sm_$(arch)) + +$(LIB): wrapper_gpu.cu $(BODY) umami.h + @mkdir -p $(dir $(LIB)) + $(NVCC) $(OPTFLAGS) -std=c++17 $(GPUARCHFLAGS) -Xcompiler -fPIC --shared $(CFLAGS) \ + -o $@ wrapper_gpu.cu $(SONAME_FLAGS) $(LIBS) + +else ifeq ($(LANGUAGE),hip) + + HIP_HOME ?= $(shell hipconfig --rocmpath 2>/dev/null) + HIPCC ?= $(HIP_HOME)/bin/hipcc + # One may pass MADGRAPH_HIP_ARCHITECTURE, as for the matrix elements + ifeq ($(origin MADGRAPH_HIP_ARCHITECTURE),undefined) + MADGRAPH_HIP_ARCHITECTURE := $(shell $(HIP_HOME)/bin/rocm_agent_enumerator 2>/dev/null | grep -v gfx000 | sort -u | head -1) + endif + ifeq ($(MADGRAPH_HIP_ARCHITECTURE),) + MADGRAPH_HIP_ARCHITECTURE := gfx942 + endif + +$(LIB): wrapper_gpu.cu $(BODY) umami.h + @mkdir -p $(dir $(LIB)) + $(HIPCC) $(OPTFLAGS) -std=c++17 --offload-arch=$(MADGRAPH_HIP_ARCHITECTURE) -fPIC -shared \ + $(CFLAGS) -x hip -o $@ wrapper_gpu.cu $(SONAME_FLAGS) $(LIBS) + +else + + ifeq ($(LANGUAGE),c) + BODY_COMPILE = $(CC) $(OPTFLAGS) -fPIC $(CFLAGS) + LINKER = $(CC) + else ifeq ($(LANGUAGE),cpp) + BODY_COMPILE = $(CXX) $(OPTFLAGS) -std=c++17 -fPIC $(CFLAGS) + LINKER = $(CXX) + else ifeq ($(LANGUAGE),fortran) + BODY_COMPILE = $(FC) $(OPTFLAGS) -fPIC $(CFLAGS) + LINKER = $(FC) + else + $(error unknown LANGUAGE '$(LANGUAGE)': expected c, cpp, fortran, cuda or hip) + endif + +wrapper_cpu.o: wrapper_cpu.c umami.h + $(CC) $(OPTFLAGS) -fPIC -c -o $@ wrapper_cpu.c + +body.o: $(BODY) + $(BODY_COMPILE) -c -o $@ $(BODY) + +$(LIB): wrapper_cpu.o body.o + @mkdir -p $(dir $(LIB)) + $(LINKER) -shared -o $@ wrapper_cpu.o body.o $(SONAME_FLAGS) $(LIBS) + +endif + +clean: + rm -f *.o *.mod $(LIB) + +.PHONY: all clean diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/bias.c b/madgraph/iolibs/template_files/mg7/custom_functions/bias.c new file mode 100644 index 0000000000..04e2464c61 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/bias.c @@ -0,0 +1,29 @@ +/* + * Custom event bias for MG7, in C. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the bias factor b for one phase-space point. Events are generated + * according to |w| b, and written with the weight divided by b, so physical + * distributions are unchanged while regions with large b get more events. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_bias.so. + */ + +#include + +double custom_bias(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + (void)momenta; + // Example: populate the tail of the pT distribution of the first + // final-state particle, b = (pT / 100 GeV)^2 with a floor of 1. + // double px = momenta[2][1], py = momenta[2][2]; + // return fmax((px * px + py * py) / (100. * 100.), 1.); + + return 1.; +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/bias.cpp b/madgraph/iolibs/template_files/mg7/custom_functions/bias.cpp new file mode 100644 index 0000000000..96d703705c --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/bias.cpp @@ -0,0 +1,33 @@ +/* + * Custom event bias for MG7, in C++. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the bias factor b for one phase-space point. Events are generated + * according to |w| b, and written with the weight divided by b, so physical + * distributions are unchanged while regions with large b get more events. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_bias.so. + */ + +#include + +using std::fmax; +using std::pow; +using std::sqrt; + +extern "C" double custom_bias(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + (void)momenta; + // Example: populate the tail of the pT distribution of the first + // final-state particle, b = (pT / 100 GeV)^2 with a floor of 1. + // double px = momenta[2][1], py = momenta[2][2]; + // return fmax((px * px + py * py) / (100. * 100.), 1.); + + return 1.; +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/bias.cu b/madgraph/iolibs/template_files/mg7/custom_functions/bias.cu new file mode 100644 index 0000000000..f555ba2305 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/bias.cu @@ -0,0 +1,28 @@ +/* + * Custom event bias for MG7, in CUDA or HIP. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the bias factor b for one phase-space point. Events are generated + * according to |w| b, and written with the weight divided by b, so physical + * distributions are unchanged while regions with large b get more events. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_bias.so. + * The function runs on the GPU, one thread per event. + */ + +__device__ double custom_bias(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + (void)momenta; + // Example: populate the tail of the pT distribution of the first + // final-state particle, b = (pT / 100 GeV)^2 with a floor of 1. + // double px = momenta[2][1], py = momenta[2][2]; + // return fmax((px * px + py * py) / (100. * 100.), 1.); + + return 1.; +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/bias.f90 b/madgraph/iolibs/template_files/mg7/custom_functions/bias.f90 new file mode 100644 index 0000000000..c2d63721d4 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/bias.f90 @@ -0,0 +1,28 @@ +! Custom event bias for MG7, in Fortran. Generated by bin/init_custom_func. +! +! It is called once per phase-space point and returns +! the bias factor b for one phase-space point. Events are generated +! according to |w| b, and written with the weight divided by b, so physical +! distributions are unchanged while regions with large b get more events. +! +! Parameters: +! particle_count number of external particles, incoming ones first +! pdg_ids PDG ids of the external particles +! momenta momenta(:, i) = (E, px, py, pz) of particle i in GeV, lab frame +! +! Run `make` in this directory to build ../lib/custom_bias.so. + +function custom_bias(particle_count, pdg_ids, momenta) bind(C, name="custom_bias") result(value) + use iso_c_binding + implicit none + integer(c_int), value :: particle_count + integer(c_int), intent(in) :: pdg_ids(particle_count) + real(c_double), intent(in) :: momenta(0:3, particle_count) + real(c_double) :: value + + ! Example: populate the tail of the pT distribution of the first + ! final-state particle, b = (pT / 100 GeV)^2 with a floor of 1. + ! value = max((momenta(1, 3)**2 + momenta(2, 3)**2) / 100d0**2, 1d0) + + value = 1d0 +end function custom_bias diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/cut.c b/madgraph/iolibs/template_files/mg7/custom_functions/cut.c new file mode 100644 index 0000000000..0affcab7b4 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/cut.c @@ -0,0 +1,27 @@ +/* + * Custom cut for MG7, in C. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the cut value for one phase-space point: 1 keeps the event, 0 rejects it. + * Values in between act as a weight factor. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_cut.so. + */ + +#include + +double custom_cut(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + (void)momenta; + // Example: require pT > 20 GeV for the first final-state particle. + // double px = momenta[2][1], py = momenta[2][2]; + // return px * px + py * py > 20. * 20. ? 1. : 0.; + + return 1.; +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/cut.cpp b/madgraph/iolibs/template_files/mg7/custom_functions/cut.cpp new file mode 100644 index 0000000000..67cd66014c --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/cut.cpp @@ -0,0 +1,31 @@ +/* + * Custom cut for MG7, in C++. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the cut value for one phase-space point: 1 keeps the event, 0 rejects it. + * Values in between act as a weight factor. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_cut.so. + */ + +#include + +using std::fmax; +using std::pow; +using std::sqrt; + +extern "C" double custom_cut(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + (void)momenta; + // Example: require pT > 20 GeV for the first final-state particle. + // double px = momenta[2][1], py = momenta[2][2]; + // return px * px + py * py > 20. * 20. ? 1. : 0.; + + return 1.; +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/cut.cu b/madgraph/iolibs/template_files/mg7/custom_functions/cut.cu new file mode 100644 index 0000000000..c20b6987fc --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/cut.cu @@ -0,0 +1,26 @@ +/* + * Custom cut for MG7, in CUDA or HIP. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the cut value for one phase-space point: 1 keeps the event, 0 rejects it. + * Values in between act as a weight factor. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_cut.so. + * The function runs on the GPU, one thread per event. + */ + +__device__ double custom_cut(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + (void)momenta; + // Example: require pT > 20 GeV for the first final-state particle. + // double px = momenta[2][1], py = momenta[2][2]; + // return px * px + py * py > 20. * 20. ? 1. : 0.; + + return 1.; +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/cut.f90 b/madgraph/iolibs/template_files/mg7/custom_functions/cut.f90 new file mode 100644 index 0000000000..3a9e69e137 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/cut.f90 @@ -0,0 +1,30 @@ +! Custom cut for MG7, in Fortran. Generated by bin/init_custom_func. +! +! It is called once per phase-space point and returns +! the cut value for one phase-space point: 1 keeps the event, 0 rejects it. +! Values in between act as a weight factor. +! +! Parameters: +! particle_count number of external particles, incoming ones first +! pdg_ids PDG ids of the external particles +! momenta momenta(:, i) = (E, px, py, pz) of particle i in GeV, lab frame +! +! Run `make` in this directory to build ../lib/custom_cut.so. + +function custom_cut(particle_count, pdg_ids, momenta) bind(C, name="custom_cut") result(value) + use iso_c_binding + implicit none + integer(c_int), value :: particle_count + integer(c_int), intent(in) :: pdg_ids(particle_count) + real(c_double), intent(in) :: momenta(0:3, particle_count) + real(c_double) :: value + + ! Example: require pT > 20 GeV for the first final-state particle. + ! if (momenta(1, 3)**2 + momenta(2, 3)**2 > 20d0**2) then + ! value = 1d0 + ! else + ! value = 0d0 + ! end if + + value = 1d0 +end function custom_cut diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/pdf.c b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.c new file mode 100644 index 0000000000..6cb5f9fc2c --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.c @@ -0,0 +1,27 @@ +/* + * Custom parton density for MG7, in C. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * x f(x, q), the momentum fraction times the parton density, for the parton + * with PDG id pdg_id (21 for the gluon) in the beam beam_index (0 or 1), at + * momentum fraction x and factorization scale q in GeV. The strong coupling + * still comes from the PDF set in the run card. + * + * Parameters: + * beam_index index of the beam, 0 or 1 + * pdg_id PDG id of the parton + * x momentum fraction + * q factorization scale in GeV + * + * Run `make` in this directory to build ../lib/custom_pdf.so. + */ + +#include + +double custom_pdf(int beam_index, int pdg_id, double x, double q) { + // default: a toy density, replace it with your own + (void)beam_index; + (void)q; + double norm = pdg_id == 21 ? 3. : 0.5; + return norm * sqrt(x) * pow(1. - x, 3); +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/pdf.cpp b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.cpp new file mode 100644 index 0000000000..c605a3bb92 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.cpp @@ -0,0 +1,36 @@ +/* + * Custom parton density for MG7, in C++. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * x f(x, q), the momentum fraction times the parton density, for the parton + * with PDG id pdg_id (21 for the gluon) in the beam beam_index (0 or 1), at + * momentum fraction x and factorization scale q in GeV. The strong coupling + * still comes from the PDF set in the run card. + * + * Parameters: + * beam_index index of the beam, 0 or 1 + * pdg_id PDG id of the parton + * x momentum fraction + * q factorization scale in GeV + * + * Run `make` in this directory to build ../lib/custom_pdf.so. + */ + +#include + +using std::fmax; +using std::pow; +using std::sqrt; + +extern "C" double custom_pdf(int beam_index, int pdg_id, double x, double q) { + // Example: call LHAPDF. Add #include "LHAPDF/LHAPDF.h" above, and the + // output of `lhapdf-config --cflags --libs` to CFLAGS and LIBS in the Makefile. + // static LHAPDF::PDF* pdf = LHAPDF::mkPDF("NNPDF23_lo_as_0130_qed", 0); + // return pdf->xfxQ(pdg_id, x, q); + + // default: a toy density, replace it with your own + (void)beam_index; + (void)q; + double norm = pdg_id == 21 ? 3. : 0.5; + return norm * sqrt(x) * pow(1. - x, 3); +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/pdf.cu b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.cu new file mode 100644 index 0000000000..dec659e12a --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.cu @@ -0,0 +1,26 @@ +/* + * Custom parton density for MG7, in CUDA or HIP. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * x f(x, q), the momentum fraction times the parton density, for the parton + * with PDG id pdg_id (21 for the gluon) in the beam beam_index (0 or 1), at + * momentum fraction x and factorization scale q in GeV. The strong coupling + * still comes from the PDF set in the run card. + * + * Parameters: + * beam_index index of the beam, 0 or 1 + * pdg_id PDG id of the parton + * x momentum fraction + * q factorization scale in GeV + * + * Run `make` in this directory to build ../lib/custom_pdf.so. + * The function runs on the GPU, one thread per event. + */ + +__device__ double custom_pdf(int beam_index, int pdg_id, double x, double q) { + // default: a toy density, replace it with your own + (void)beam_index; + (void)q; + double norm = pdg_id == 21 ? 3. : 0.5; + return norm * sqrt(x) * pow(1. - x, 3); +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/pdf.f90 b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.f90 new file mode 100644 index 0000000000..f9a9e5ebfd --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.f90 @@ -0,0 +1,32 @@ +! Custom parton density for MG7, in Fortran. Generated by bin/init_custom_func. +! +! It is called once per phase-space point and returns +! x f(x, q), the momentum fraction times the parton density, for the parton +! with PDG id pdg_id (21 for the gluon) in the beam beam_index (0 or 1), at +! momentum fraction x and factorization scale q in GeV. The strong coupling +! still comes from the PDF set in the run card. +! +! Parameters: +! beam_index index of the beam, 0 or 1 +! pdg_id PDG id of the parton +! x momentum fraction +! q factorization scale in GeV +! +! Run `make` in this directory to build ../lib/custom_pdf.so. + +function custom_pdf(beam_index, pdg_id, x, q) bind(C, name="custom_pdf") result(value) + use iso_c_binding + implicit none + integer(c_int), value :: beam_index, pdg_id + real(c_double), value :: x, q + real(c_double) :: value + real(c_double) :: norm + + ! default: a toy density, replace it with your own + if (pdg_id == 21) then + norm = 3d0 + else + norm = 0.5d0 + end if + value = norm * sqrt(x) * (1d0 - x)**3 +end function custom_pdf diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/scale.c b/madgraph/iolibs/template_files/mg7/custom_functions/scale.c new file mode 100644 index 0000000000..5ec5fc2a19 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/scale.c @@ -0,0 +1,36 @@ +/* + * Custom energy scale for MG7, in C. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the renormalization and factorization scale in GeV for one phase-space + * point. The run card scale_factor is applied on top, and fixed scales still + * take precedence. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_scale.so. + */ + +#include + +double custom_scale(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + // Example: H_T / 2, the scalar sum of the final-state transverse masses. + // double ht = 0.; + // for (int i = 2; i < particle_count; ++i) { + // double e = momenta[i][0], pz = momenta[i][3]; + // ht += sqrt(fmax(e * e - pz * pz, 0.)); + // } + // return ht / 2.; + + // default: partonic center-of-mass energy sqrt(s_hat) + double e = momenta[0][0] + momenta[1][0]; + double px = momenta[0][1] + momenta[1][1]; + double py = momenta[0][2] + momenta[1][2]; + double pz = momenta[0][3] + momenta[1][3]; + return sqrt(fmax(e * e - px * px - py * py - pz * pz, 0.)); +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/scale.cpp b/madgraph/iolibs/template_files/mg7/custom_functions/scale.cpp new file mode 100644 index 0000000000..b91fc25c85 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/scale.cpp @@ -0,0 +1,40 @@ +/* + * Custom energy scale for MG7, in C++. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the renormalization and factorization scale in GeV for one phase-space + * point. The run card scale_factor is applied on top, and fixed scales still + * take precedence. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_scale.so. + */ + +#include + +using std::fmax; +using std::pow; +using std::sqrt; + +extern "C" double custom_scale(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + // Example: H_T / 2, the scalar sum of the final-state transverse masses. + // double ht = 0.; + // for (int i = 2; i < particle_count; ++i) { + // double e = momenta[i][0], pz = momenta[i][3]; + // ht += sqrt(fmax(e * e - pz * pz, 0.)); + // } + // return ht / 2.; + + // default: partonic center-of-mass energy sqrt(s_hat) + double e = momenta[0][0] + momenta[1][0]; + double px = momenta[0][1] + momenta[1][1]; + double py = momenta[0][2] + momenta[1][2]; + double pz = momenta[0][3] + momenta[1][3]; + return sqrt(fmax(e * e - px * px - py * py - pz * pz, 0.)); +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/scale.cu b/madgraph/iolibs/template_files/mg7/custom_functions/scale.cu new file mode 100644 index 0000000000..05679b90ce --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/scale.cu @@ -0,0 +1,35 @@ +/* + * Custom energy scale for MG7, in CUDA or HIP. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the renormalization and factorization scale in GeV for one phase-space + * point. The run card scale_factor is applied on top, and fixed scales still + * take precedence. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_scale.so. + * The function runs on the GPU, one thread per event. + */ + +__device__ double custom_scale(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + // Example: H_T / 2, the scalar sum of the final-state transverse masses. + // double ht = 0.; + // for (int i = 2; i < particle_count; ++i) { + // double e = momenta[i][0], pz = momenta[i][3]; + // ht += sqrt(fmax(e * e - pz * pz, 0.)); + // } + // return ht / 2.; + + // default: partonic center-of-mass energy sqrt(s_hat) + double e = momenta[0][0] + momenta[1][0]; + double px = momenta[0][1] + momenta[1][1]; + double py = momenta[0][2] + momenta[1][2]; + double pz = momenta[0][3] + momenta[1][3]; + return sqrt(fmax(e * e - px * px - py * py - pz * pz, 0.)); +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/scale.f90 b/madgraph/iolibs/template_files/mg7/custom_functions/scale.f90 new file mode 100644 index 0000000000..6fcc401b1d --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/scale.f90 @@ -0,0 +1,35 @@ +! Custom energy scale for MG7, in Fortran. Generated by bin/init_custom_func. +! +! It is called once per phase-space point and returns +! the renormalization and factorization scale in GeV for one phase-space +! point. The run card scale_factor is applied on top, and fixed scales still +! take precedence. +! +! Parameters: +! particle_count number of external particles, incoming ones first +! pdg_ids PDG ids of the external particles +! momenta momenta(:, i) = (E, px, py, pz) of particle i in GeV, lab frame +! +! Run `make` in this directory to build ../lib/custom_scale.so. + +function custom_scale(particle_count, pdg_ids, momenta) bind(C, name="custom_scale") result(value) + use iso_c_binding + implicit none + integer(c_int), value :: particle_count + integer(c_int), intent(in) :: pdg_ids(particle_count) + real(c_double), intent(in) :: momenta(0:3, particle_count) + real(c_double) :: value + real(c_double) :: p(0:3) + + ! Example: H_T / 2, the scalar sum of the final-state transverse masses. + ! integer :: i + ! value = 0d0 + ! do i = 3, particle_count + ! value = value + sqrt(max(momenta(0, i)**2 - momenta(3, i)**2, 0d0)) + ! end do + ! value = value / 2d0 + + ! default: partonic center-of-mass energy sqrt(s_hat) + p = momenta(:, 1) + momenta(:, 2) + value = sqrt(max(p(0)**2 - p(1)**2 - p(2)**2 - p(3)**2, 0d0)) +end function custom_scale diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/wrapper_cpu.c b/madgraph/iolibs/template_files/mg7/custom_functions/wrapper_cpu.c new file mode 100644 index 0000000000..dbca693923 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/wrapper_cpu.c @@ -0,0 +1,195 @@ +/* + * UMAMI wrapper for the custom %(kind)s function, generated by bin/init_custom_func. + * It unpacks the batched UMAMI inputs and calls %(function_name)s once per event. + * Write your function in %(body_file)s, there is no need to edit this file. + */ + +#include "umami.h" + +#include + +#define CUSTOM_IS_PDF %(is_pdf)d +#define CUSTOM_OUTPUT_KEY %(output_key)s + +#if CUSTOM_IS_PDF +double %(function_name)s(int beam_index, int pdg_id, double x, double q); +static const int input_keys[] = { + UMAMI_IN_MOMENTUM_FRACTION, UMAMI_IN_SCALE, UMAMI_IN_PDG_ID, UMAMI_IN_BEAM_INDEX +}; +#else +double %(function_name)s(int particle_count, const int* pdg_ids, const double (*momenta)[4]); +static const int input_keys[] = {UMAMI_IN_MOMENTA, UMAMI_IN_PARTICLE_COUNT, UMAMI_IN_PDG_IDS}; +#endif +#define INPUT_COUNT (sizeof(input_keys) / sizeof(input_keys[0])) + +static bool supported_inputs[UMAMI_INPUT_KEY_COUNT]; +static bool supported_outputs[UMAMI_OUTPUT_KEY_COUNT]; +static int dummy_instance; + +static void fill_key_tables(void) { + for (size_t i = 0; i < INPUT_COUNT; ++i) { + supported_inputs[input_keys[i]] = true; + } + supported_outputs[CUSTOM_OUTPUT_KEY] = true; +} + +UmamiStatus umami_get_meta(UmamiMetaKey meta_key, void* result) { + if (meta_key == UMAMI_META_DEVICE) { + *(UmamiDevice*)result = UMAMI_DEVICE_CPU; + return UMAMI_SUCCESS; + } + return UMAMI_ERROR_UNSUPPORTED_META; +} + +UmamiStatus umami_supported_inputs(bool const** supported, int* count) { + fill_key_tables(); + *supported = supported_inputs; + *count = UMAMI_INPUT_KEY_COUNT; + return UMAMI_SUCCESS; +} + +UmamiStatus umami_required_inputs(bool const** required, int* count) { + return umami_supported_inputs(required, count); +} + +UmamiStatus umami_supported_outputs(bool const** supported, int* count) { + fill_key_tables(); + *supported = supported_outputs; + *count = UMAMI_OUTPUT_KEY_COUNT; + return UMAMI_SUCCESS; +} + +UmamiStatus umami_initialize(UmamiHandle* handle, char const* param_card_path) { + (void)param_card_path; + *handle = &dummy_instance; + return UMAMI_SUCCESS; +} + +UmamiStatus umami_set_parameter( + UmamiHandle handle, char const* name, double parameter_real, double parameter_imag +) { + (void)handle; + (void)name; + (void)parameter_real; + (void)parameter_imag; + return UMAMI_ERROR_NOT_IMPLEMENTED; +} + +UmamiStatus umami_get_parameter( + UmamiHandle handle, char const* name, double* parameter_real, double* parameter_imag +) { + (void)handle; + (void)name; + (void)parameter_real; + (void)parameter_imag; + return UMAMI_ERROR_NOT_IMPLEMENTED; +} + +UmamiStatus umami_matrix_element( + UmamiHandle handle, + size_t count, + size_t stride, + size_t offset, + size_t input_count, + UmamiInputKey const* input_keys_in, + void const* const* inputs, + size_t output_count, + UmamiOutputKey const* output_keys, + void* const* outputs +) { + (void)handle; + double* output = NULL; + for (size_t i = 0; i < output_count; ++i) { + if (output_keys[i] == CUSTOM_OUTPUT_KEY) { + output = (double*)outputs[i]; + } else if (output_keys[i] != UMAMI_OUT_GPU_STREAM) { + return UMAMI_ERROR_UNSUPPORTED_OUTPUT; + } + } + if (output == NULL) { + return UMAMI_SUCCESS; + } + +#if CUSTOM_IS_PDF + const double* x = NULL; + const double* q = NULL; + const int* pdg_id = NULL; + const int* beam_index = NULL; + for (size_t i = 0; i < input_count; ++i) { + switch (input_keys_in[i]) { + case UMAMI_IN_MOMENTUM_FRACTION: + x = (const double*)inputs[i]; + break; + case UMAMI_IN_SCALE: + q = (const double*)inputs[i]; + break; + case UMAMI_IN_PDG_ID: + pdg_id = (const int*)inputs[i]; + break; + case UMAMI_IN_BEAM_INDEX: + beam_index = (const int*)inputs[i]; + break; + default: + return UMAMI_ERROR_UNSUPPORTED_INPUT; + } + } + if (x == NULL || q == NULL || pdg_id == NULL || beam_index == NULL) { + return UMAMI_ERROR_MISSING_INPUT; + } + for (size_t i = offset; i < offset + count; ++i) { + output[i] = %(function_name)s(beam_index[i], pdg_id[i], x[i], q[i]); + } +#else + const double* momenta = NULL; + const int* particle_count = NULL; + const int* pdg_ids = NULL; + for (size_t i = 0; i < input_count; ++i) { + switch (input_keys_in[i]) { + case UMAMI_IN_MOMENTA: + momenta = (const double*)inputs[i]; + break; + case UMAMI_IN_PARTICLE_COUNT: + particle_count = (const int*)inputs[i]; + break; + case UMAMI_IN_PDG_IDS: + pdg_ids = (const int*)inputs[i]; + break; + default: + return UMAMI_ERROR_UNSUPPORTED_INPUT; + } + } + if (momenta == NULL || particle_count == NULL || pdg_ids == NULL) { + return UMAMI_ERROR_MISSING_INPUT; + } + if (count == 0) { + return UMAMI_SUCCESS; + } + // the particle count is the same for every event of a batch + int n = particle_count[offset]; + int* event_pdg_ids = malloc(n * sizeof(int)); + double (*event_momenta)[4] = malloc(n * sizeof(double[4])); + if (event_pdg_ids == NULL || event_momenta == NULL) { + free(event_pdg_ids); + free(event_momenta); + return UMAMI_ERROR; + } + for (size_t i = offset; i < offset + count; ++i) { + // column-major layout with the batch dimension contiguous + for (int j = 0; j < n; ++j) { + event_pdg_ids[j] = pdg_ids[stride * j + i]; + for (int k = 0; k < 4; ++k) { + event_momenta[j][k] = momenta[stride * (n * k + j) + i]; + } + } + output[i] = %(function_name)s(n, event_pdg_ids, (const double (*)[4])event_momenta); + } + free(event_pdg_ids); + free(event_momenta); +#endif + return UMAMI_SUCCESS; +} + +UmamiStatus umami_free(UmamiHandle handle) { + (void)handle; + return UMAMI_SUCCESS; +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/wrapper_gpu.cu b/madgraph/iolibs/template_files/mg7/custom_functions/wrapper_gpu.cu new file mode 100644 index 0000000000..c0033ce034 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/wrapper_gpu.cu @@ -0,0 +1,226 @@ +// UMAMI wrapper for the custom %(kind)s function, generated by bin/init_custom_func. +// It unpacks the batched UMAMI inputs and calls %(function_name)s once per event, +// one GPU thread per event. Write your function in %(body_file)s, there is no need +// to edit this file. + +#include "umami.h" + +#ifdef __HIPCC__ +#include +typedef hipStream_t gpuStream_t; +#define gpuGetLastError hipGetLastError +#define gpuSuccess hipSuccess +#define UMAMI_GPU_DEVICE UMAMI_DEVICE_HIP +#else +#include +typedef cudaStream_t gpuStream_t; +#define gpuGetLastError cudaGetLastError +#define gpuSuccess cudaSuccess +#define UMAMI_GPU_DEVICE UMAMI_DEVICE_CUDA +#endif + +#define CUSTOM_IS_PDF %(is_pdf)d +#define CUSTOM_OUTPUT_KEY %(output_key)s + +// the function body is compiled into this translation unit +#include "%(body_file)s" + +namespace { + +constexpr int thread_count = 256; +int dummy_instance; + +#if CUSTOM_IS_PDF +constexpr UmamiInputKey input_keys[] = { + UMAMI_IN_MOMENTUM_FRACTION, UMAMI_IN_SCALE, UMAMI_IN_PDG_ID, UMAMI_IN_BEAM_INDEX +}; + +__global__ void evaluate( + std::size_t count, + const double* x, + const double* q, + const int* pdg_id, + const int* beam_index, + double* output +) { + std::size_t i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < count) { + output[i] = %(function_name)s(beam_index[i], pdg_id[i], x[i], q[i]); + } +} +#else +constexpr UmamiInputKey input_keys[] = { + UMAMI_IN_MOMENTA, UMAMI_IN_PARTICLE_COUNT, UMAMI_IN_PDG_IDS +}; +// largest particle count supported by the GPU wrapper +constexpr int max_particle_count = 32; + +__global__ void evaluate( + std::size_t count, + std::size_t stride, + const double* momenta, + const int* particle_count, + const int* pdg_ids, + double* output +) { + std::size_t i = blockDim.x * blockIdx.x + threadIdx.x; + if (i >= count) { + return; + } + int n = particle_count[i]; + int event_pdg_ids[max_particle_count]; + double event_momenta[max_particle_count][4]; + // column-major layout with the batch dimension contiguous + for (int j = 0; j < n; ++j) { + event_pdg_ids[j] = pdg_ids[stride * j + i]; + for (int k = 0; k < 4; ++k) { + event_momenta[j][k] = momenta[stride * (n * k + j) + i]; + } + } + output[i] = %(function_name)s(n, event_pdg_ids, event_momenta); +} +#endif + +struct KeyTables { + bool inputs[UMAMI_INPUT_KEY_COUNT] = {}; + bool outputs[UMAMI_OUTPUT_KEY_COUNT] = {}; + KeyTables() { + for (auto key : input_keys) { + inputs[key] = true; + } + outputs[CUSTOM_OUTPUT_KEY] = true; + outputs[UMAMI_OUT_GPU_STREAM] = true; + } +}; +const KeyTables key_tables; + +} // namespace + +extern "C" { + +UmamiStatus umami_get_meta(UmamiMetaKey meta_key, void* result) { + if (meta_key == UMAMI_META_DEVICE) { + *static_cast(result) = UMAMI_GPU_DEVICE; + return UMAMI_SUCCESS; + } + return UMAMI_ERROR_UNSUPPORTED_META; +} + +UmamiStatus umami_supported_inputs(bool const** supported, int* count) { + *supported = key_tables.inputs; + *count = UMAMI_INPUT_KEY_COUNT; + return UMAMI_SUCCESS; +} + +UmamiStatus umami_required_inputs(bool const** required, int* count) { + return umami_supported_inputs(required, count); +} + +UmamiStatus umami_supported_outputs(bool const** supported, int* count) { + *supported = key_tables.outputs; + *count = UMAMI_OUTPUT_KEY_COUNT; + return UMAMI_SUCCESS; +} + +UmamiStatus umami_initialize(UmamiHandle* handle, char const*) { + *handle = &dummy_instance; + return UMAMI_SUCCESS; +} + +UmamiStatus umami_set_parameter(UmamiHandle, char const*, double, double) { + return UMAMI_ERROR_NOT_IMPLEMENTED; +} + +UmamiStatus umami_get_parameter(UmamiHandle, char const*, double*, double*) { + return UMAMI_ERROR_NOT_IMPLEMENTED; +} + +UmamiStatus umami_matrix_element( + UmamiHandle, + std::size_t count, + std::size_t stride, + std::size_t offset, + std::size_t input_count, + UmamiInputKey const* input_keys_in, + void const* const* inputs, + std::size_t output_count, + UmamiOutputKey const* output_keys, + void* const* outputs +) { + double* output = nullptr; + gpuStream_t stream = nullptr; + for (std::size_t i = 0; i < output_count; ++i) { + if (output_keys[i] == CUSTOM_OUTPUT_KEY) { + output = static_cast(outputs[i]) + offset; + } else if (output_keys[i] == UMAMI_OUT_GPU_STREAM) { + stream = static_cast(outputs[i]); + } else { + return UMAMI_ERROR_UNSUPPORTED_OUTPUT; + } + } + if (output == nullptr || count == 0) { + return UMAMI_SUCCESS; + } + unsigned int block_count = (count + thread_count - 1) / thread_count; + +#if CUSTOM_IS_PDF + const double* x = nullptr; + const double* q = nullptr; + const int* pdg_id = nullptr; + const int* beam_index = nullptr; + for (std::size_t i = 0; i < input_count; ++i) { + switch (input_keys_in[i]) { + case UMAMI_IN_MOMENTUM_FRACTION: + x = static_cast(inputs[i]) + offset; + break; + case UMAMI_IN_SCALE: + q = static_cast(inputs[i]) + offset; + break; + case UMAMI_IN_PDG_ID: + pdg_id = static_cast(inputs[i]) + offset; + break; + case UMAMI_IN_BEAM_INDEX: + beam_index = static_cast(inputs[i]) + offset; + break; + default: + return UMAMI_ERROR_UNSUPPORTED_INPUT; + } + } + if (!x || !q || !pdg_id || !beam_index) { + return UMAMI_ERROR_MISSING_INPUT; + } + evaluate<<>>( + count, x, q, pdg_id, beam_index, output + ); +#else + const double* momenta = nullptr; + const int* particle_count = nullptr; + const int* pdg_ids = nullptr; + for (std::size_t i = 0; i < input_count; ++i) { + switch (input_keys_in[i]) { + case UMAMI_IN_MOMENTA: + momenta = static_cast(inputs[i]) + offset; + break; + case UMAMI_IN_PARTICLE_COUNT: + particle_count = static_cast(inputs[i]) + offset; + break; + case UMAMI_IN_PDG_IDS: + pdg_ids = static_cast(inputs[i]) + offset; + break; + default: + return UMAMI_ERROR_UNSUPPORTED_INPUT; + } + } + if (!momenta || !particle_count || !pdg_ids) { + return UMAMI_ERROR_MISSING_INPUT; + } + evaluate<<>>( + count, stride, momenta, particle_count, pdg_ids, output + ); +#endif + return gpuGetLastError() == gpuSuccess ? UMAMI_SUCCESS : UMAMI_ERROR; +} + +UmamiStatus umami_free(UmamiHandle) { return UMAMI_SUCCESS; } + +} // extern "C" diff --git a/madgraph/iolibs/template_files/mg7/init_custom_func.py b/madgraph/iolibs/template_files/mg7/init_custom_func.py new file mode 100644 index 0000000000..af50527977 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/init_custom_func.py @@ -0,0 +1,222 @@ +"""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_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, 1 keeps the event and 0 rejects it"), + "scale": ("UMAMI_OUT_SCALE", "the renormalization and factorization scale"), + "pdf": ("UMAMI_OUT_PDF", "a parton density x f(x, Q)"), + "bias": ("UMAMI_OUT_BIAS", "an event bias for the unweighting"), +} +# 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"), +} + +_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] + for i, name in enumerate(names, 1): + lines.append(f" {i}. {name:8s} {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_/ and build it into lib/custom_.so." + ), + ) + 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 write_files(out_dir, function, language): + """Write the wrapper, the function body, umami.h and the Makefile.""" + extension = LANGUAGES[language][0] + body_file = f"custom_{function}.{extension}" + values = { + "kind": function, + "function_name": f"custom_{function}", + "output_key": FUNCTIONS[function][0], + "is_pdf": int(function == "pdf"), + "body_file": body_file, + "language": language, + } + 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"): + 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 = os.path.join("lib", f"custom_{function}.so") + 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 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: + 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/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 7c571d12d7..23989166dd 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -114,6 +114,14 @@ bw_cutoff = %(phasespace.bw_cutoff)s cut_decays = %(phasespace.cut_decays)s # apply the [cuts] to decay products of on-shell particles adaptive_symmetry_sampling = %(phasespace.adaptive_symmetry_sampling)s +[custom_functions] +# shared libraries built with bin/init_custom_func, e.g. "lib/custom_cut.so" +# paths are relative to the process directory; "" disables the function +cut = %(custom_functions.cut)s +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 + [multiparticles] $multiparticles diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 7619d334b8..6a899f7762 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6883,6 +6883,17 @@ def default_setup(self): self.add_toml_param('phasespace', 'adaptive_symmetry_sampling', True, comment="sample symmetric channels adaptively") + # ------------------------- [custom_functions] ----------------- + # user functions built with bin/init_custom_func; "" disables them + self.add_toml_param('custom_functions', 'cut', "", + comment="custom cut library, e.g. lib/custom_cut.so") + self.add_toml_param('custom_functions', 'scale', "", + comment="custom scale library, replaces dynamical_scale_choice") + self.add_toml_param('custom_functions', 'pdf', "", + comment="custom PDF library, replaces the PDF set (not alpha_s)") + self.add_toml_param('custom_functions', 'bias', "", + comment="custom event bias library") + # ----------------------------- [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 dca9029ccff13c7a90eddad77bff78c013f0c13a Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 12 Oct 2026 13:25:52 +0200 Subject: [PATCH 17/37] mg7: load [custom_functions] libraries and pass them to madspace The libraries are built from custom_/ if present, and loaded before the matrix elements in every context and gridpack. Custom cuts and scales get the PDG ids of the first flavor, the bias those of the sampled flavor. A custom PDF replaces the PDF grid (alpha_s still comes from the PDF set) and disables the systematics. --- .../iolibs/template_files/mg7/gridpack.py | 2 + madgraph/iolibs/template_files/mg7/launch.py | 131 +++++++++++++++--- 2 files changed, 117 insertions(+), 16 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index e7081fa56d..2f120b4ea7 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -214,6 +214,8 @@ def main() -> None: global_dir = os.path.join("data", "globals") for context, backend in zip(contexts, backends): context.load_globals(global_dir) + for path in madspace_data.get("custom_functions", []): + context.load_matrix_element(path, param_card_path) for me_path in madspace_data["matrix_elements"]: context.load_matrix_element( me_path.format(device=backend), param_card_path, me_parameters diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index f19ba110ef..bcc6a9b053 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -285,6 +285,7 @@ def __init__(self): self.init_backend() self.init_event_dir() self.init_context() + self.init_custom_functions() self.init_cuts() self.init_histograms() self.init_generator_config() @@ -514,6 +515,55 @@ def init_context(self) -> None: 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: + """Build and load the libraries of [custom_functions]. 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 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}: library not found. " + "Create it with bin/init_custom_func." + ) + apis = [context.load_matrix_element(path, "") for context in self.contexts] + if len({api.index() for api in apis}) != 1: + raise RuntimeError(f"custom {kind} library loaded at different indices") + self.custom_functions[kind] = apis[0] + self.custom_function_paths.append(path) + logger.info("Using the custom %s function %s", kind, path) + + @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 @@ -692,6 +742,10 @@ def init_beam(self) -> None: 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. @@ -720,10 +774,17 @@ def init_beam(self) -> None: self.pdf_dir = self.lhapdf.find_set(pdf_set) if self.pdf_dir is None: raise RuntimeError(self._no_pdf_message(pdf_set)) - self.pdf_grid = ms.PdfGrid(os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}_0000.dat")) 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: - self.pdf_grid.initialize_globals(context) + 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) @@ -734,14 +795,21 @@ def systematics_enabled(self) -> bool: """The native scale/PDF weights are requested ([systematics] enable, or the legacy [generation] systematics alias).""" try: - if self.run_card["systematics"]["enable"]: - return True - except Exception: - pass - try: - return bool(self.run_card["generation"]["systematics"]) + 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: @@ -1915,6 +1983,11 @@ def save_gridpack(self) -> None: shutil.copytree("SubProcesses", os.path.join(gridpack_path, "SubProcesses")) 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( @@ -1988,6 +2061,8 @@ def save_gridpack(self) -> None: data = { "channels": channel_files, "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, } @@ -2407,8 +2482,14 @@ def cut_pids(observable_kwargs): # 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)] - self.cuts = ( - ms.Cuts([ + # 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), @@ -2419,10 +2500,11 @@ def cut_pids(observable_kwargs): mode=cut_item.mode, ) for cut_item in self.process.cut_data - ]) - if len(self.process.cut_data) > 0 - else None - ) + ], 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 @@ -2452,7 +2534,11 @@ def cut_pids(observable_kwargs): ) self.scale = ms.EnergyScale( - particle_count=self.particle_count, **self.process.scale_kwargs + 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"]: @@ -2467,6 +2553,17 @@ def cut_pids(observable_kwargs): ) self.matrix_elements.append(mat) + 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 @@ -3141,6 +3238,7 @@ def build_integrands( ) 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( @@ -3171,7 +3269,8 @@ def build_integrands( flavor_diff_xs_indices, flavor_subproc_indices, flavor_per_subproc_remap, - madnis_args["compressed_channel_weight_count"] + madnis_args["compressed_channel_weight_count"], + custom_bias, )) #print(integrands[1].function()) #for i in integrands: print(i.function()) From 7da585d3431cc83468f1d4cca527d5c2f28a96cd Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 12 Oct 2026 13:31:49 +0200 Subject: [PATCH 18/37] madspace: divide the custom bias out of the cross section and event weights The integrand weight includes the bias b, so adaptation, unweighting and the channel allocation see |w| b. The signed running integral, which gives the reported cross section, accumulates w / b, and the event records carry b so that the combined events are written with weight w_unw / b. --- .../madspace/driver/channel_generator.hpp | 1 + .../madspace/driver/generator_data.hpp | 2 ++ madspace/include/madspace/driver/io.hpp | 13 +++++++---- madspace/src/driver/channel_generator.cpp | 23 ++++++++++++++++++- madspace/src/driver/event_generator.cpp | 14 ++++++++--- madspace/src/phasespace/integrand.cpp | 5 ++-- 6 files changed, 48 insertions(+), 10 deletions(-) diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 1830817719..e5c1dfb70b 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -181,6 +181,7 @@ class ChannelEventGenerator { int ren_scale, alpha_qcd; int x1, fact_scale1, x2, fact_scale2, partial_weight_product; int subprocess_index; + int bias; int random, rest; }; diff --git a/madspace/include/madspace/driver/generator_data.hpp b/madspace/include/madspace/driver/generator_data.hpp index ee2c4e87e9..aedf5e398b 100644 --- a/madspace/include/madspace/driver/generator_data.hpp +++ b/madspace/include/madspace/driver/generator_data.hpp @@ -243,6 +243,8 @@ struct GeneratorBatchJob { // for them and isn't read. std::size_t split_job_count; Tensor weights; + // per-event bias, empty without a custom bias + Tensor biases; TensorVec events; TensorVec unweighted_events; TensorVec hists; diff --git a/madspace/include/madspace/driver/io.hpp b/madspace/include/madspace/driver/io.hpp index 5e631ad78d..efcdd3444c 100644 --- a/madspace/include/madspace/driver/io.hpp +++ b/madspace/include/madspace/driver/io.hpp @@ -149,13 +149,12 @@ class EventRecord { static constexpr int f_beam2 = 32; static constexpr int f_partial_weights = 64; static constexpr int f_syst_weights = 128; + static constexpr int f_bias = 256; // `syst_weight_count` columns of systematic variation weights are appended // when f_syst_weights is set; `syst_weight_ids` names them (rwgt_). - static std::vector layout( - int fields, - const std::vector& syst_weight_ids = {} - ) { + static std::vector + layout(int fields, const std::vector& syst_weight_ids = {}) { std::vector ret; if (fields & f_weight) { ret.push_back({"weight", FieldLayout::f64, 0}); @@ -214,6 +213,9 @@ class EventRecord { ret.push_back({std::format("rwgt_{}", id), FieldLayout::f64, 7}); } } + if (fields & f_bias) { + ret.push_back({"bias", FieldLayout::f64, 8}); + } return ret; } @@ -257,6 +259,9 @@ class EventRecord { return &_data[_offsets[6] + 8 * index]; } + // event bias, divided out of the weight when the channels are combined + UnalignedRef bias() { return &_data[_offsets[7] + 0]; } + void from_lhe_event(const LHEEvent& event) { lhe_process_id() = event.process_id; lhe_weight() = event.weight; diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 86f471c5a2..36be9a206f 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -21,6 +21,9 @@ int event_extra_flags(const std::unordered_map& index_ if (index_map.contains("subprocess_index")) { flags |= EventRecord::f_subproc_index; } + if (index_map.contains("bias")) { + flags |= EventRecord::f_bias; + } return flags; } @@ -323,6 +326,7 @@ void ChannelEventGenerator::init_field_indices() { } else { _field_indices.subprocess_index = -1; } + _field_indices.bias = index_map.contains("bias") ? index_map.at("bias") : -1; _field_indices.random = index_map.at("random"); _field_indices.rest = _field_indices.random + 1; } @@ -361,7 +365,14 @@ void ChannelEventGenerator::integrate(const GeneratorBatchJob& job) { if (w_view[i] != 0) { ++sample_count_after_cuts; } - _cross_section.push(w_view[i]); + // With a bias b the weights are w b: the absolute integral, which drives + // the adaptation, unweighting and channel allocation, keeps the bias, + // while the reported cross section divides it out. + double weight = w_view[i]; + if (_field_indices.bias != -1 && weight != 0.) { + weight /= job.biases.view()[i]; + } + _cross_section.push(weight); _abs_cross_section.push(std::abs(w_view[i])); } _status.mean = _cross_section.mean(); @@ -556,6 +567,9 @@ void ChannelEventGenerator::start_job( throw std::runtime_error("job cancelled"); } job.weights = job.events.at(_field_indices.weight).cpu(); + if (_field_indices.bias != -1) { + job.biases = job.events.at(_field_indices.bias).cpu(); + } // observable_histograms/vegas_histogram/discrete_histogram don't consume // random numbers, so they're never seeded. if (runtimes.observable_histograms) { @@ -791,6 +805,13 @@ void ChannelEventGenerator::write_events( } } + if (_field_indices.bias != -1) { + auto bias_view = unweighted_events.at(_field_indices.bias).view(); + for (std::size_t i = 0; i < bias_view.size(); ++i) { + event_buffer.event(i).bias() = bias_view[i]; + } + } + if (_field_indices.subprocess_index != -1) { auto subproc_view = unweighted_events.at(_field_indices.subprocess_index).view(); diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index 39b55d84d4..bca833ca07 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -851,7 +851,8 @@ 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() + EventRecord::f_event_data | + (_channels.at(0)->event_layout_extra_flags() & ~EventRecord::f_bias) ), ParticleRecord::layout( ParticleRecord::f_particle_data | @@ -862,7 +863,8 @@ DataLayout EventGenerator::combined_layout() const { DataLayout EventGenerator::output_layout(int event_flags, std::size_t particle_flags) const { - int extra_flags = _channels.at(0)->event_layout_extra_flags(); + int extra_flags = + _channels.at(0)->event_layout_extra_flags() & ~EventRecord::f_bias; std::vector syst_ids; if (_systematics) { if (!_systematics->config().write_inputs) { @@ -1115,6 +1117,7 @@ void EventGenerator::read_and_combine( _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; + bool has_bias = _channels.at(0)->event_layout_extra_flags() & EventRecord::f_bias; 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); @@ -1152,6 +1155,10 @@ void EventGenerator::read_and_combine( std::max(1., std::abs(weight) / channel->max_weight()), weight ) * norm_factor; + if (has_bias) { + // unweighted with the bias b, so the event carries 1 / b + event_out.weight() = event_out.weight() / event_in.bias().value(); + } event_out.subprocess_index() = has_subproc_index ? event_in.subprocess_index().value() : static_cast(channel->status().subprocess); @@ -1442,7 +1449,8 @@ void EventGenerator::print_gen_update_pretty(bool done) { if (!std::isnan(_status.error)) { int_str = format_with_error(_status.mean, _status.error); if (std::abs(_status.error) < std::abs(_status.mean)) { - rel_str = std::format("{:.4f} %", _status.error / std::abs(_status.mean) * 100); + rel_str = + std::format("{:.4f} %", _status.error / std::abs(_status.mean) * 100); } rsd_str = std::format("{:.3f}", _status.rel_std_dev); uweff_str = std::format( diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index b8f5fde48e..b05999051c 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -150,10 +150,11 @@ Integrand::Integrand( ret_types.push_back("partial_weight_product", batch_float); } } - ret_types.push_back("random", batch_float_array(mapping.random_dim())); + // before "random": the outputs after it feed the discrete histograms if (custom_bias) { ret_types.push_back("bias", batch_float); } + ret_types.push_back("random", batch_float_array(mapping.random_dim())); if (mapping.channel_count() > 1 && !std::holds_alternative(discrete_sym)) { ret_types.push_back("channel_index_in_group", batch_int); @@ -1004,12 +1005,12 @@ NamedVector Integrand::build_common_part( ); } } - outputs.push_back("random", optional_cut(args.at("r"))); if (bias_acc) { outputs.push_back( "bias", scatter_or_drop(fb.full({1., batch_size_val}), bias_acc) ); } + outputs.push_back("random", optional_cut(args.at("r"))); if (has_permutations && !std::holds_alternative(_discrete_sym)) { outputs.push_back( From 1d6be6c98c766562a963c3c3b304e4645f41a549 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 12 Oct 2026 13:38:27 +0200 Subject: [PATCH 19/37] Acceptance tests and documentation for custom functions test_mg7_custom_functions.py compares a custom cut, scale, PDF and bias with the equivalent built-in settings on p p > t t~, and builds and evaluates every function in every language with a toolchain. A new Usage page describes bin/init_custom_func and the function signatures. --- docs/source/custom_functions.rst | 234 +++++++++++ docs/source/index.rst | 1 + docs/source/run_card.rst | 22 + .../test_mg7_custom_functions.py | 381 ++++++++++++++++++ 4 files changed, 638 insertions(+) create mode 100644 docs/source/custom_functions.rst create mode 100644 tests/acceptance_tests/test_mg7_custom_functions.py diff --git a/docs/source/custom_functions.rst b/docs/source/custom_functions.rst new file mode 100644 index 0000000000..1f16a48104 --- /dev/null +++ b/docs/source/custom_functions.rst @@ -0,0 +1,234 @@ +Custom functions +================ + +MG7 can call your own code for four parts of the calculation: + +``cut`` + Decides if a phase-space point passes the cuts. It is applied on top of the + ``[cuts]`` of the :doc:`run card `. + +``scale`` + The renormalization and factorization scale. It replaces ``dynamical_scale_choice``, + like ``user_dynamical_scale`` in MadEvent. + +``pdf`` + The parton densities. It replaces the LHAPDF set of ``[beam] pdf``. + +``bias`` + An event bias. More events are generated where the bias is large, and their weights + are lowered to compensate. + +You write the function for a single phase-space point in C, C++, Fortran, CUDA or HIP. +MG7 builds it into a shared library and calls it on batches of events. The library +implements the :doc:`UMAMI interface `, the same interface the +matrix elements use. A generated wrapper takes care of that, so you only write the +function itself. + +Creating a custom function +-------------------------- + +Run ``bin/init_custom_func`` in the process directory. Without options, it asks which +function and language you want: + +.. code-block:: text + + $ ./bin/init_custom_func + Which custom function do you want to create? + 1. cut a cut, 1 keeps the event and 0 rejects it + 2. scale the renormalization and factorization scale + 3. pdf a parton density x f(x, Q) + 4. bias an event bias for the unweighting + Which language do you want to write it in? + 1. c C + 2. cpp C++ + 3. fortran Fortran + 4. cuda CUDA, for runs on NVIDIA GPUs + 5. hip HIP, for runs on AMD GPUs + +The same choices are available as options: + +.. code-block:: bash + + ./bin/init_custom_func --function cut --language cpp + +``-f`` and ``-l`` are short for ``--function`` and ``--language``. ``-y`` takes the +defaults (``cut`` and ``cpp``) for the options you leave out instead of asking. +``--force`` overwrites an existing directory. ``--no-build`` skips the build. +``--no-set-run-card`` leaves the run card unchanged. + +The script creates a directory ``custom_``: + +``custom_.`` + The function body. This is the only file you edit. + +``wrapper_cpu.c`` or ``wrapper_gpu.cu`` + The UMAMI wrapper. It unpacks the batched inputs and calls your function once per + event. + +``umami.h`` + The UMAMI header. + +``Makefile`` + Builds ``lib/custom_.so``. + +By default, the script builds the default function right away and enables it in +``Cards/run_card.toml``: + +.. code-block:: toml + + [custom_functions] + cut = "lib/custom_cut.so" + +``bin/generate_events`` runs ``make`` in ``custom_`` before each run. Your +changes are therefore picked up without a manual rebuild. The build output goes to +``custom_/compile.log``. + +Function signatures +------------------- + +The cut, scale and bias functions receive the external particles of one event. The +incoming particles come first. The momenta are in the lab frame, in GeV, with components +``(E, px, py, pz)``. The PDG ids are those of the first flavor of the subprocess for cuts +and scales, like the built-in cuts, and those of the sampled flavor for the bias. + +C, C++ and CUDA/HIP: + +.. code-block:: c + + double custom_cut(int particle_count, const int* pdg_ids, const double momenta[][4]); + double custom_scale(int particle_count, const int* pdg_ids, const double momenta[][4]); + double custom_bias(int particle_count, const int* pdg_ids, const double momenta[][4]); + double custom_pdf(int beam_index, int pdg_id, double x, double q); + +In C++, the functions are declared ``extern "C"``. In CUDA and HIP, they are +``__device__`` functions and run on the GPU, one thread per event. + +Fortran: + +.. code-block:: fortran + + function custom_cut(particle_count, pdg_ids, momenta) bind(C, name="custom_cut") result(value) + use iso_c_binding + integer(c_int), value :: particle_count + integer(c_int), intent(in) :: pdg_ids(particle_count) + real(c_double), intent(in) :: momenta(0:3, particle_count) + real(c_double) :: value + +The PDF takes ``integer(c_int), value :: beam_index, pdg_id`` and +``real(c_double), value :: x, q``. + +What the functions return: + +``cut`` + 1 to keep the event, 0 to reject it. The value multiplies the event weight, so values + in between act as a weight factor. + +``scale`` + The scale in GeV, used for the renormalization scale and both factorization scales. + ``scale_factor`` is applied on top. ``fixed_ren_scale`` and ``fixed_fact_scale`` still + take precedence. + +``pdf`` + :math:`x f(x, Q)` for the parton ``pdg_id`` (21 for the gluon) in beam ``beam_index`` + (0 or 1), at momentum fraction ``x`` and factorization scale ``q``. + +``bias`` + The bias factor :math:`b > 0`. The phase space is sampled and unweighted according to + :math:`|w|\,b`. The written events carry the weight divided by :math:`b`. Distributions + filled with the event weights therefore stay unbiased, while regions with a large bias + get more events. The reported cross section does not include the bias. + +Examples +-------- + +A cut requiring :math:`p_T > 150` GeV for both particles of a ``p p > t t~`` process, in +C++: + +.. code-block:: c++ + + #include + + extern "C" double custom_cut(int particle_count, const int* pdg_ids, const double momenta[][4]) { + for (int i = 2; i < 4; ++i) { + if (std::hypot(momenta[i][1], momenta[i][2]) < 150.) { + return 0.; + } + } + return 1.; + } + +:math:`H_T/2` as the scale, in Fortran: + +.. code-block:: fortran + + function custom_scale(particle_count, pdg_ids, momenta) bind(C, name="custom_scale") result(value) + use iso_c_binding + implicit none + integer(c_int), value :: particle_count + integer(c_int), intent(in) :: pdg_ids(particle_count) + real(c_double), intent(in) :: momenta(0:3, particle_count) + real(c_double) :: value + integer :: i + value = 0d0 + do i = 3, particle_count + value = value + sqrt(max(momenta(0, i)**2 - momenta(3, i)**2, 0d0)) + end do + value = value / 2d0 + end function custom_scale + +A bias that populates the :math:`p_T` tail of the first final-state particle, in C: + +.. code-block:: c + + #include + + double custom_bias(int particle_count, const int* pdg_ids, const double momenta[][4]) { + double pt2 = momenta[2][1] * momenta[2][1] + momenta[2][2] * momenta[2][2]; + return fmax(pt2 / (100. * 100.), 1.); + } + +An LHAPDF set called through the LHAPDF C++ API: + +.. code-block:: c++ + + #include + #include "LHAPDF/LHAPDF.h" + + extern "C" double custom_pdf(int beam_index, int pdg_id, double x, double q) { + static LHAPDF::PDF* pdf = LHAPDF::mkPDF("NNPDF40_lo_as_01180", 0); + static std::mutex mutex; + std::lock_guard lock(mutex); + return pdf->xfxQ(pdg_id, x, q); + } + +The function is called from several threads, so the call into LHAPDF is protected by a +lock. To build it, pass the LHAPDF flags to ``make``: + +.. code-block:: bash + + make -C custom_pdf CFLAGS="$(lhapdf-config --cflags)" \ + LIBS="$(lhapdf-config --libs) -Wl,-rpath,$(lhapdf-config --libdir)" + +Building +-------- + +The ``Makefile`` takes the usual ``CC``, ``CXX`` and ``FC`` compilers, with ``gfortran`` +as the default for Fortran. ``OPTFLAGS`` defaults to ``-O3``. Extra compiler flags go into +``CFLAGS`` and extra libraries into ``LIBS``. The CUDA and HIP builds detect the GPU +architecture like the matrix elements, and accept the same ``MADGRAPH_CUDA_ARCHITECTURE`` +and ``MADGRAPH_HIP_ARCHITECTURE`` overrides. + +The language decides the device the function runs on. C, C++ and Fortran functions run on +the CPU, and CUDA and HIP functions run on the GPU. Use the language that matches the +``device`` of the run. + +Limitations +----------- + +- With a custom PDF, :math:`\alpha_s` still comes from the PDF set of ``[beam] pdf``, and + the scale and PDF variations of ``[systematics]`` are disabled. +- With a custom bias, the weighted integration histograms of ``[run] + weighted_histograms`` include the bias. The event weights do not. +- Decays ignore a custom scale and a custom PDF. +- A gridpack stores the custom libraries it was made with. They cannot be changed in the + gridpack. diff --git a/docs/source/index.rst b/docs/source/index.rst index 3c597d70cb..2c5952a259 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -27,6 +27,7 @@ If you need a stable release, use the installation first_steps run_card + custom_functions gridpacks madevent_to_mg7 diff --git a/docs/source/run_card.rst b/docs/source/run_card.rst index 0192a54b98..b0beadb725 100644 --- a/docs/source/run_card.rst +++ b/docs/source/run_card.rst @@ -308,6 +308,28 @@ VEGAS grids adapt the phase-space sampling during the survey. Learn the probabilities of the permutations of identical particles instead of using equal ones. +[custom_functions] +------------------ + +Shared libraries with user-written functions, created with ``bin/init_custom_func``. Paths +are relative to the process directory. An empty string disables the function. See +:doc:`custom_functions` for details. + +``cut`` (``""``) + Custom cut, for example ``"lib/custom_cut.so"``. It is applied on top of the + ``[cuts]``. + +``scale`` (``""``) + Custom scale. It replaces ``dynamical_scale_choice``. + +``pdf`` (``""``) + Custom parton density. It replaces the PDF set, but :math:`\alpha_s` still comes from + ``[beam] pdf``. The systematics are disabled. + +``bias`` (``""``) + Custom event bias. Events are generated with the weight times the bias and written with + the weight divided by the bias. + .. _run-card-multiparticles: [multiparticles] diff --git a/tests/acceptance_tests/test_mg7_custom_functions.py b/tests/acceptance_tests/test_mg7_custom_functions.py new file mode 100644 index 0000000000..9aabac95ae --- /dev/null +++ b/tests/acceptance_tests/test_mg7_custom_functions.py @@ -0,0 +1,381 @@ +################################################################################ +# +# Copyright (c) 2026 The MadGraph7 Development team and Contributors +# +# This file is a part of the MadGraph7 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 MadGraph7 license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""MG7 custom cut, scale, PDF and bias functions ([custom_functions]). + +Each custom function is created with bin/init_custom_func in a copy of one +``p p > t t~`` output, given a body that reproduces a built-in setting, and +compared with a run using that built-in setting: + + * cut: pT > 150 GeV for both tops against ``top-pt.min = 150``, and every + event passes the cut; + * scale: H_T/2 against ``dynamical_scale_choice = half_transverse_mass``, + and the LHE scale of every event is H_T/2 of its momenta; + * pdf: the LHAPDF set of the run card called through the LHAPDF C++ API + against the built-in interpolation of the same set; + * bias: b = max((pT / 100 GeV)^2, 1) keeps the cross section and the + weighted pT tail, while putting more events into the tail; + * every function in every language with a toolchain here builds, loads and + returns its default value when evaluated through madspace. + +Run locally with:: + + ./tests/test_manager.py test_.*custom.*_mg7 -pA -t0 + +``MG7_CUSTOM_EVENTS`` sets the events per run (default 5000). +""" + +from __future__ import absolute_import +from __future__ import division + +import glob +import gzip +import json +import math +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 +from tests.acceptance_tests.test_check_xsec_processes_mg7 import ( + _REFERENCE_PDF, _edit_run_card, _require_mg7_runtime, _tail) + +pjoin = os.path.join + +_PROCESS = 'p p > t t~' +_EVENTS = int(os.environ.get('MG7_CUSTOM_EVENTS', 5000)) +_NSIGMA = 3 +_SEED = 1234 +_TOP_PT_CUT = 150. + +# bodies of the custom functions, in C++ +_CUT_BODY = r''' +#include +extern "C" double custom_cut(int particle_count, const int* pdg_ids, const double momenta[][4]) { + for (int i = 2; i < 4; ++i) { + if (std::hypot(momenta[i][1], momenta[i][2]) < %(cut)r) { + return 0.; + } + } + return 1.; +} +''' % {'cut': _TOP_PT_CUT} + +_SCALE_BODY = r''' +#include +extern "C" double custom_scale(int particle_count, const int* pdg_ids, const double momenta[][4]) { + double ht = 0.; + for (int i = 2; i < particle_count; ++i) { + double e = momenta[i][0], pz = momenta[i][3]; + ht += std::sqrt(std::fmax(e * e - pz * pz, 0.)); + } + return ht / 2.; +} +''' + +_PDF_BODY = r''' +#include +#include "LHAPDF/LHAPDF.h" +extern "C" double custom_pdf(int beam_index, int pdg_id, double x, double q) { + static LHAPDF::PDF* pdf = LHAPDF::mkPDF("%(pdf)s", 0); + static std::mutex mutex; + std::lock_guard lock(mutex); + return pdf->xfxQ(pdg_id, x, q); +} +''' % {'pdf': _REFERENCE_PDF} + +_BIAS_BODY = r''' +#include +extern "C" double custom_bias(int particle_count, const int* pdg_ids, const double momenta[][4]) { + double pt2 = momenta[2][1] * momenta[2][1] + momenta[2][2] * momenta[2][2]; + return std::fmax(pt2 / (100. * 100.), 1.); +} +''' + + +def _set(text, key, value): + """Set ``key = value`` (first occurrence at the start of a line).""" + text, count = re.subn(r'(?m)^%s = .*$' % re.escape(key), + '%s = %s' % (key, value), text, count=1) + assert count == 1, 'no %s in the run card' % key + return text + + +def _read_lhe(path): + """(weights, scales, momenta) of an LHE file; momenta[i] lists the + (E, px, py, pz) of every particle of event i.""" + weights, scales, momenta = [], [], [] + with gzip.open(path, 'rt') as f: + lines = iter(f) + for line in lines: + if not line.startswith(' t t~`` output shared by the tests: fixed scales, the + PDF set of the xsec references, LHE output and a fixed seed. The first + call also generates (and compiles) it with a reference run.""" + _require_mg7_runtime(self) + cls = type(self) + if cls.reference is None: + mg = MGCmd.MasterCmd() + mg.no_notification() + mg.exec_cmd('set automatic_html_opening False --no_save') + mg.exec_cmd('generate %s' % _PROCESS) + mg.exec_cmd('output mg7 %s' % self.base) + toml = pjoin(self.base, 'Cards', 'run_card.toml') + _edit_run_card(toml, _EVENTS, disable_jet_cuts=False) + text = open(toml).read() + text = _set(text, 'output_format', '"lhe"') + text = _set(text, 'seed', _SEED) + text = text.replace('[multiparticles]\n', '[multiparticles]\ntop = [6, -6]\n') + open(toml, 'w').write(text) + cls.reference = self.run_process(self.base) + return self.base + + def copy_process(self, name): + target = pjoin(self.path, name) + shutil.copytree(self.base_process(), target, symlinks=True) + # the copy starts without the events of the reference run + shutil.rmtree(pjoin(target, 'Events'), ignore_errors=True) + os.mkdir(pjoin(target, 'Events')) + return target + + def edit_card(self, proc_dir, edit): + toml = pjoin(proc_dir, 'Cards', 'run_card.toml') + text = open(toml).read() + open(toml, 'w').write(edit(text)) + + def init_custom_function(self, proc_dir, kind, body): + """bin/init_custom_func for ``kind`` in C++, with ``body``; the run + builds it.""" + log = pjoin(proc_dir, 'init_custom_func.log') + with open(log, 'w') as logfh: + ret = subprocess.call( + [sys.executable, pjoin(proc_dir, 'bin', 'init_custom_func'), + '-f', kind, '-l', 'cpp', '-y', '--no-build'], + cwd=proc_dir, stdout=logfh, stderr=subprocess.STDOUT) + self.assertEqual(ret, 0, 'bin/init_custom_func failed:\n%s' % _tail(log)) + with open(pjoin(proc_dir, 'custom_%s' % kind, 'custom_%s.cpp' % kind), 'w') as f: + f.write(body) + + def run_process(self, proc_dir): + """Run bin/generate_events: (cross section, error, LHE path).""" + log = pjoin(proc_dir, 'mg7_gen.log') + with open(log, 'w') as logfh: + ret = subprocess.call( + [sys.executable, pjoin(proc_dir, 'bin', 'generate_events'), '-f'], + cwd=proc_dir, stdout=logfh, stderr=subprocess.STDOUT) + self.assertEqual(ret, 0, 'mg7 generate_events failed:\n%s' % _tail(log)) + infos = sorted(glob.glob(pjoin(proc_dir, 'Events', '*', 'info.json'))) + self.assertTrue(infos, 'no info.json produced:\n%s' % _tail(log)) + with open(infos[-1]) as f: + info = json.load(f)['process'] + lhe = pjoin(os.path.dirname(infos[-1]), 'events.lhe.gz') + self.assertTrue(os.path.exists(lhe), 'no LHE file produced') + return float(info['mean']), float(info['error']), lhe + + def assertCompatible(self, got, ref, what): + """Cross sections (mean, error) agree within _NSIGMA.""" + sigma = math.hypot(got[1], ref[1]) + self.assertLessEqual( + abs(got[0] - ref[0]), _NSIGMA * sigma, + '%s: %.5g +- %.3g pb against %.5g +- %.3g pb' + % (what, got[0], got[1], ref[0], ref[1])) + + def test_custom_cut_mg7(self): + builtin = self.copy_process('CUT_BUILTIN') + self.edit_card(builtin, lambda t: t.replace( + '[cuts]\n', '[cuts]\ntop-pt.min = %r\n' % _TOP_PT_CUT)) + ref = self.run_process(builtin) + + custom = self.copy_process('CUT_CUSTOM') + self.init_custom_function(custom, 'cut', _CUT_BODY) + got = self.run_process(custom) + self.assertCompatible(got, ref, 'custom cut') + + _, _, momenta = _read_lhe(got[2]) + for event in momenta: + for p in event[2:4]: + # up to the precision of the LHE file + self.assertGreaterEqual(_pt(p), _TOP_PT_CUT * (1 - 1e-6)) + + def test_custom_scale_mg7(self): + def dynamic(text): + text = _set(text, 'fixed_ren_scale', 'false') + text = _set(text, 'fixed_fact_scale', 'false') + return _set(text, 'dynamical_scale_choice', '"half_transverse_mass"') + + builtin = self.copy_process('SCALE_BUILTIN') + self.edit_card(builtin, dynamic) + ref = self.run_process(builtin) + + custom = self.copy_process('SCALE_CUSTOM') + # partonic_energy differs from H_T/2, so the custom function must win + self.edit_card(custom, lambda t: _set( + dynamic(t), 'dynamical_scale_choice', '"partonic_energy"')) + self.init_custom_function(custom, 'scale', _SCALE_BODY) + got = self.run_process(custom) + self.assertCompatible(got, ref, 'custom scale') + + _, scales, momenta = _read_lhe(got[2]) + for scale, event in zip(scales, momenta): + ht = sum(math.sqrt(max(p[0] ** 2 - p[3] ** 2, 0.)) for p in event[2:]) + self.assertAlmostEqual(scale / (ht / 2), 1., places=5) + + def test_custom_pdf_mg7(self): + lhapdf_config = misc.resolve_lhapdf().config + if not lhapdf_config: + self.skipTest('lhapdf-config not found') + flags = {} + for flag in ('cflags', 'libs', 'libdir'): + flags[flag] = subprocess.check_output( + [lhapdf_config, '--%s' % flag], text=True).strip() + + custom = self.copy_process('PDF_CUSTOM') + self.init_custom_function(custom, 'pdf', _PDF_BODY) + build = subprocess.run( + ['make', 'CFLAGS=%s' % flags['cflags'], + 'LIBS=%s -Wl,-rpath,%s' % (flags['libs'], flags['libdir'])], + cwd=pjoin(custom, 'custom_pdf'), capture_output=True, text=True) + self.assertEqual(build.returncode, 0, build.stdout + build.stderr) + got = self.run_process(custom) + self.assertCompatible(got, type(self).reference, 'custom pdf') + + def test_custom_bias_mg7(self): + custom = self.copy_process('BIAS_CUSTOM') + ref = type(self).reference + self.init_custom_function(custom, 'bias', _BIAS_BODY) + got = self.run_process(custom) + self.assertCompatible(got, ref, 'biased cross section') + + def tail(lhe, threshold=250.): + weights, _, momenta = _read_lhe(lhe) + n = len(weights) + in_tail = [w for w, event in zip(weights, momenta) + if _pt(event[2]) > threshold] + mean = sum(weights) / n + error = math.sqrt(sum((w - mean) ** 2 for w in weights) / n / (n - 1)) + xsec = sum(in_tail) / n + xsec_error = math.sqrt(sum(w ** 2 for w in in_tail)) / n + return weights, (mean, error), (xsec, xsec_error), len(in_tail) / n + + ref_weights, _, ref_tail, ref_fraction = tail(ref[2]) + weights, mean, tail_xsec, fraction = tail(got[2]) + self.assertGreater(max(weights) / min(weights), 2., + 'the biased events should not have a constant weight') + self.assertCompatible(mean, ref[:2], 'mean biased event weight') + self.assertCompatible(tail_xsec, ref_tail, 'pT > 250 GeV cross section') + self.assertGreater(fraction, 2 * ref_fraction, + 'the bias should put more events into the tail') + + def test_init_custom_func_languages_mg7(self): + ms = type(self).ms + if ms is None or not hasattr(ms, 'CustomFunction'): + self.skipTest('madspace with custom functions unavailable') + import numpy as np + compilers = {'c': 'cc', 'cpp': 'c++', 'fortran': 'gfortran', + 'cuda': 'nvcc', 'hip': 'hipcc'} + from madgraph.iolibs.template_files.mg7 import init_custom_func + + expected = {'cut': 1., 'bias': 1.} + momenta = np.array([[[500., 0., 0., 500.], [500., 0., 0., -500.], + [500., 100., 50., 200.], [500., -100., -50., -200.]], + [[300., 0., 0., 300.], [200., 0., 0., -200.], + [250., 30., 0., 100.], [250., -30., 0., 0.]]]) + tested = [] + for language, compiler in compilers.items(): + if shutil.which(compiler) is None: + continue + me_dir = pjoin(self.path, 'LANG_%s' % language) + os.makedirs(me_dir) + for kind in ('cut', 'scale', 'pdf', 'bias'): + ret = init_custom_func.main(me_dir=me_dir, argv=[ + '-f', kind, '-l', language, '-y', '--build', '--no-set-run-card']) + self.assertEqual(ret, 0, '%s %s did not build' % (language, kind)) + self.assertTrue(os.path.exists(pjoin(me_dir, 'lib', 'custom_%s.so' % kind))) + tested.append(language) + if language in ('cuda', 'hip'): + continue + context = ms.Context(device=ms.cpu_device(), thread_count=1) + for kind in ('cut', 'scale', 'bias'): + api = context.load_matrix_element( + pjoin(me_dir, 'lib', 'custom_%s.so' % kind), '') + function = ms.CustomFunction( + api, getattr(ms.CustomFunction, kind), 4, [[21, 21, 6, -6]]) + values = ms.FunctionRuntime(function.function(), context)(momenta) + for value, event in zip(np.array(values), momenta): + if kind == 'scale': + p = [a + b for a, b in zip(event[0], event[1])] + target = math.sqrt(p[0] ** 2 - p[1] ** 2 - p[2] ** 2 - p[3] ** 2) + else: + target = expected[kind] + self.assertAlmostEqual(value, target, places=9) + api = context.load_matrix_element(pjoin(me_dir, 'lib', 'custom_pdf.so'), '') + pdf = ms.PartonDensity(ms.CustomPdf(api), [21, 2], 0) + values = ms.FunctionRuntime(pdf.function(), context)( + np.array([0.1, 0.5]), np.array([100., 100.])) + for x, row in zip([0.1, 0.5], np.array(values)): + toy = math.sqrt(x) * (1 - x) ** 3 + self.assertAlmostEqual(row[0], 3 * toy, places=12) + self.assertAlmostEqual(row[1], 0.5 * toy, places=12) + if not tested: + self.skipTest('no compiler found') From 3874b5ae9c2a8784668e92b4e05657ebe32fee18 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 12 Oct 2026 13:38:43 +0200 Subject: [PATCH 20/37] docs: list CustomFunction and CustomPdf in the madspace API pages --- docs/source/madspace/cpp-api.rst | 2 ++ docs/source/madspace/python-api.rst | 6 ++++++ 2 files changed, 8 insertions(+) diff --git a/docs/source/madspace/cpp-api.rst b/docs/source/madspace/cpp-api.rst index 64384364ef..3fc8856dce 100644 --- a/docs/source/madspace/cpp-api.rst +++ b/docs/source/madspace/cpp-api.rst @@ -43,6 +43,7 @@ Builders that assemble compute-graph functions for channel weights, energy scale cpp/BatchSampler cpp/BufferUnweighter cpp/ChannelWeightNetwork + cpp/CustomFunction cpp/Cuts cpp/DifferentialCrossSection cpp/DiscreteHistogram @@ -79,6 +80,7 @@ Supporting types shared by the mappings: topologies, cuts, observables and densi :maxdepth: 1 cpp/AlphaSGrid + cpp/CustomPdf cpp/Diagram cpp/PdfGrid cpp/Propagator diff --git a/docs/source/madspace/python-api.rst b/docs/source/madspace/python-api.rst index ec5d0ed19c..424b0b038f 100644 --- a/docs/source/madspace/python-api.rst +++ b/docs/source/madspace/python-api.rst @@ -40,9 +40,11 @@ Builders that assemble compute-graph functions for channel weights, energy scale .. toctree:: :maxdepth: 1 + python/BufferUnweighter python/CachedPdf python/CachedScale python/ChannelWeightNetwork + python/CustomFunction python/CutItem python/Cuts python/DifferentialCrossSection @@ -77,6 +79,7 @@ Supporting types shared by the mappings: topologies, cuts, observables and densi :maxdepth: 1 python/AlphaSGrid + python/CustomPdf python/Decay python/Diagram python/LineRef @@ -177,3 +180,6 @@ Lower-level helpers and containers. python/NamedTypes python/NamedValues + python/get_stream + python/set_stream + python/stream From d1fdf13c35b8a91ce33caff4a04983cc69c7823b Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 12 Oct 2026 16:17:57 +0200 Subject: [PATCH 21/37] madspace: build a MatrixElementApi from a table of UMAMI functions UmamiFunctions holds one std::function per UMAMI entry point. The dlopen path fills it from the library symbols, and Context::add_matrix_element takes it directly, so matrix elements can be implemented without a shared library. --- madspace/include/madspace/driver/context.hpp | 86 +++++++--- madspace/src/driver/context.cpp | 168 ++++++++++--------- 2 files changed, 160 insertions(+), 94 deletions(-) diff --git a/madspace/include/madspace/driver/context.hpp b/madspace/include/madspace/driver/context.hpp index 521cfb4493..764f0c2886 100644 --- a/madspace/include/madspace/driver/context.hpp +++ b/madspace/include/madspace/driver/context.hpp @@ -1,7 +1,9 @@ #pragma once #include +#include #include +#include #include #include @@ -12,6 +14,31 @@ namespace madspace { +/** + * Table of UMAMI entry points, one per function of the UMAMI C interface. + * + * Used to build a @ref MatrixElementApi without a shared library, with @ref + * Context::add_matrix_element. Each entry accepts a plain C function pointer + * or a capturing callable. `get_meta`, `initialize`, `matrix_element` and + * `free` are required. The others are optional: missing key queries report + * `UMAMI_ERROR_NOT_IMPLEMENTED`, and a missing `set_parameter` only fails if + * parameters have to be set. + */ +struct UmamiFunctions { + std::function> get_meta; + std::function> + supported_inputs; + std::function> + required_inputs; + std::function> + supported_outputs; + std::function> initialize; + std::function> set_parameter; + std::function> + matrix_element; + std::function> free; +}; + /** * Loaded UMAMI shared library for one matrix element. * @@ -20,8 +47,9 @@ namespace madspace { * MatrixElement for the compute-graph function built on top of it. A * `MatrixElementApi` dynamically loads the library, resolves its exported * UMAMI entry points, and exposes the metadata queries and the raw @ref call. - * Obtained from @ref Context::load_matrix_element, never constructed - * directly. + * Obtained from @ref Context::load_matrix_element, or from @ref + * Context::add_matrix_element for entry points given as @ref UmamiFunctions, + * never constructed directly. * * **References** * - [1] S. Hageböck et al., "Data-parallel leading-order event generation in @@ -36,7 +64,7 @@ class MatrixElementApi { /// Device the underlying library was built for. DeviceType device_type() const { UmamiDevice dev; - check_umami_status(_get_meta(UMAMI_META_DEVICE, &dev)); + check_umami_status(_functions.get_meta(UMAMI_META_DEVICE, &dev)); switch (dev) { case UMAMI_DEVICE_CPU: return DeviceType::cpu; @@ -51,30 +79,31 @@ class MatrixElementApi { /// Number of external particles. std::size_t particle_count() const { int count; - check_umami_status(_get_meta(UMAMI_META_PARTICLE_COUNT, &count)); + check_umami_status(_functions.get_meta(UMAMI_META_PARTICLE_COUNT, &count)); return count; } /// Number of Feynman diagrams. std::size_t diagram_count() const { int count; - check_umami_status(_get_meta(UMAMI_META_DIAGRAM_COUNT, &count)); + check_umami_status(_functions.get_meta(UMAMI_META_DIAGRAM_COUNT, &count)); return count; } /// Number of helicity configurations. std::size_t helicity_count() const { int count; - check_umami_status(_get_meta(UMAMI_META_HELICITY_COUNT, &count)); + check_umami_status(_functions.get_meta(UMAMI_META_HELICITY_COUNT, &count)); return count; } /// Index of this matrix element in its @ref Context, in load order. std::size_t index() const { return _index; } - /// Path of the loaded UMAMI shared library. + /// Path of the loaded UMAMI shared library, or the name given to @ref + /// Context::add_matrix_element. const std::string& file_name() const { return _file_name; } /// UMAMI input keys this library can accept. std::vector supported_inputs() const { bool const* data; int count; - check_umami_status(_supported_inputs(&data, &count)); + check_umami_status(_functions.supported_inputs(&data, &count)); std::vector result(UMAMI_INPUT_KEY_COUNT, false); for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i) { result[i] = data[i]; @@ -85,7 +114,7 @@ class MatrixElementApi { std::vector required_inputs() const { bool const* data; int count; - check_umami_status(_required_inputs(&data, &count)); + check_umami_status(_functions.required_inputs(&data, &count)); std::vector supported = supported_inputs(); for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i) { if (data[i] && !supported[i]) { @@ -106,7 +135,7 @@ class MatrixElementApi { std::vector supported_outputs() const { bool const* data; int count; - check_umami_status(_supported_outputs(&data, &count)); + check_umami_status(_functions.supported_outputs(&data, &count)); std::vector result(UMAMI_OUTPUT_KEY_COUNT, false); for (int i = 0; i < count && i < UMAMI_OUTPUT_KEY_COUNT; ++i) { result[i] = data[i]; @@ -128,7 +157,7 @@ class MatrixElementApi { UmamiOutputKey const* output_keys, void* const* outputs ) const { - check_umami_status(_matrix_element( + check_umami_status(_functions.matrix_element( handle, count, stride, @@ -154,18 +183,27 @@ class MatrixElementApi { std::size_t index = 0, const std::unordered_map& parameters = {} ); + MatrixElementApi( + UmamiFunctions functions, + const std::string& name, + const std::string& param_card, + ThreadPool& thread_pool, + DevicePtr device, + std::size_t index = 0, + const std::unordered_map& parameters = {} + ); + void init_functions(const std::unordered_map& parameters); + void init_instances( + ThreadPool& thread_pool, + DevicePtr device, + const std::string& param_card, + const std::unordered_map& parameters + ); void check_umami_status(UmamiStatus status) const; [[noreturn]] void throw_error(const std::string& message) const; std::unique_ptr> _shared_lib; - decltype(&umami_get_meta) _get_meta; - decltype(&umami_supported_inputs) _supported_inputs; - decltype(&umami_required_inputs) _required_inputs; - decltype(&umami_supported_outputs) _supported_outputs; - decltype(&umami_initialize) _initialize; - decltype(&umami_set_parameter) _set_parameter; - decltype(&umami_matrix_element) _matrix_element; - decltype(&umami_free) _free; + UmamiFunctions _functions; using InstanceType = std::unique_ptr>; ThreadResource _instances; std::string _file_name; @@ -217,6 +255,16 @@ class Context { const std::string& param_card, const std::unordered_map& parameters = {} ); + /// Add a matrix element from the UMAMI entry points in `functions`, + /// without a shared library. `name` identifies it in error messages, and + /// `param_card` and `parameters` are passed on as in @ref + /// load_matrix_element. + const MatrixElementApi& add_matrix_element( + UmamiFunctions functions, + const std::string& name, + const std::string& param_card = "", + const std::unordered_map& parameters = {} + ); /// Create and zero-initialize a new global named `name`. Tensor define_global( const std::string& name, diff --git a/madspace/src/driver/context.cpp b/madspace/src/driver/context.cpp index f988b8f731..372e359fbf 100644 --- a/madspace/src/driver/context.cpp +++ b/madspace/src/driver/context.cpp @@ -12,7 +12,18 @@ using namespace madspace; using json = nlohmann::json; namespace { -UmamiStatus umami_key_query_not_implemented(bool const**, int*) { return UMAMI_ERROR_NOT_IMPLEMENTED; } +UmamiStatus umami_key_query_not_implemented(bool const**, int*) { + return UMAMI_ERROR_NOT_IMPLEMENTED; +} + +// Fills `function` with the symbol `name` of `lib`, if it exists. +template +void resolve_symbol(void* lib, const char* name, std::function& function) { + void* address = dlsym(lib, name); + if (address != nullptr) { + function = reinterpret_cast(address); + } +} thread_local std::optional current_caller_stream; } // namespace @@ -34,97 +45,84 @@ MatrixElementApi::MatrixElementApi( ); } - _get_meta = reinterpret_cast( - dlsym(_shared_lib.get(), "umami_get_meta") - ); - if (_get_meta == nullptr) { - throw std::runtime_error( - std::format("Did not find symbol umami_get_meta in shared object {}", file) - ); - } + void* lib = _shared_lib.get(); + resolve_symbol(lib, "umami_get_meta", _functions.get_meta); + resolve_symbol(lib, "umami_supported_inputs", _functions.supported_inputs); + resolve_symbol(lib, "umami_required_inputs", _functions.required_inputs); + resolve_symbol(lib, "umami_supported_outputs", _functions.supported_outputs); + resolve_symbol(lib, "umami_initialize", _functions.initialize); + resolve_symbol(lib, "umami_set_parameter", _functions.set_parameter); + resolve_symbol(lib, "umami_matrix_element", _functions.matrix_element); + resolve_symbol(lib, "umami_free", _functions.free); + init_functions(parameters); + init_instances(thread_pool, device, param_card, parameters); +} - // These symbols are optional: implementations may not report which keys they - // support, in which case queries for that information fail with - // UMAMI_ERROR_NOT_IMPLEMENTED instead of the constructor throwing. - _supported_inputs = reinterpret_cast( - dlsym(_shared_lib.get(), "umami_supported_inputs") - ); - if (_supported_inputs == nullptr) { - _supported_inputs = umami_key_query_not_implemented; - } +MatrixElementApi::MatrixElementApi( + UmamiFunctions functions, + const std::string& name, + const std::string& param_card, + ThreadPool& thread_pool, + DevicePtr device, + std::size_t index, + const std::unordered_map& parameters +) : + _functions(std::move(functions)), _file_name(name), _index(index) { + init_functions(parameters); + init_instances(thread_pool, device, param_card, parameters); +} - _required_inputs = reinterpret_cast( - dlsym(_shared_lib.get(), "umami_required_inputs") - ); - if (_required_inputs == nullptr) { - _required_inputs = umami_key_query_not_implemented; +void MatrixElementApi::init_functions( + const std::unordered_map& parameters +) { + auto require = [&](bool present, const char* name) { + if (!present) { + throw_error(std::format("missing UMAMI function {}", name)); + } + }; + require(bool(_functions.get_meta), "umami_get_meta"); + require(bool(_functions.initialize), "umami_initialize"); + require(bool(_functions.matrix_element), "umami_matrix_element"); + require(bool(_functions.free), "umami_free"); + // Implementations may not report which keys they support; queries for that + // information then fail with UMAMI_ERROR_NOT_IMPLEMENTED. + if (!_functions.supported_inputs) { + _functions.supported_inputs = umami_key_query_not_implemented; } - - _supported_outputs = reinterpret_cast( - dlsym(_shared_lib.get(), "umami_supported_outputs") - ); - if (_supported_outputs == nullptr) { - _supported_outputs = umami_key_query_not_implemented; + if (!_functions.required_inputs) { + _functions.required_inputs = umami_key_query_not_implemented; } - - _initialize = reinterpret_cast( - dlsym(_shared_lib.get(), "umami_initialize") - ); - if (_initialize == nullptr) { - throw std::runtime_error( - std::format( - "Did not find symbol umami_initialize in shared object {}", file - ) - ); + if (!_functions.supported_outputs) { + _functions.supported_outputs = umami_key_query_not_implemented; } - - // Optional too, unless parameters have to be set (checked below). - _set_parameter = reinterpret_cast( - dlsym(_shared_lib.get(), "umami_set_parameter") - ); - if (!parameters.empty() && _set_parameter == nullptr) { - throw std::runtime_error( - std::format( - "Did not find symbol umami_set_parameter in shared object {}", file - ) - ); - } - - _matrix_element = reinterpret_cast( - dlsym(_shared_lib.get(), "umami_matrix_element") - ); - if (_matrix_element == nullptr) { - throw std::runtime_error( - std::format( - "Did not find symbol umami_matrix_element in shared object {}", file - ) - ); - } - - _free = - reinterpret_cast(dlsym(_shared_lib.get(), "umami_free")); - if (_free == nullptr) { - throw std::runtime_error( - std::format("Did not find symbol umami_free in shared object {}", file) - ); + if (!parameters.empty()) { + require(bool(_functions.set_parameter), "umami_set_parameter"); } +} +void MatrixElementApi::init_instances( + ThreadPool& thread_pool, + DevicePtr device, + const std::string& param_card, + const std::unordered_map& parameters +) { _instances = ThreadResource( thread_pool, [this, device, param_card, parameters] { device->activate(); void* instance; - check_umami_status(_initialize(&instance, param_card.c_str())); + check_umami_status(_functions.initialize(&instance, param_card.c_str())); for (const auto& [name, value] : parameters) { - UmamiStatus status = _set_parameter(instance, name.c_str(), value, 0.); + UmamiStatus status = + _functions.set_parameter(instance, name.c_str(), value, 0.); if (status != UMAMI_SUCCESS) { - _free(instance); + _functions.free(instance); throw_error(std::format("could not set parameter {}", name)); } } return InstanceType(instance, [this, device](void* proc) { device->activate(); - _free(proc); + _functions.free(proc); }); }, std::nullopt, @@ -160,8 +158,7 @@ void MatrixElementApi::throw_error(const std::string& message) const { ); } -const MatrixElementApi& -Context::load_matrix_element( +const MatrixElementApi& Context::load_matrix_element( const std::string& file, const std::string& param_card, const std::unordered_map& parameters @@ -180,6 +177,27 @@ Context::load_matrix_element( return *_matrix_elements.back().get(); } +const MatrixElementApi& Context::add_matrix_element( + UmamiFunctions functions, + const std::string& name, + const std::string& param_card, + const std::unordered_map& parameters +) { + _param_card_paths.push_back(param_card); + _matrix_elements.push_back( + std::unique_ptr(new MatrixElementApi( + std::move(functions), + name, + param_card, + *_thread_pool, + _device, + _matrix_elements.size(), + parameters + )) + ); + return *_matrix_elements.back().get(); +} + Tensor Context::define_global( const std::string& name, DataType dtype, const SizeVec& shape, bool requires_grad ) { From 6b232c5612a58b1803c14d3c290fe8ba7aa038b6 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 12 Oct 2026 16:22:14 +0200 Subject: [PATCH 22/37] madspace: release the GIL while runtimes and generators execute Worker threads can then call back into Python. The abort check and the DLPack deleters of Python-owned tensors take the GIL themselves. --- madspace/src/python/function_runtime.cpp | 17 +++++++++++++---- madspace/src/python/function_runtime.hpp | 22 ++++++++++++++++++++++ madspace/src/python/madspace.cpp | 10 ++++++++++ 3 files changed, 45 insertions(+), 4 deletions(-) diff --git a/madspace/src/python/function_runtime.cpp b/madspace/src/python/function_runtime.cpp index 161c4e5ec5..16169c420c 100644 --- a/madspace/src/python/function_runtime.cpp +++ b/madspace/src/python/function_runtime.cpp @@ -580,6 +580,8 @@ Tensor madspace_py::dlpack_to_tensor( deleter = [managed_ptr]() { auto ptr = static_cast(managed_ptr); if (ptr->deleter) { + // the producer may be freed on a worker thread without the GIL + PythonGuard guard; ptr->deleter(ptr); } }; @@ -587,6 +589,7 @@ Tensor madspace_py::dlpack_to_tensor( deleter = [managed_ptr]() { auto ptr = static_cast(managed_ptr); if (ptr->deleter) { + PythonGuard guard; ptr->deleter(ptr); } }; @@ -618,6 +621,8 @@ std::vector FunctionRuntime::call(std::vector args) { ProducerStream producer_stream(args); auto [inputs, runtime] = check_and_convert_args(args, *this, &_dlpack_version_cache); + // a Python matrix element called from a worker thread needs the GIL + py::gil_scoped_release release; return runtime->run(inputs); } @@ -628,8 +633,10 @@ FunctionRuntime::call_with_grad( ProducerStream producer_stream(args); auto [inputs, runtime] = check_and_convert_args(args, *this, &_dlpack_version_cache); - auto [outputs, loc_grad, eval_grad] = - runtime->run_with_grad(inputs, input_requires_grad); + auto [outputs, loc_grad, eval_grad] = [&] { + py::gil_scoped_release release; + return runtime->run_with_grad(inputs, input_requires_grad); + }(); std::vector> local_grads; for (auto& grad : loc_grad) { if (grad) { @@ -689,8 +696,10 @@ FunctionRuntime::call_backward( } // TODO: checks here Runtime* runtime = get_runtime(*this, expected_device); - auto [ret_in_grads, ret_glob_grads] = - runtime->run_backward(arg_out, arg_locals, eval_grad); + auto [ret_in_grads, ret_glob_grads] = [&] { + py::gil_scoped_release release; + return runtime->run_backward(arg_out, arg_locals, eval_grad); + }(); std::vector> input_grads; for (auto& grad : ret_in_grads) { if (grad) { diff --git a/madspace/src/python/function_runtime.hpp b/madspace/src/python/function_runtime.hpp index 6d189be7d4..e0ee582443 100644 --- a/madspace/src/python/function_runtime.hpp +++ b/madspace/src/python/function_runtime.hpp @@ -13,6 +13,28 @@ using namespace madspace; namespace madspace_py { +/// Holds the GIL while alive, for code that can run on madspace worker threads. +/// Does nothing once the interpreter is finalized. +class PythonGuard { +public: + PythonGuard() : _active(Py_IsInitialized()) { + if (_active) { + _state = PyGILState_Ensure(); + } + } + ~PythonGuard() { + if (_active) { + PyGILState_Release(_state); + } + } + PythonGuard(const PythonGuard&) = delete; + PythonGuard& operator=(const PythonGuard&) = delete; + +private: + bool _active; + PyGILState_STATE _state; +}; + std::tuple dlpack_device(Tensor tensor); py::object tensor_to_dlpack( Tensor tensor, diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 9b267ec9d0..75a328b044 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1863,6 +1863,7 @@ PYBIND11_MODULE(_madspace_py, m) { ); ++i; } + py::gil_scoped_release release; return opt.step(tensors); }, py::arg("inputs"), @@ -2481,6 +2482,7 @@ PYBIND11_MODULE(_madspace_py, m) { .def( "train_step", &MadnisTraining::train_step, + py::call_guard(), py::arg("batch_index"), pydoc::doc("MadnisTraining::train_step") ) @@ -2538,6 +2540,7 @@ PYBIND11_MODULE(_madspace_py, m) { .def( "train", &MultiMadnisTraining::train, + py::call_guard(), pydoc::doc("MultiMadnisTraining::train") ) .def( @@ -3710,17 +3713,20 @@ PYBIND11_MODULE(_madspace_py, m) { .def( "survey", &EventGenerator::survey, + py::call_guard(), py::arg("survey_pass") = 0, pydoc::doc("EventGenerator::survey") ) .def( "generate", &EventGenerator::generate, + py::call_guard(), pydoc::doc("EventGenerator::generate") ) .def( "combine_to_compact_npy", &EventGenerator::combine_to_compact_npy, + py::call_guard(), py::arg("file_name"), py::arg("systematics") = nullptr, py::arg("histograms") = nullptr, @@ -3729,6 +3735,7 @@ PYBIND11_MODULE(_madspace_py, m) { .def( "combine_to_lhe_npy", &EventGenerator::combine_to_lhe_npy, + py::call_guard(), py::arg("file_name"), py::arg("lhe_completer"), py::arg("systematics") = nullptr, @@ -3738,6 +3745,7 @@ PYBIND11_MODULE(_madspace_py, m) { .def( "combine_to_lhe", &EventGenerator::combine_to_lhe, + py::call_guard(), py::arg("file_name"), py::arg("lhe_completer"), py::arg_v("meta", LHEMeta{}, "LHEMeta()"), @@ -3838,6 +3846,8 @@ PYBIND11_MODULE(_madspace_py, m) { ); auto abort_check_function = [] { + // the long-running calls release the GIL + py::gil_scoped_acquire acquire; if (PyErr_CheckSignals() != 0) { throw py::error_already_set(); } From 3c31f82f9cbba48e2153860d6a01929d1f9d700a Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 12 Oct 2026 16:27:09 +0200 Subject: [PATCH 23/37] madspace: PythonFunction, a UMAMI implementation in Python ms.PythonFunction calls a Python function on numpy arrays or torch tensors with selectable UMAMI inputs and outputs, and Context.add_matrix_element adds it to a context. The inputs are zero-copy views of the madspace buffers. The results are checked for type, dtype, shape and device before they are copied back. torch functions run on the stream of the GPU runtime. --- madspace/CMakeLists.txt | 2 + madspace/madspace/_madspace_py_loader.py | 1 + madspace/madspace/python_function.py | 247 +++++++++++++++++++++++ madspace/src/python/madspace.cpp | 57 ++++++ madspace/src/python/python_function.cpp | 223 ++++++++++++++++++++ madspace/src/python/python_function.hpp | 37 ++++ madspace/tests/test_python_function.py | 160 +++++++++++++++ 7 files changed, 727 insertions(+) create mode 100644 madspace/madspace/python_function.py create mode 100644 madspace/src/python/python_function.cpp create mode 100644 madspace/src/python/python_function.hpp create mode 100644 madspace/tests/test_python_function.py diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index d96db0a477..5aad0c520f 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -534,6 +534,8 @@ python_add_library( ${MADSPACE_GENERATED_DIR}/python/docstrings.hpp src/python/function_runtime.cpp src/python/function_runtime.hpp + src/python/python_function.cpp + src/python/python_function.hpp WITH_SOABI ) target_include_directories(_madspace_py PRIVATE ${MADSPACE_GENERATED_DIR}/python) diff --git a/madspace/madspace/_madspace_py_loader.py b/madspace/madspace/_madspace_py_loader.py index 5c8744b981..b68cff15f7 100644 --- a/madspace/madspace/_madspace_py_loader.py +++ b/madspace/madspace/_madspace_py_loader.py @@ -16,6 +16,7 @@ ) from ._madspace_py import * +from .python_function import PythonFunction @contextlib.contextmanager diff --git a/madspace/madspace/python_function.py b/madspace/madspace/python_function.py new file mode 100644 index 0000000000..4d0b8253fc --- /dev/null +++ b/madspace/madspace/python_function.py @@ -0,0 +1,247 @@ +import contextlib +import ctypes + +from ._madspace_py import ( + _UMAMI_INPUT_KEYS, + _UMAMI_OUTPUT_KEYS, + Context, + DataType, + _tensor_view, +) + +# input name: (data type, shape after the batch dimension); "n" is the particle count +_INPUT_SHAPES = { + "momenta": (DataType.float, ("n", 4)), + "alpha_s": (DataType.float, ()), + "flavor_index": (DataType.int, ()), + "random_color": (DataType.float, ()), + "random_helicity": (DataType.float, ()), + "random_diagram": (DataType.float, ()), + "helicity_index": (DataType.int, ()), + "diagram_index": (DataType.int, ()), + "channel_index": (DataType.int, ()), + "particle_count": (DataType.int, ()), + "pdg_ids": (DataType.int, ("n",)), + "momentum_fraction": (DataType.float, ()), + "scale": (DataType.float, ()), + "pdg_id": (DataType.int, ()), + "beam_index": (DataType.int, ()), +} +_OUTPUT_TYPES = { + "matrix_element": DataType.float, + "color_index": DataType.int, + "helicity_index": DataType.int, + "diagram_index": DataType.int, + "cut": DataType.float, + "scale": DataType.float, + "pdf": DataType.float, + "bias": DataType.float, +} + + +class PythonFunction: + """ + A UMAMI implementation that calls a Python function on numpy arrays or torch + tensors, for example a custom cut. Add it to a context with + ``Context.add_matrix_element``, then use the returned ``MatrixElementApi`` like a + loaded library. + + ``function`` is called once per batch, with one array per entry of ``inputs``, + in that order. The arrays have the batch size as first dimension: + ``momenta`` has shape ``(batch, n, 4)``, ``pdg_ids`` has shape ``(batch, n)``, + and all other inputs have shape ``(batch,)``. Floats are 64-bit, integers 32-bit. + The function returns one array of shape ``(batch,)`` per entry of ``outputs``, a + single array for one output and a tuple for several. + + The arrays are views of madspace's buffers on the device of the context, so they + must not be kept after the call. numpy functions only run on the CPU, torch + functions also on CUDA and HIP GPUs, on the stream of the madspace runtime. + + Args: + function: The Python function. + inputs: Names of the UMAMI inputs passed to ``function``. + outputs: Names of the UMAMI outputs ``function`` returns. + framework: ``"numpy"`` or ``"torch"``. + particle_count: The particle count ``n`` for ``momenta`` and ``pdg_ids``, if + the caller does not pass the ``particle_count`` input. + name: Name used in error messages, by default the function's name. + """ + + def __init__( + self, + function, + inputs, + outputs, + framework="numpy", + particle_count=None, + name=None, + ): + if framework not in ("numpy", "torch"): + raise ValueError(f"framework must be 'numpy' or 'torch', not {framework!r}") + for key in inputs: + if key not in _INPUT_SHAPES: + raise ValueError( + f"unknown input {key!r}, expected one of {', '.join(_INPUT_SHAPES)}" + ) + for key in outputs: + if key not in _OUTPUT_TYPES: + raise ValueError( + f"unknown output {key!r}, expected one of {', '.join(_OUTPUT_TYPES)}" + ) + if len(outputs) == 0: + raise ValueError("at least one output is needed") + self.function = function + self.inputs = list(inputs) + self.outputs = list(outputs) + self.framework = framework + self.particle_count = particle_count + self.name = name or getattr(function, "__name__", "python_function") + self._input_keys = [_UMAMI_INPUT_KEYS[key] for key in self.inputs] + self._output_keys = [_UMAMI_OUTPUT_KEYS[key] for key in self.outputs] + self._needs_particle_count = any( + "n" in _INPUT_SHAPES[key][1] for key in self.inputs + ) + if self._needs_particle_count and particle_count is None: + # read from the particle_count input, even if the function does not take it + self._input_keys.append(_UMAMI_INPUT_KEYS["particle_count"]) + if framework == "torch": + import torch + + self._torch = torch + self._from_dlpack = torch.from_dlpack + else: + import numpy + + self._numpy = numpy + self._from_dlpack = numpy.from_dlpack + + def _view(self, address, dtype, shape, stride, offset, device): + return self._from_dlpack( + _tensor_view(address, dtype, shape, stride, offset, device) + ) + + def _numpy_output(self, address, dtype, count, offset): + # numpy.from_dlpack gives read-only arrays, so map the CPU buffer directly + ctype = ctypes.c_double if dtype == DataType.float else ctypes.c_int32 + size = ctypes.sizeof(ctype) + buffer = (ctype * count).from_address(address + offset * size) + return self._numpy.frombuffer(buffer, dtype=self._numpy.dtype(ctype)) + + def _stream(self, device, stream): + if self.framework == "torch" and device != "cpu" and stream is not None: + torch = self._torch + return torch.cuda.stream(torch.cuda.ExternalStream(stream)) + return contextlib.nullcontext() + + def _call(self, device, count, stride, offset, inputs, outputs, stream): + """Called by the UMAMI implementation with the raw buffer addresses.""" + with self._stream(device, stream): + n = self.particle_count + if self._needs_particle_count and n is None: + counts = self._view( + inputs[_UMAMI_INPUT_KEYS["particle_count"]], + DataType.int, + [count], + [1], + offset, + device, + ) + n = int(counts[0]) + args = [] + for key in self.inputs: + dtype, shape = _INPUT_SHAPES[key] + shape = [count] + [n if size == "n" else size for size in shape] + # the batch dimension is contiguous, the others are strided by it + strides = [1] + [ + stride * _prod(shape[1:i]) for i in range(1, len(shape)) + ] + args.append( + self._view( + inputs[_UMAMI_INPUT_KEYS[key]], + dtype, + shape, + strides, + offset, + device, + ) + ) + results = self.function(*args) + if len(self.outputs) == 1: + results = (results,) + elif not isinstance(results, (tuple, list)) or len(results) != len( + self.outputs + ): + raise TypeError( + f"{self.name} must return a tuple of {len(self.outputs)} arrays " + f"({', '.join(self.outputs)})" + ) + for key, result in zip(self.outputs, results): + address = outputs.get(_UMAMI_OUTPUT_KEYS[key]) + if address is None: + continue + dtype = _OUTPUT_TYPES[key] + self._check(key, result, dtype, count, device) + if self.framework == "torch": + self._view(address, dtype, [count], [1], offset, device).copy_( + result + ) + else: + self._numpy_output(address, dtype, count, offset)[:] = result + + def _check(self, key, result, dtype, count, device): + prefix = f"output {key!r} of {self.name}" + if self.framework == "torch": + torch = self._torch + if not isinstance(result, torch.Tensor): + raise TypeError( + f"{prefix} must be a torch.Tensor, not {type(result).__name__}" + ) + expected = torch.float64 if dtype == DataType.float else torch.int32 + result_device = "cpu" if result.device.type == "cpu" else "gpu" + if result_device != ("cpu" if device == "cpu" else "gpu"): + raise ValueError( + f"{prefix} must be on the {device} device, not {result.device}" + ) + else: + numpy = self._numpy + if not isinstance(result, numpy.ndarray): + raise TypeError( + f"{prefix} must be a numpy.ndarray, not {type(result).__name__}" + ) + expected = numpy.dtype( + numpy.float64 if dtype == DataType.float else numpy.int32 + ) + if result.dtype != expected: + raise TypeError(f"{prefix} must have dtype {expected}, not {result.dtype}") + if tuple(result.shape) != (count,): + raise ValueError( + f"{prefix} must have shape ({count},), not {tuple(result.shape)}" + ) + + +def _prod(values): + result = 1 + for value in values: + result *= value + return result + + +def _add_matrix_element(self, function, name=None): + """ + Add a matrix element implemented in Python, see ``PythonFunction``. Returns its + ``MatrixElementApi``. + """ + if not isinstance(function, PythonFunction): + raise TypeError("add_matrix_element takes a madspace.PythonFunction") + required = list(function._input_keys) + return self._add_python_function( + function, + function._input_keys, + required, + function._output_keys, + function.framework, + name or function.name, + ) + + +Context.add_matrix_element = _add_matrix_element diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 75a328b044..0eb9e273b4 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -10,6 +10,7 @@ #include "madspace/driver.hpp" #include "madspace/phasespace.hpp" #include "madspace/util.hpp" +#include "python_function.hpp" namespace py = pybind11; using namespace madspace; @@ -357,6 +358,17 @@ PYBIND11_MODULE(_madspace_py, m) { py::return_value_policy::reference_internal, pydoc::doc("Context::load_matrix_element") ) + .def( + "_add_python_function", + &add_python_function, + py::arg("function"), + py::arg("input_keys"), + py::arg("required_input_keys"), + py::arg("output_keys"), + py::arg("framework"), + py::arg("name"), + py::return_value_policy::reference_internal + ) .def( "define_global", &Context::define_global, @@ -445,6 +457,51 @@ PYBIND11_MODULE(_madspace_py, m) { pydoc::doc("set_caller_stream") ); + m.def( + "_tensor_view", + &tensor_view, + py::arg("address"), + py::arg("dtype"), + py::arg("shape"), + py::arg("stride"), + py::arg("offset"), + py::arg("device") + ); + py::dict umami_input_keys, umami_output_keys; + for (auto [name, key] : std::initializer_list>{ + {"momenta", UMAMI_IN_MOMENTA}, + {"alpha_s", UMAMI_IN_ALPHA_S}, + {"flavor_index", UMAMI_IN_FLAVOR_INDEX}, + {"random_color", UMAMI_IN_RANDOM_COLOR}, + {"random_helicity", UMAMI_IN_RANDOM_HELICITY}, + {"random_diagram", UMAMI_IN_RANDOM_DIAGRAM}, + {"helicity_index", UMAMI_IN_HELICITY_INDEX}, + {"diagram_index", UMAMI_IN_DIAGRAM_INDEX}, + {"channel_index", UMAMI_IN_CHANNEL_INDEX}, + {"particle_count", UMAMI_IN_PARTICLE_COUNT}, + {"pdg_ids", UMAMI_IN_PDG_IDS}, + {"momentum_fraction", UMAMI_IN_MOMENTUM_FRACTION}, + {"scale", UMAMI_IN_SCALE}, + {"pdg_id", UMAMI_IN_PDG_ID}, + {"beam_index", UMAMI_IN_BEAM_INDEX}, + }) { + umami_input_keys[name] = key; + } + for (auto [name, key] : std::initializer_list>{ + {"matrix_element", UMAMI_OUT_MATRIX_ELEMENT}, + {"color_index", UMAMI_OUT_COLOR_INDEX}, + {"helicity_index", UMAMI_OUT_HELICITY_INDEX}, + {"diagram_index", UMAMI_OUT_DIAGRAM_INDEX}, + {"cut", UMAMI_OUT_CUT}, + {"scale", UMAMI_OUT_SCALE}, + {"pdf", UMAMI_OUT_PDF}, + {"bias", UMAMI_OUT_BIAS}, + }) { + umami_output_keys[name] = key; + } + m.attr("_UMAMI_INPUT_KEYS") = umami_input_keys; + m.attr("_UMAMI_OUTPUT_KEYS") = umami_output_keys; + py::classh(m, "FunctionRuntime", py::dynamic_attr()) .def(py::init(), py::arg("function")) .def(py::init(), py::arg("function"), py::arg("context")) diff --git a/madspace/src/python/python_function.cpp b/madspace/src/python/python_function.cpp new file mode 100644 index 0000000000..e2bde8c847 --- /dev/null +++ b/madspace/src/python/python_function.cpp @@ -0,0 +1,223 @@ +#include "python_function.hpp" + +#include +#include +#include + +#include "function_runtime.hpp" + +using namespace madspace; + +namespace madspace_py { + +namespace { + +// Shared by the UMAMI entry points of one Python function. Python calls the +// `_call` method of `function` (a madspace.PythonFunction). +struct PythonFunctionState { + py::object function; + std::string name; + std::string device_name; + UmamiDevice umami_device; + std::array supported_inputs{}; + std::array required_inputs{}; + std::array supported_outputs{}; + + ~PythonFunctionState() { + if (Py_IsInitialized()) { + PythonGuard guard; + function = py::object(); + } else { + function.release(); + } + } +}; + +} // namespace + +Tensor tensor_view( + std::uintptr_t address, + DataType dtype, + const std::vector& shape, + const std::vector& stride, + std::size_t offset, + const std::string& device_name +) { + DevicePtr device; + if (device_name == "cpu") { + device = cpu_device(); + } else if (device_name == "cuda") { + device = cuda_device(); + } else if (device_name == "hip") { + device = hip_device(); + } else { + throw std::invalid_argument(std::format("unknown device {}", device_name)); + } + std::size_t item_size = + dtype == DataType::dt_float ? sizeof(double) : sizeof(me_int_t); + void* data = reinterpret_cast(address + offset * item_size); + return Tensor(dtype, Sizes(shape), Sizes(stride), device, data, [] {}); +} + +const MatrixElementApi& add_python_function( + Context& context, + py::object function, + const std::vector& input_keys, + const std::vector& required_input_keys, + const std::vector& output_keys, + const std::string& framework, + const std::string& name +) { + auto state = std::make_shared(); + state->function = function; + state->name = name; + switch (context.device()->device_type()) { + case DeviceType::cpu: + state->umami_device = UMAMI_DEVICE_CPU; + state->device_name = "cpu"; + break; + case DeviceType::cuda: + state->umami_device = UMAMI_DEVICE_CUDA; + state->device_name = "cuda"; + break; + case DeviceType::hip: + state->umami_device = UMAMI_DEVICE_HIP; + state->device_name = "hip"; + break; + } + if (framework == "numpy" && state->umami_device != UMAMI_DEVICE_CPU) { + throw std::invalid_argument( + std::format( + "Python function {}: numpy functions run on the CPU only, use torch " + "for GPU runs", + name + ) + ); + } + for (int key : input_keys) { + if (key < 0 || key >= UMAMI_INPUT_KEY_COUNT) { + throw std::invalid_argument("invalid UMAMI input key"); + } + state->supported_inputs[key] = true; + } + for (int key : required_input_keys) { + if (key < 0 || key >= UMAMI_INPUT_KEY_COUNT || !state->supported_inputs[key]) { + throw std::invalid_argument("required input key is not supported"); + } + state->required_inputs[key] = true; + } + for (int key : output_keys) { + if (key < 0 || key >= UMAMI_OUTPUT_KEY_COUNT || key == UMAMI_OUT_GPU_STREAM) { + throw std::invalid_argument("invalid UMAMI output key"); + } + state->supported_outputs[key] = true; + } + + UmamiFunctions functions; + functions.get_meta = [state](UmamiMetaKey meta_key, void* result) { + if (meta_key == UMAMI_META_DEVICE) { + *static_cast(result) = state->umami_device; + return UMAMI_SUCCESS; + } + return UMAMI_ERROR_UNSUPPORTED_META; + }; + functions.supported_inputs = [state](bool const** supported, int* count) { + *supported = state->supported_inputs.data(); + *count = UMAMI_INPUT_KEY_COUNT; + return UMAMI_SUCCESS; + }; + functions.required_inputs = [state](bool const** required, int* count) { + *required = state->required_inputs.data(); + *count = UMAMI_INPUT_KEY_COUNT; + return UMAMI_SUCCESS; + }; + functions.supported_outputs = [state](bool const** supported, int* count) { + *supported = state->supported_outputs.data(); + *count = UMAMI_OUTPUT_KEY_COUNT; + return UMAMI_SUCCESS; + }; + functions.initialize = [state](UmamiHandle* handle, char const*) { + *handle = state.get(); + return UMAMI_SUCCESS; + }; + functions.free = [](UmamiHandle) { return UMAMI_SUCCESS; }; + functions.matrix_element = + [state]( + UmamiHandle, + std::size_t count, + std::size_t stride, + std::size_t offset, + std::size_t input_count, + UmamiInputKey const* input_keys, + void const* const* inputs, + std::size_t output_count, + UmamiOutputKey const* output_keys, + void* const* outputs + ) { + std::array has_input{}; + for (std::size_t i = 0; i < input_count; ++i) { + if (input_keys[i] < 0 || input_keys[i] >= UMAMI_INPUT_KEY_COUNT || + !state->supported_inputs[input_keys[i]]) { + return UMAMI_ERROR_UNSUPPORTED_INPUT; + } + has_input[input_keys[i]] = true; + } + for (int key = 0; key < UMAMI_INPUT_KEY_COUNT; ++key) { + if (state->required_inputs[key] && !has_input[key]) { + return UMAMI_ERROR_MISSING_INPUT; + } + } + std::optional stream; + for (std::size_t i = 0; i < output_count; ++i) { + if (output_keys[i] == UMAMI_OUT_GPU_STREAM) { + // the GPU runtime passes the stream handle itself + stream = reinterpret_cast(outputs[i]); + } else if (output_keys[i] < 0 || + output_keys[i] >= UMAMI_OUTPUT_KEY_COUNT || + !state->supported_outputs[output_keys[i]]) { + return UMAMI_ERROR_UNSUPPORTED_OUTPUT; + } + } + if (count == 0) { + return UMAMI_SUCCESS; + } + + std::string error; + { + PythonGuard guard; + try { + py::dict input_dict, output_dict; + for (std::size_t i = 0; i < input_count; ++i) { + input_dict[py::int_(static_cast(input_keys[i]))] = + py::int_(reinterpret_cast(inputs[i])); + } + for (std::size_t i = 0; i < output_count; ++i) { + if (output_keys[i] != UMAMI_OUT_GPU_STREAM) { + output_dict[py::int_(static_cast(output_keys[i]))] = + py::int_(reinterpret_cast(outputs[i])); + } + } + state->function.attr("_call")( + state->device_name, + count, + stride, + offset, + input_dict, + output_dict, + stream ? py::object(py::int_(*stream)) : py::object(py::none()) + ); + } catch (py::error_already_set& e) { + error = e.what(); + } + } + if (!error.empty()) { + throw std::runtime_error( + std::format("Error in Python function {}: {}", state->name, error) + ); + } + return UMAMI_SUCCESS; + }; + return context.add_matrix_element(std::move(functions), name); +} + +} // namespace madspace_py diff --git a/madspace/src/python/python_function.hpp b/madspace/src/python/python_function.hpp new file mode 100644 index 0000000000..fe76343c60 --- /dev/null +++ b/madspace/src/python/python_function.hpp @@ -0,0 +1,37 @@ +#pragma once + +#include +#include +#include +#include + +#include "madspace/driver/context.hpp" + +namespace py = pybind11; + +namespace madspace_py { + +/// Non-owning view of `shape` and element `stride` onto the memory at +/// `address`, starting `offset` elements in. +madspace::Tensor tensor_view( + std::uintptr_t address, + madspace::DataType dtype, + const std::vector& shape, + const std::vector& stride, + std::size_t offset, + const std::string& device_name +); + +/// Add a UMAMI matrix element to `context` that calls `function._call` with +/// the raw UMAMI buffers, see madspace.PythonFunction. +const madspace::MatrixElementApi& add_python_function( + madspace::Context& context, + py::object function, + const std::vector& input_keys, + const std::vector& required_input_keys, + const std::vector& output_keys, + const std::string& framework, + const std::string& name +); + +} // namespace madspace_py diff --git a/madspace/tests/test_python_function.py b/madspace/tests/test_python_function.py new file mode 100644 index 0000000000..4f290637d5 --- /dev/null +++ b/madspace/tests/test_python_function.py @@ -0,0 +1,160 @@ +import numpy as np +import pytest + +import madspace as ms + +PDG_IDS = [21, 21, 6, -6] + + +@pytest.fixture +def context(): + # several threads, so the function is called from worker threads + return ms.Context(device=ms.cpu_device(), thread_count=4) + + +@pytest.fixture +def momenta(): + rng = np.random.default_rng(1234) + return rng.normal(size=(5000, 4, 4)) * 100.0 + + +def pt(momenta, index): + return np.hypot(momenta[:, index, 1], momenta[:, index, 2]) + + +def run_custom(context, function, kind, momenta, framework="numpy"): + api = context.add_matrix_element( + ms.PythonFunction(function, ["momenta", "pdg_ids"], [kind], framework=framework) + ) + generator = ms.CustomFunction(api, getattr(ms.CustomFunction, kind), 4, [PDG_IDS]) + return ms.FunctionRuntime(generator.function(), context)(momenta) + + +def test_numpy_cut(context, momenta): + def cut(p, pdg_ids): + assert p.shape == (len(p), 4, 4) + assert (pdg_ids == PDG_IDS).all() + return (pt(p, 2) > 100.0).astype(np.float64) + + values = run_custom(context, cut, "cut", momenta) + assert np.array_equal(np.asarray(values), (pt(momenta, 2) > 100.0).astype(float)) + + +def test_numpy_scale(context, momenta): + def scale(p, pdg_ids): + return (pt(p, 2) + pt(p, 3)) / 2 + + values = run_custom(context, scale, "scale", momenta) + assert np.allclose(values, (pt(momenta, 2) + pt(momenta, 3)) / 2) + + +def test_numpy_pdf(context): + def pdf(beam_index, pdg_id, x, q): + return np.where(pdg_id == 21, 3.0, 0.5) * np.sqrt(x) * (1 - x) ** 3 + beam_index + + api = context.add_matrix_element( + ms.PythonFunction( + pdf, ["beam_index", "pdg_id", "momentum_fraction", "scale"], ["pdf"] + ) + ) + x = np.array([0.1, 0.5, 0.9]) + for beam_index in (0, 1): + density = ms.PartonDensity(ms.CustomPdf(api), [21, 2], beam_index) + values = np.asarray( + ms.FunctionRuntime(density.function(), context)(x, np.full(3, 100.0)) + ) + toy = np.sqrt(x) * (1 - x) ** 3 + assert np.allclose(values[:, 0], 3 * toy + beam_index) + assert np.allclose(values[:, 1], 0.5 * toy + beam_index) + + +def test_multiple_outputs_fixed_particle_count(context, momenta): + def matrix_element(p): + return p[:, 2, 0] ** 2, np.full(len(p), 3, dtype=np.int32) + + api = context.add_matrix_element( + ms.PythonFunction( + matrix_element, + ["momenta"], + ["matrix_element", "helicity_index"], + particle_count=4, + ) + ) + generator = ms.MatrixElement( + api.index(), + 4, + [ms.MatrixElement.momenta_in], + [ms.MatrixElement.matrix_element_out, ms.MatrixElement.helicity_index_out], + ) + me, helicity = ms.FunctionRuntime(generator.function(), context)(momenta) + assert np.allclose(me, momenta[:, 2, 0] ** 2) + assert (np.asarray(helicity) == 3).all() + + +@pytest.mark.parametrize( + "result, message", + [ + (lambda p: np.ones(len(p), dtype=np.float32), "must have dtype float64"), + (lambda p: np.ones(3), "must have shape"), + (lambda p: list(np.ones(len(p))), "must be a numpy.ndarray"), + ], +) +def test_result_checks(context, momenta, result, message): + with pytest.raises(RuntimeError, match=message): + run_custom(context, lambda p, pdg_ids: result(p), "cut", momenta) + + +def test_exception_propagates(context, momenta): + def broken(p, pdg_ids): + raise KeyError("something went wrong") + + with pytest.raises( + RuntimeError, match="Error in Python function broken.*something" + ): + run_custom(context, broken, "cut", momenta) + + +def test_invalid_arguments(): + with pytest.raises(ValueError, match="unknown input"): + ms.PythonFunction(lambda: None, ["momentum"], ["cut"]) + with pytest.raises(ValueError, match="unknown output"): + ms.PythonFunction(lambda: None, ["momenta"], ["cuts"]) + with pytest.raises(ValueError, match="framework"): + ms.PythonFunction(lambda: None, ["momenta"], ["cut"], framework="jax") + + +def test_release_inputs_on_worker_threads(context, momenta): + # numpy-backed inputs are freed by madspace, possibly on a worker thread + def cut(p, pdg_ids): + return np.ones(len(p)) + + api = context.add_matrix_element( + ms.PythonFunction(cut, ["momenta", "pdg_ids"], ["cut"]) + ) + generator = ms.CustomFunction(api, ms.CustomFunction.cut, 4, [PDG_IDS]) + runtime = ms.FunctionRuntime(generator.function(), context) + for _ in range(20): + runtime(momenta.copy()) + runtime.release_inputs() + + +def test_torch_cpu(context, momenta): + torch = pytest.importorskip("torch") + + def scale(p, pdg_ids): + assert isinstance(p, torch.Tensor) + return torch.hypot(p[:, 2:, 1], p[:, 2:, 2]).sum(dim=1) / 2 + + expected = (pt(momenta, 2) + pt(momenta, 3)) / 2 + values = run_custom(context, scale, "scale", momenta, framework="torch") + assert np.allclose(values, expected) + values = run_custom( + context, scale, "scale", torch.from_numpy(momenta), framework="torch" + ) + assert torch.allclose(values, torch.from_numpy(expected)) + + +def test_torch_wrong_type(context, momenta): + pytest.importorskip("torch") + with pytest.raises(RuntimeError, match="must be a torch.Tensor"): + run_custom(context, lambda p, ids: np.ones(len(p)), "cut", momenta, "torch") From b6ede35f11bf58752a26b5e9461c836054c18e4f Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 12 Oct 2026 16:48:16 +0200 Subject: [PATCH 24/37] madspace: stage python_function.py for the .pyi stub generation _madspace_py_loader imports it, so pybind11-stubgen failed to import the staged package in the ENABLE_DOCS builds of CI. --- madspace/generate_pyi.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/madspace/generate_pyi.py b/madspace/generate_pyi.py index c1581c1d6b..2663d9363a 100644 --- a/madspace/generate_pyi.py +++ b/madspace/generate_pyi.py @@ -31,7 +31,7 @@ def main(): shutil.rmtree(stage_pkg) (stage_pkg / "lib").mkdir(parents=True) - for name in ("__init__.py", "_madspace_py_loader.py"): + for name in ("__init__.py", "_madspace_py_loader.py", "python_function.py"): shutil.copy2(args.package_dir / name, stage_pkg / name) shutil.copy2(args.module, stage_pkg / args.module.name) shutil.copy2(args.lib, stage_pkg / "lib" / args.lib.name) From bd8150af86bbc634b9126c9f96763f8c309f9367 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 12 Oct 2026 16:48:19 +0200 Subject: [PATCH 25/37] mg7: custom functions in Python with numpy or torch [custom_functions] takes a .so library or a .py file. custom_function_loader loads both, for bin/generate_events and gridpacks. bin/init_custom_func gains the python-numpy and python-torch languages, which need no build. torch is limited to one thread on the CPU, where madspace already splits the batches. --- .../mg7/custom_function_loader.py | 72 +++++++++++++++++++ .../mg7/custom_functions/bias_numpy.py | 34 +++++++++ .../mg7/custom_functions/bias_torch.py | 35 +++++++++ .../mg7/custom_functions/cut_numpy.py | 32 +++++++++ .../mg7/custom_functions/cut_torch.py | 33 +++++++++ .../mg7/custom_functions/pdf_numpy.py | 37 ++++++++++ .../mg7/custom_functions/pdf_torch.py | 31 ++++++++ .../mg7/custom_functions/scale_numpy.py | 35 +++++++++ .../mg7/custom_functions/scale_torch.py | 36 ++++++++++ .../iolibs/template_files/mg7/gridpack.py | 12 +++- .../template_files/mg7/init_custom_func.py | 40 +++++++++-- madgraph/iolibs/template_files/mg7/launch.py | 25 ++++--- .../iolibs/template_files/mg7/run_card.toml | 3 +- madgraph/various/banner.py | 8 +-- 14 files changed, 409 insertions(+), 24 deletions(-) create mode 100644 madgraph/iolibs/template_files/mg7/custom_function_loader.py create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/bias_numpy.py create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/bias_torch.py create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/cut_numpy.py create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/cut_torch.py create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/pdf_numpy.py create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/pdf_torch.py create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/scale_numpy.py create mode 100644 madgraph/iolibs/template_files/mg7/custom_functions/scale_torch.py diff --git a/madgraph/iolibs/template_files/mg7/custom_function_loader.py b/madgraph/iolibs/template_files/mg7/custom_function_loader.py new file mode 100644 index 0000000000..b6b1da919c --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_function_loader.py @@ -0,0 +1,72 @@ +"""Loads the functions of the [custom_functions] run card section. + +Shared by bin/generate_events of a process directory and of a gridpack. A path +ending in .so is a UMAMI library built by bin/init_custom_func, a path ending in +.py a Python module defining custom_ and FRAMEWORK. +""" + +import importlib.util +import logging +import os + +logger = logging.getLogger("madevent") + +# the arguments passed to custom_ in a Python module +PYTHON_INPUTS = { + "cut": ["momenta", "pdg_ids"], + "scale": ["momenta", "pdg_ids"], + "bias": ["momenta", "pdg_ids"], + "pdf": ["beam_index", "pdg_id", "momentum_fraction", "scale"], +} + +_modules = {} + + +def _import_module(path): + """Import the module at ``path`` once per process.""" + path = os.path.abspath(path) + if path not in _modules: + spec = importlib.util.spec_from_file_location( + f"mg7_custom_function_{len(_modules)}", path + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + _modules[path] = module + 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" + ) + + 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}") + framework = getattr(module, "FRAMEWORK", "numpy") + if framework not in ("numpy", "torch"): + raise RuntimeError(f"{path}: FRAMEWORK must be 'numpy' or 'torch', not {framework!r}") + if framework == "numpy" and any(device != "cpu" for device in device_types): + raise RuntimeError( + f"{path}: numpy functions run on the CPU only, use FRAMEWORK = 'torch' " + "for GPU runs" + ) + if framework == "torch" and "cpu" in device_types: + import torch + + # madspace already splits CPU batches over its threads + if torch.get_num_threads() != 1: + logger.info("Limiting torch to one thread per madspace thread") + torch.set_num_threads(1) + + python_function = ms.PythonFunction( + function, PYTHON_INPUTS[kind], [kind], framework=framework, name=name + ) + return [context.add_matrix_element(python_function) for context in contexts] diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/bias_numpy.py b/madgraph/iolibs/template_files/mg7/custom_functions/bias_numpy.py new file mode 100644 index 0000000000..f2e97421b1 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/bias_numpy.py @@ -0,0 +1,34 @@ +""" +Custom event bias for MG7, in Python with numpy. Generated by bin/init_custom_func. + +custom_bias is called once per batch of phase-space points, with one array +entry per point. It returns a float64 array of shape (batch,) with +the bias factor b per event. Events are generated according to |w| b, +and written with the weight divided by b, so physical distributions are +unchanged while regions with large b get more events. + +It runs on the CPU only. Use the python-torch template for GPU runs. + +Parameters: +momenta float64 array of shape (batch, particle_count, 4), the momenta + (E, px, py, pz) of the external particles in GeV, lab frame, + incoming ones first +pdg_ids int32 array of shape (batch, particle_count), their PDG ids + +Enable it in Cards/run_card.toml with +[custom_functions] +bias = "custom_bias/custom_bias.py" +""" + +import numpy as np + +FRAMEWORK = "numpy" + + +def custom_bias(momenta, pdg_ids): + # Example: populate the tail of the pT distribution of the first + # final-state particle, b = (pT / 100 GeV)^2 with a floor of 1. + # pt2 = momenta[:, 2, 1] ** 2 + momenta[:, 2, 2] ** 2 + # return np.maximum(pt2 / 100. ** 2, 1.) + + return np.ones(len(momenta)) diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/bias_torch.py b/madgraph/iolibs/template_files/mg7/custom_functions/bias_torch.py new file mode 100644 index 0000000000..79777caca8 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/bias_torch.py @@ -0,0 +1,35 @@ +""" +Custom event bias for MG7, in Python with PyTorch. Generated by bin/init_custom_func. + +custom_bias is called once per batch of phase-space points, with one array +entry per point. It returns a float64 tensor of shape (batch,) with +the bias factor b per event. Events are generated according to |w| b, +and written with the weight divided by b, so physical distributions are +unchanged while regions with large b get more events. + +It runs on the device of the run, CPU or GPU. On the CPU, MG7 limits torch +to one thread, since it already calls the function from several threads. + +Parameters: +momenta float64 tensor of shape (batch, particle_count, 4), the momenta + (E, px, py, pz) of the external particles in GeV, lab frame, + incoming ones first +pdg_ids int32 tensor of shape (batch, particle_count), their PDG ids + +Enable it in Cards/run_card.toml with +[custom_functions] +bias = "custom_bias/custom_bias.py" +""" + +import torch + +FRAMEWORK = "torch" + + +def custom_bias(momenta, pdg_ids): + # Example: populate the tail of the pT distribution of the first + # final-state particle, b = (pT / 100 GeV)^2 with a floor of 1. + # pt2 = momenta[:, 2, 1] ** 2 + momenta[:, 2, 2] ** 2 + # return torch.clamp(pt2 / 100. ** 2, min=1.) + + return torch.ones(len(momenta), dtype=torch.float64, device=momenta.device) diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/cut_numpy.py b/madgraph/iolibs/template_files/mg7/custom_functions/cut_numpy.py new file mode 100644 index 0000000000..9cc1af1d0e --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/cut_numpy.py @@ -0,0 +1,32 @@ +""" +Custom cut for MG7, in Python with numpy. Generated by bin/init_custom_func. + +custom_cut is called once per batch of phase-space points, with one array +entry per point. It returns a float64 array of shape (batch,) with +the cut value per event: 1 keeps the event, 0 rejects it. Values in +between act as a weight factor. + +It runs on the CPU only. Use the python-torch template for GPU runs. + +Parameters: +momenta float64 array of shape (batch, particle_count, 4), the momenta + (E, px, py, pz) of the external particles in GeV, lab frame, + incoming ones first +pdg_ids int32 array of shape (batch, particle_count), their PDG ids + +Enable it in Cards/run_card.toml with +[custom_functions] +cut = "custom_cut/custom_cut.py" +""" + +import numpy as np + +FRAMEWORK = "numpy" + + +def custom_cut(momenta, pdg_ids): + # Example: require pT > 20 GeV for the first final-state particle. + # pt = np.hypot(momenta[:, 2, 1], momenta[:, 2, 2]) + # return np.where(pt > 20., 1., 0.) + + return np.ones(len(momenta)) diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/cut_torch.py b/madgraph/iolibs/template_files/mg7/custom_functions/cut_torch.py new file mode 100644 index 0000000000..9010d689a1 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/cut_torch.py @@ -0,0 +1,33 @@ +""" +Custom cut for MG7, in Python with PyTorch. Generated by bin/init_custom_func. + +custom_cut is called once per batch of phase-space points, with one array +entry per point. It returns a float64 tensor of shape (batch,) with +the cut value per event: 1 keeps the event, 0 rejects it. Values in +between act as a weight factor. + +It runs on the device of the run, CPU or GPU. On the CPU, MG7 limits torch +to one thread, since it already calls the function from several threads. + +Parameters: +momenta float64 tensor of shape (batch, particle_count, 4), the momenta + (E, px, py, pz) of the external particles in GeV, lab frame, + incoming ones first +pdg_ids int32 tensor of shape (batch, particle_count), their PDG ids + +Enable it in Cards/run_card.toml with +[custom_functions] +cut = "custom_cut/custom_cut.py" +""" + +import torch + +FRAMEWORK = "torch" + + +def custom_cut(momenta, pdg_ids): + # Example: require pT > 20 GeV for the first final-state particle. + # pt = torch.hypot(momenta[:, 2, 1], momenta[:, 2, 2]) + # return (pt > 20.).double() + + return torch.ones(len(momenta), dtype=torch.float64, device=momenta.device) diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/pdf_numpy.py b/madgraph/iolibs/template_files/mg7/custom_functions/pdf_numpy.py new file mode 100644 index 0000000000..309fd73369 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/pdf_numpy.py @@ -0,0 +1,37 @@ +""" +Custom parton density for MG7, in Python with numpy. Generated by bin/init_custom_func. + +custom_pdf is called once per batch of phase-space points, with one array +entry per point. It returns a float64 array of shape (batch,) with +x f(x, q), the momentum fraction times the parton density, per event. +The strong coupling still comes from the PDF set in the run card. + +It runs on the CPU only. Use the python-torch template for GPU runs. + +Parameters: +beam_index int32 array of shape (batch,), the beam, 0 or 1 +pdg_id int32 array of shape (batch,), the PDG id of the parton +x float64 array of shape (batch,), the momentum fraction +q float64 array of shape (batch,), the factorization scale in GeV + +Enable it in Cards/run_card.toml with +[custom_functions] +pdf = "custom_pdf/custom_pdf.py" +""" + +import numpy as np + +# Example: LHAPDF through its Python module. Uncomment these lines and the +# return statement in custom_pdf below. +# import lhapdf +# pdf_set = lhapdf.mkPDF("NNPDF23_lo_as_0130_qed", 0) + +FRAMEWORK = "numpy" + + +def custom_pdf(beam_index, pdg_id, x, q): + # return np.array([pdf_set.xfxQ(int(i), a, b) for i, a, b in zip(pdg_id, x, q)]) + + # default: a toy density, replace it with your own + norm = np.where(pdg_id == 21, 3., 0.5) + return norm * np.sqrt(x) * (1. - x) ** 3 diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/pdf_torch.py b/madgraph/iolibs/template_files/mg7/custom_functions/pdf_torch.py new file mode 100644 index 0000000000..b3d6c3ac1d --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/pdf_torch.py @@ -0,0 +1,31 @@ +""" +Custom parton density for MG7, in Python with PyTorch. Generated by bin/init_custom_func. + +custom_pdf is called once per batch of phase-space points, with one array +entry per point. It returns a float64 tensor of shape (batch,) with +x f(x, q), the momentum fraction times the parton density, per event. +The strong coupling still comes from the PDF set in the run card. + +It runs on the device of the run, CPU or GPU. On the CPU, MG7 limits torch +to one thread, since it already calls the function from several threads. + +Parameters: +beam_index int32 tensor of shape (batch,), the beam, 0 or 1 +pdg_id int32 tensor of shape (batch,), the PDG id of the parton +x float64 tensor of shape (batch,), the momentum fraction +q float64 tensor of shape (batch,), the factorization scale in GeV + +Enable it in Cards/run_card.toml with +[custom_functions] +pdf = "custom_pdf/custom_pdf.py" +""" + +import torch + +FRAMEWORK = "torch" + + +def custom_pdf(beam_index, pdg_id, x, q): + # default: a toy density, replace it with your own + norm = torch.where(pdg_id == 21, 3., 0.5).double() + return norm * torch.sqrt(x) * (1. - x) ** 3 diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/scale_numpy.py b/madgraph/iolibs/template_files/mg7/custom_functions/scale_numpy.py new file mode 100644 index 0000000000..5d8b3bea49 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/scale_numpy.py @@ -0,0 +1,35 @@ +""" +Custom energy scale for MG7, in Python with numpy. Generated by bin/init_custom_func. + +custom_scale is called once per batch of phase-space points, with one array +entry per point. It returns a float64 array of shape (batch,) with +the renormalization and factorization scale in GeV per event. The run +card scale_factor is applied on top, and fixed scales still take precedence. + +It runs on the CPU only. Use the python-torch template for GPU runs. + +Parameters: +momenta float64 array of shape (batch, particle_count, 4), the momenta + (E, px, py, pz) of the external particles in GeV, lab frame, + incoming ones first +pdg_ids int32 array of shape (batch, particle_count), their PDG ids + +Enable it in Cards/run_card.toml with +[custom_functions] +scale = "custom_scale/custom_scale.py" +""" + +import numpy as np + +FRAMEWORK = "numpy" + + +def custom_scale(momenta, pdg_ids): + # Example: H_T / 2, the scalar sum of the final-state transverse masses. + # final = momenta[:, 2:] + # mt = np.sqrt(np.maximum(final[..., 0] ** 2 - final[..., 3] ** 2, 0.)) + # return mt.sum(axis=1) / 2. + + # default: partonic center-of-mass energy sqrt(s_hat) + p = momenta[:, 0] + momenta[:, 1] + return np.sqrt(np.maximum(p[:, 0] ** 2 - (p[:, 1:] ** 2).sum(axis=1), 0.)) diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/scale_torch.py b/madgraph/iolibs/template_files/mg7/custom_functions/scale_torch.py new file mode 100644 index 0000000000..d746a3b6e3 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/scale_torch.py @@ -0,0 +1,36 @@ +""" +Custom energy scale for MG7, in Python with PyTorch. Generated by bin/init_custom_func. + +custom_scale is called once per batch of phase-space points, with one array +entry per point. It returns a float64 tensor of shape (batch,) with +the renormalization and factorization scale in GeV per event. The run +card scale_factor is applied on top, and fixed scales still take precedence. + +It runs on the device of the run, CPU or GPU. On the CPU, MG7 limits torch +to one thread, since it already calls the function from several threads. + +Parameters: +momenta float64 tensor of shape (batch, particle_count, 4), the momenta + (E, px, py, pz) of the external particles in GeV, lab frame, + incoming ones first +pdg_ids int32 tensor of shape (batch, particle_count), their PDG ids + +Enable it in Cards/run_card.toml with +[custom_functions] +scale = "custom_scale/custom_scale.py" +""" + +import torch + +FRAMEWORK = "torch" + + +def custom_scale(momenta, pdg_ids): + # Example: H_T / 2, the scalar sum of the final-state transverse masses. + # final = momenta[:, 2:] + # mt = torch.sqrt(torch.clamp(final[..., 0] ** 2 - final[..., 3] ** 2, min=0.)) + # return mt.sum(dim=1) / 2. + + # default: partonic center-of-mass energy sqrt(s_hat) + p = momenta[:, 0] + momenta[:, 1] + return torch.sqrt(torch.clamp(p[:, 0] ** 2 - (p[:, 1:] ** 2).sum(dim=1), min=0.)) diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 2f120b4ea7..a6cb03a93e 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -28,6 +28,7 @@ sys.path.insert(0, str(_LOCAL_INSTALL_DIR)) import madspace as ms +from custom_function_loader import load_custom_function import glob import json import tomllib @@ -212,10 +213,15 @@ def main() -> None: # gridpack; gridpacks written before they were recorded used the default me_parameters = madspace_data.get("me_parameters", {}) global_dir = os.path.join("data", "globals") - for context, backend in zip(contexts, backends): + for context in contexts: context.load_globals(global_dir) - for path in madspace_data.get("custom_functions", []): - context.load_matrix_element(path, param_card_path) + # custom functions come first, as in the run that made the gridpack + device_types = [b if b in ("cuda", "hip") else "cpu" for b in backends] + for entry in madspace_data.get("custom_functions", []): + # older gridpacks store only the path of a .so library + kind, path = entry if isinstance(entry, list) else ("", entry) + load_custom_function(ms, contexts, device_types, kind, path) + for context, backend in zip(contexts, backends): for me_path in madspace_data["matrix_elements"]: context.load_matrix_element( me_path.format(device=backend), param_card_path, me_parameters diff --git a/madgraph/iolibs/template_files/mg7/init_custom_func.py b/madgraph/iolibs/template_files/mg7/init_custom_func.py index af50527977..9a40959056 100644 --- a/madgraph/iolibs/template_files/mg7/init_custom_func.py +++ b/madgraph/iolibs/template_files/mg7/init_custom_func.py @@ -42,6 +42,8 @@ "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:]) @@ -69,8 +71,9 @@ 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:8s} {options[name]}") + 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: @@ -98,7 +101,8 @@ def build_parser(): prog="bin/init_custom_func", description=( "Create a skeleton for a custom cut, scale, PDF or bias function in " - "custom_/ and build it into lib/custom_.so." + "custom_/. Compiled languages are built into " + "lib/custom_.so, Python functions need no build." ), ) parser.add_argument("-f", "--function", choices=list(FUNCTIONS)) @@ -121,10 +125,30 @@ def build_parser(): 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 wrapper, the function body, umami.h and the Makefile.""" + """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}", @@ -194,12 +218,16 @@ def main(me_dir=None, argv=None): shutil.rmtree(out_dir) body_file = write_files(out_dir, function, language) - lib_path = os.path.join("lib", f"custom_{function}.so") + 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 do_build is None: + 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 @@ -212,7 +240,7 @@ def main(me_dir=None, argv=None): 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: + 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 diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index bcc6a9b053..2345846cdf 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -38,6 +38,7 @@ 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.various.banner import RunCardMG7 from madgraph.various import misc from madgraph.interface.extended_cmd import Cmd @@ -518,9 +519,9 @@ def init_context(self) -> None: CUSTOM_FUNCTION_KINDS = ("cut", "scale", "pdf", "bias") def init_custom_functions(self) -> None: - """Build and load the libraries of [custom_functions]. 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.""" + """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: @@ -528,18 +529,18 @@ def init_custom_functions(self) -> None: if not path: continue build_dir = f"custom_{kind}" - if os.path.isfile(os.path.join(build_dir, "Makefile")): + 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}: library not found. " - "Create it with bin/init_custom_func." + f"[custom_functions] {kind} = {path!r}: file not found. " + "Create a .so library or a .py file with bin/init_custom_func." ) - apis = [context.load_matrix_element(path, "") for context in self.contexts] + 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} library loaded at different indices") + raise RuntimeError(f"custom {kind} function loaded at different indices") self.custom_functions[kind] = apis[0] - self.custom_function_paths.append(path) + self.custom_function_paths.append([kind, path]) logger.info("Using the custom %s function %s", kind, path) @staticmethod @@ -1983,7 +1984,7 @@ def save_gridpack(self) -> None: shutil.copytree("SubProcesses", os.path.join(gridpack_path, "SubProcesses")) else: shutil.copytree("lib", lib_path) - for path in self.custom_function_paths: + 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) @@ -2056,6 +2057,10 @@ def save_gridpack(self) -> None: shutil.copy( os.path.join(os.path.dirname(__file__), "gridpack.py"), gen_events_file ) + shutil.copy( + os.path.join(os.path.dirname(__file__), "custom_function_loader.py"), + os.path.join(bin_path, "custom_function_loader.py"), + ) os.chmod(gen_events_file, 0o755) data = { diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 23989166dd..b5cb5d34e7 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -115,7 +115,8 @@ cut_decays = %(phasespace.cut_decays)s # apply the [cuts] to decay products of o adaptive_symmetry_sampling = %(phasespace.adaptive_symmetry_sampling)s [custom_functions] -# shared libraries built with bin/init_custom_func, e.g. "lib/custom_cut.so" +# functions created with bin/init_custom_func: a .so library, e.g. "lib/custom_cut.so", +# or a Python file, e.g. "custom_cut/custom_cut.py" # paths are relative to the process directory; "" disables the function cut = %(custom_functions.cut)s scale = %(custom_functions.scale)s # replaces dynamical_scale_choice diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 6a899f7762..d526302ce2 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6886,13 +6886,13 @@ def default_setup(self): # ------------------------- [custom_functions] ----------------- # user functions built with bin/init_custom_func; "" disables them self.add_toml_param('custom_functions', 'cut', "", - comment="custom cut library, e.g. lib/custom_cut.so") + comment="custom cut, a .so library or a .py file, e.g. lib/custom_cut.so") self.add_toml_param('custom_functions', 'scale', "", - comment="custom scale library, replaces dynamical_scale_choice") + comment="custom scale, replaces dynamical_scale_choice") self.add_toml_param('custom_functions', 'pdf', "", - comment="custom PDF library, replaces the PDF set (not alpha_s)") + comment="custom PDF, replaces the PDF set (not alpha_s)") self.add_toml_param('custom_functions', 'bias', "", - comment="custom event bias library") + comment="custom event bias") # ----------------------------- [madnis] ----------------------- self.add_toml_param('madnis', 'enable', False, allowed=["auto", True, False], auto=True, From e6b3346fdcb2fb021ddc3c09229e4b0f03b57537 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 12 Oct 2026 16:48:19 +0200 Subject: [PATCH 26/37] Acceptance tests and documentation for Python custom functions --- docs/source/custom_functions.rst | 118 ++++++++-- docs/source/madspace/cpp-api.rst | 1 + docs/source/madspace/python-api.rst | 1 + docs/source/run_card.rst | 10 +- .../test_mg7_custom_functions.py | 208 +++++++++++++----- 5 files changed, 269 insertions(+), 69 deletions(-) diff --git a/docs/source/custom_functions.rst b/docs/source/custom_functions.rst index 1f16a48104..a9d17bb0da 100644 --- a/docs/source/custom_functions.rst +++ b/docs/source/custom_functions.rst @@ -18,11 +18,12 @@ MG7 can call your own code for four parts of the calculation: An event bias. More events are generated where the bias is large, and their weights are lowered to compensate. -You write the function for a single phase-space point in C, C++, Fortran, CUDA or HIP. -MG7 builds it into a shared library and calls it on batches of events. The library -implements the :doc:`UMAMI interface `, the same interface the -matrix elements use. A generated wrapper takes care of that, so you only write the -function itself. +You write the function either for a single phase-space point in C, C++, Fortran, CUDA or +HIP, or for a whole batch of points in Python with numpy or PyTorch. MG7 builds a compiled +function into a shared library and calls it on batches of events. The library implements +the :doc:`UMAMI interface `, the same interface the matrix elements +use. A generated wrapper takes care of that, so you only write the function itself. A +Python function needs no build. MG7 imports it and calls it through the same interface. Creating a custom function -------------------------- @@ -39,11 +40,13 @@ function and language you want: 3. pdf a parton density x f(x, Q) 4. bias an event bias for the unweighting Which language do you want to write it in? - 1. c C - 2. cpp C++ - 3. fortran Fortran - 4. cuda CUDA, for runs on NVIDIA GPUs - 5. hip HIP, for runs on AMD GPUs + 1. c C + 2. cpp C++ + 3. fortran Fortran + 4. cuda CUDA, for runs on NVIDIA GPUs + 5. hip HIP, for runs on AMD GPUs + 6. python-numpy Python with numpy, for CPU runs, no compilation + 7. python-torch Python with PyTorch, for CPU or GPU runs, no compilation The same choices are available as options: @@ -56,7 +59,7 @@ defaults (``cut`` and ``cpp``) for the options you leave out instead of asking. ``--force`` overwrites an existing directory. ``--no-build`` skips the build. ``--no-set-run-card`` leaves the run card unchanged. -The script creates a directory ``custom_``: +For a compiled language, the script creates a directory ``custom_``: ``custom_.`` The function body. This is the only file you edit. @@ -71,6 +74,9 @@ The script creates a directory ``custom_``: ``Makefile`` Builds ``lib/custom_.so``. +For ``python-numpy`` and ``python-torch``, the directory only holds the function, +``custom_.py``. + By default, the script builds the default function right away and enables it in ``Cards/run_card.toml``: @@ -79,6 +85,13 @@ By default, the script builds the default function right away and enables it in [custom_functions] cut = "lib/custom_cut.so" +A Python function is enabled with the path of its file: + +.. code-block:: toml + + [custom_functions] + cut = "custom_cut/custom_cut.py" + ``bin/generate_events`` runs ``make`` in ``custom_`` before each run. Your changes are therefore picked up without a manual rebuild. The build output goes to ``custom_/compile.log``. @@ -138,6 +151,80 @@ What the functions return: filled with the event weights therefore stay unbiased, while regions with a large bias get more events. The reported cross section does not include the bias. +Python functions +---------------- + +A Python function is called once per batch of phase-space points, with arrays that hold +one entry per point. The file defines the function ``custom_`` and the +framework it uses: + +.. code-block:: python + + import numpy as np + + FRAMEWORK = "numpy" # or "torch" + + def custom_cut(momenta, pdg_ids): + ... + +The cut, scale and bias functions take ``(momenta, pdg_ids)``: + +``momenta`` + float64, shape ``(batch, particle_count, 4)``. The momenta ``(E, px, py, pz)`` of the + external particles in GeV, lab frame, incoming particles first. + +``pdg_ids`` + int32, shape ``(batch, particle_count)``. Their PDG ids, as for the compiled + functions. + +The PDF function takes ``(beam_index, pdg_id, x, q)``, each with shape ``(batch,)``. +``beam_index`` and ``pdg_id`` are int32, ``x`` and ``q`` are float64. + +Every function returns a float64 array of shape ``(batch,)``, with the same meaning as for +the compiled functions. MG7 checks the type, the dtype, the shape and the device of the +result. With numpy the arguments are numpy arrays, and with ``FRAMEWORK = "torch"`` they are +torch tensors. They are views of MG7's own buffers, so they must not be modified or kept +after the call. + +numpy functions run on the CPU only. torch functions run on the device of the run, so the +same file works for CPU and GPU runs. On the GPU the function runs on the CUDA or HIP stream +of the MG7 runtime, and the arguments and the result stay on the GPU. + +The same cut as in the C++ example below, with numpy: + +.. code-block:: python + + import numpy as np + + FRAMEWORK = "numpy" + + def custom_cut(momenta, pdg_ids): + pt = np.hypot(momenta[:, 2:4, 1], momenta[:, 2:4, 2]) + return np.where((pt >= 150.).all(axis=1), 1., 0.) + +:math:`H_T/2` with torch, for CPU or GPU runs: + +.. code-block:: python + + import torch + + FRAMEWORK = "torch" + + def custom_scale(momenta, pdg_ids): + final = momenta[:, 2:] + mt = torch.sqrt(torch.clamp(final[..., 0] ** 2 - final[..., 3] ** 2, min=0.)) + return mt.sum(dim=1) / 2. + +On the CPU, MG7 splits each batch over the threads of ``cpu_thread_pool_size`` and calls +the function from these threads. Python runs one call at a time, but numpy and torch +release this lock inside their array operations, so vectorized functions still run +partly in parallel. MG7 limits torch to one thread on the CPU, so that it does not start +threads of its own on top. numpy uses several threads only for linear algebra, such as +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. + Examples -------- @@ -220,7 +307,8 @@ and ``MADGRAPH_HIP_ARCHITECTURE`` overrides. The language decides the device the function runs on. C, C++ and Fortran functions run on the CPU, and CUDA and HIP functions run on the GPU. Use the language that matches the -``device`` of the run. +``device`` of the run. Python functions are not built. numpy functions run on the CPU, and +torch functions run on the device of the run. Limitations ----------- @@ -230,5 +318,7 @@ Limitations - With a custom bias, the weighted integration histograms of ``[run] weighted_histograms`` include the bias. The event weights do not. - Decays ignore a custom scale and a custom PDF. -- A gridpack stores the custom libraries it was made with. They cannot be changed in the - gridpack. +- A gridpack stores the custom libraries and Python files it was made with. They cannot + be changed in the gridpack. A Python function in a gridpack needs the same Python + packages, such as numpy or torch, where the gridpack runs. +- Python functions with torch support only the first GPU of a machine. diff --git a/docs/source/madspace/cpp-api.rst b/docs/source/madspace/cpp-api.rst index 3fc8856dce..bdd3137dd5 100644 --- a/docs/source/madspace/cpp-api.rst +++ b/docs/source/madspace/cpp-api.rst @@ -165,6 +165,7 @@ The high-level driver plus the tensor and device layer it runs on: integration, cpp/TensorView cpp/ThreadPool cpp/ThreadResource + cpp/UmamiFunctions cpp/UnalignedRef cpp/Variation cpp/VegasGridOptimizer diff --git a/docs/source/madspace/python-api.rst b/docs/source/madspace/python-api.rst index 424b0b038f..2ddbbbbf6c 100644 --- a/docs/source/madspace/python-api.rst +++ b/docs/source/madspace/python-api.rst @@ -180,6 +180,7 @@ Lower-level helpers and containers. python/NamedTypes python/NamedValues + python/PythonFunction python/get_stream python/set_stream python/stream diff --git a/docs/source/run_card.rst b/docs/source/run_card.rst index b0beadb725..11e7825365 100644 --- a/docs/source/run_card.rst +++ b/docs/source/run_card.rst @@ -311,13 +311,13 @@ VEGAS grids adapt the phase-space sampling during the survey. [custom_functions] ------------------ -Shared libraries with user-written functions, created with ``bin/init_custom_func``. Paths -are relative to the process directory. An empty string disables the function. See -:doc:`custom_functions` for details. +User-written functions, created with ``bin/init_custom_func``. Each one is a shared +library (``.so``) or a Python file (``.py``). Paths are relative to the process directory. +An empty string disables the function. See :doc:`custom_functions` for details. ``cut`` (``""``) - Custom cut, for example ``"lib/custom_cut.so"``. It is applied on top of the - ``[cuts]``. + Custom cut, for example ``"lib/custom_cut.so"`` or ``"custom_cut/custom_cut.py"``. It + is applied on top of the ``[cuts]``. ``scale`` (``""``) Custom scale. It replaces ``dynamical_scale_choice``. diff --git a/tests/acceptance_tests/test_mg7_custom_functions.py b/tests/acceptance_tests/test_mg7_custom_functions.py index 9aabac95ae..c15b31e1b1 100644 --- a/tests/acceptance_tests/test_mg7_custom_functions.py +++ b/tests/acceptance_tests/test_mg7_custom_functions.py @@ -108,6 +108,44 @@ } ''' +# the same functions in Python with numpy, and the cut with torch +_PY_HEADER = 'import numpy as np\nFRAMEWORK = "numpy"\n' + +_CUT_NUMPY = _PY_HEADER + ''' +def custom_cut(momenta, pdg_ids): + pt = np.hypot(momenta[:, 2:4, 1], momenta[:, 2:4, 2]) + return np.where((pt >= %(cut)r).all(axis=1), 1., 0.) +''' % {'cut': _TOP_PT_CUT} + +_CUT_TORCH = '''import torch +FRAMEWORK = "torch" + +def custom_cut(momenta, pdg_ids): + pt = torch.hypot(momenta[:, 2:4, 1], momenta[:, 2:4, 2]) + return (pt >= %(cut)r).all(dim=1).double() +''' % {'cut': _TOP_PT_CUT} + +_SCALE_NUMPY = _PY_HEADER + ''' +def custom_scale(momenta, pdg_ids): + final = momenta[:, 2:] + return np.sqrt(np.maximum(final[..., 0] ** 2 - final[..., 3] ** 2, 0.)).sum(axis=1) / 2. +''' + +_PDF_NUMPY = _PY_HEADER + '''import lhapdf +pdf_set = lhapdf.mkPDF("%(pdf)s", 0) + +def custom_pdf(beam_index, pdg_id, x, q): + return np.array([pdf_set.xfxQ(int(i), a, b) for i, a, b in zip(pdg_id, x, q)]) +''' % {'pdf': _REFERENCE_PDF} + +_BIAS_NUMPY = _PY_HEADER + ''' +def custom_bias(momenta, pdg_ids): + pt2 = momenta[:, 2, 1] ** 2 + momenta[:, 2, 2] ** 2 + return np.maximum(pt2 / 100. ** 2, 1.) +''' + +_EXTENSIONS = {'cpp': 'cpp', 'python-numpy': 'py', 'python-torch': 'py'} + def _set(text, key, value): """Set ``key = value`` (first occurrence at the start of a line).""" @@ -204,19 +242,34 @@ def edit_card(self, proc_dir, edit): text = open(toml).read() open(toml, 'w').write(edit(text)) - def init_custom_function(self, proc_dir, kind, body): - """bin/init_custom_func for ``kind`` in C++, with ``body``; the run - builds it.""" + def init_custom_function(self, proc_dir, kind, body, language='cpp'): + """bin/init_custom_func for ``kind`` in ``language``, with ``body``; + the run builds C++ functions.""" log = pjoin(proc_dir, 'init_custom_func.log') with open(log, 'w') as logfh: ret = subprocess.call( [sys.executable, pjoin(proc_dir, 'bin', 'init_custom_func'), - '-f', kind, '-l', 'cpp', '-y', '--no-build'], + '-f', kind, '-l', language, '-y', '--no-build'], cwd=proc_dir, stdout=logfh, stderr=subprocess.STDOUT) self.assertEqual(ret, 0, 'bin/init_custom_func failed:\n%s' % _tail(log)) - with open(pjoin(proc_dir, 'custom_%s' % kind, 'custom_%s.cpp' % kind), 'w') as f: + body_file = 'custom_%s.%s' % (kind, _EXTENSIONS[language]) + with open(pjoin(proc_dir, 'custom_%s' % kind, body_file), 'w') as f: f.write(body) + def builtin_reference(self, name, edit): + """The run with the built-in setting ``edit``, shared by the languages.""" + cls = type(self) + if not hasattr(cls, 'builtin_runs'): + cls.builtin_runs = {} + if name not in cls.builtin_runs: + builtin = self.copy_process(name) + self.edit_card(builtin, edit) + cls.builtin_runs[name] = self.run_process(builtin) + return cls.builtin_runs[name] + + def copy_for(self, kind, language): + return self.copy_process('%s_%s' % (kind.upper(), language.upper().replace('-', '_'))) + def run_process(self, proc_dir): """Run bin/generate_events: (cross section, error, LHE path).""" log = pjoin(proc_dir, 'mg7_gen.log') @@ -241,16 +294,13 @@ def assertCompatible(self, got, ref, what): '%s: %.5g +- %.3g pb against %.5g +- %.3g pb' % (what, got[0], got[1], ref[0], ref[1])) - def test_custom_cut_mg7(self): - builtin = self.copy_process('CUT_BUILTIN') - self.edit_card(builtin, lambda t: t.replace( + def check_cut(self, language, body): + ref = self.builtin_reference('CUT_BUILTIN', lambda t: t.replace( '[cuts]\n', '[cuts]\ntop-pt.min = %r\n' % _TOP_PT_CUT)) - ref = self.run_process(builtin) - - custom = self.copy_process('CUT_CUSTOM') - self.init_custom_function(custom, 'cut', _CUT_BODY) + custom = self.copy_for('cut', language) + self.init_custom_function(custom, 'cut', body, language) got = self.run_process(custom) - self.assertCompatible(got, ref, 'custom cut') + self.assertCompatible(got, ref, 'custom cut in %s' % language) _, _, momenta = _read_lhe(got[2]) for event in momenta: @@ -258,29 +308,45 @@ def test_custom_cut_mg7(self): # up to the precision of the LHE file self.assertGreaterEqual(_pt(p), _TOP_PT_CUT * (1 - 1e-6)) - def test_custom_scale_mg7(self): + def test_custom_cut_mg7(self): + self.check_cut('cpp', _CUT_BODY) + + def test_custom_cut_numpy_mg7(self): + self.check_cut('python-numpy', _CUT_NUMPY) + + def test_custom_cut_torch_mg7(self): + check = subprocess.run([sys.executable, '-c', 'import numpy, torch'], + capture_output=True) + if check.returncode != 0: + self.skipTest('torch is unavailable') + self.check_cut('python-torch', _CUT_TORCH) + + def check_scale(self, language, body): def dynamic(text): text = _set(text, 'fixed_ren_scale', 'false') text = _set(text, 'fixed_fact_scale', 'false') return _set(text, 'dynamical_scale_choice', '"half_transverse_mass"') - builtin = self.copy_process('SCALE_BUILTIN') - self.edit_card(builtin, dynamic) - ref = self.run_process(builtin) - - custom = self.copy_process('SCALE_CUSTOM') + ref = self.builtin_reference('SCALE_BUILTIN', dynamic) + custom = self.copy_for('scale', language) # partonic_energy differs from H_T/2, so the custom function must win self.edit_card(custom, lambda t: _set( dynamic(t), 'dynamical_scale_choice', '"partonic_energy"')) - self.init_custom_function(custom, 'scale', _SCALE_BODY) + self.init_custom_function(custom, 'scale', body, language) got = self.run_process(custom) - self.assertCompatible(got, ref, 'custom scale') + self.assertCompatible(got, ref, 'custom scale in %s' % language) _, scales, momenta = _read_lhe(got[2]) for scale, event in zip(scales, momenta): ht = sum(math.sqrt(max(p[0] ** 2 - p[3] ** 2, 0.)) for p in event[2:]) self.assertAlmostEqual(scale / (ht / 2), 1., places=5) + def test_custom_scale_mg7(self): + self.check_scale('cpp', _SCALE_BODY) + + def test_custom_scale_numpy_mg7(self): + self.check_scale('python-numpy', _SCALE_NUMPY) + def test_custom_pdf_mg7(self): lhapdf_config = misc.resolve_lhapdf().config if not lhapdf_config: @@ -300,12 +366,28 @@ def test_custom_pdf_mg7(self): got = self.run_process(custom) self.assertCompatible(got, type(self).reference, 'custom pdf') + def test_custom_pdf_numpy_mg7(self): + check = subprocess.run([sys.executable, '-c', 'import lhapdf'], + capture_output=True) + if check.returncode != 0: + self.skipTest('the lhapdf Python module is unavailable') + custom = self.copy_for('pdf', 'python-numpy') + self.init_custom_function(custom, 'pdf', _PDF_NUMPY, 'python-numpy') + got = self.run_process(custom) + self.assertCompatible(got, type(self).reference, 'custom pdf in numpy') + def test_custom_bias_mg7(self): - custom = self.copy_process('BIAS_CUSTOM') + self.check_bias('cpp', _BIAS_BODY) + + def test_custom_bias_numpy_mg7(self): + self.check_bias('python-numpy', _BIAS_NUMPY) + + def check_bias(self, language, body): + custom = self.copy_for('bias', language) ref = type(self).reference - self.init_custom_function(custom, 'bias', _BIAS_BODY) + self.init_custom_function(custom, 'bias', body, language) got = self.run_process(custom) - self.assertCompatible(got, ref, 'biased cross section') + self.assertCompatible(got, ref, 'biased cross section in %s' % language) def tail(lhe, threshold=250.): weights, _, momenta = _read_lhe(lhe) @@ -327,20 +409,41 @@ def tail(lhe, threshold=250.): self.assertGreater(fraction, 2 * ref_fraction, 'the bias should put more events into the tail') + def check_defaults(self, ms, context, apis): + """The default functions of bin/init_custom_func evaluated through + madspace: 1 for cut and bias, sqrt(s_hat) for scale, the toy pdf.""" + import numpy as np + momenta = np.array([[[500., 0., 0., 500.], [500., 0., 0., -500.], + [500., 100., 50., 200.], [500., -100., -50., -200.]], + [[300., 0., 0., 300.], [200., 0., 0., -200.], + [250., 30., 0., 100.], [250., -30., 0., 0.]]]) + for kind in ('cut', 'scale', 'bias'): + function = ms.CustomFunction( + apis[kind], getattr(ms.CustomFunction, kind), 4, [[21, 21, 6, -6]]) + values = ms.FunctionRuntime(function.function(), context)(momenta) + for value, event in zip(np.array(values), momenta): + if kind == 'scale': + p = [a + b for a, b in zip(event[0], event[1])] + target = math.sqrt(p[0] ** 2 - p[1] ** 2 - p[2] ** 2 - p[3] ** 2) + else: + target = 1. + self.assertAlmostEqual(value, target, places=9) + pdf = ms.PartonDensity(ms.CustomPdf(apis['pdf']), [21, 2], 0) + values = ms.FunctionRuntime(pdf.function(), context)( + np.array([0.1, 0.5]), np.array([100., 100.])) + for x, row in zip([0.1, 0.5], np.array(values)): + toy = math.sqrt(x) * (1 - x) ** 3 + self.assertAlmostEqual(row[0], 3 * toy, places=12) + self.assertAlmostEqual(row[1], 0.5 * toy, places=12) + def test_init_custom_func_languages_mg7(self): ms = type(self).ms if ms is None or not hasattr(ms, 'CustomFunction'): self.skipTest('madspace with custom functions unavailable') - import numpy as np compilers = {'c': 'cc', 'cpp': 'c++', 'fortran': 'gfortran', 'cuda': 'nvcc', 'hip': 'hipcc'} from madgraph.iolibs.template_files.mg7 import init_custom_func - expected = {'cut': 1., 'bias': 1.} - momenta = np.array([[[500., 0., 0., 500.], [500., 0., 0., -500.], - [500., 100., 50., 200.], [500., -100., -50., -200.]], - [[300., 0., 0., 300.], [200., 0., 0., -200.], - [250., 30., 0., 100.], [250., -30., 0., 0.]]]) tested = [] for language, compiler in compilers.items(): if shutil.which(compiler) is None: @@ -356,26 +459,31 @@ def test_init_custom_func_languages_mg7(self): if language in ('cuda', 'hip'): continue context = ms.Context(device=ms.cpu_device(), thread_count=1) - for kind in ('cut', 'scale', 'bias'): - api = context.load_matrix_element( + self.check_defaults(ms, context, { + kind: context.load_matrix_element( pjoin(me_dir, 'lib', 'custom_%s.so' % kind), '') - function = ms.CustomFunction( - api, getattr(ms.CustomFunction, kind), 4, [[21, 21, 6, -6]]) - values = ms.FunctionRuntime(function.function(), context)(momenta) - for value, event in zip(np.array(values), momenta): - if kind == 'scale': - p = [a + b for a, b in zip(event[0], event[1])] - target = math.sqrt(p[0] ** 2 - p[1] ** 2 - p[2] ** 2 - p[3] ** 2) - else: - target = expected[kind] - self.assertAlmostEqual(value, target, places=9) - api = context.load_matrix_element(pjoin(me_dir, 'lib', 'custom_pdf.so'), '') - pdf = ms.PartonDensity(ms.CustomPdf(api), [21, 2], 0) - values = ms.FunctionRuntime(pdf.function(), context)( - np.array([0.1, 0.5]), np.array([100., 100.])) - for x, row in zip([0.1, 0.5], np.array(values)): - toy = math.sqrt(x) * (1 - x) ** 3 - self.assertAlmostEqual(row[0], 3 * toy, places=12) - self.assertAlmostEqual(row[1], 0.5 * toy, places=12) + for kind in ('cut', 'scale', 'pdf', 'bias')}) + + from madgraph.iolibs.template_files.mg7.custom_function_loader import ( + load_custom_function) + for language in ('python-numpy', 'python-torch'): + if language == 'python-torch': + try: + import torch + except Exception: + continue + me_dir = pjoin(self.path, 'LANG_%s' % language) + os.makedirs(me_dir) + context = ms.Context(device=ms.cpu_device(), thread_count=2) + apis = {} + for kind in ('cut', 'scale', 'pdf', 'bias'): + ret = init_custom_func.main(me_dir=me_dir, argv=[ + '-f', kind, '-l', language, '-y', '--no-set-run-card']) + self.assertEqual(ret, 0, '%s %s failed' % (language, kind)) + path = pjoin(me_dir, 'custom_%s' % kind, 'custom_%s.py' % kind) + self.assertTrue(os.path.exists(path)) + apis[kind] = load_custom_function(ms, [context], ['cpu'], kind, path)[0] + self.check_defaults(ms, context, apis) + tested.append(language) if not tested: self.skipTest('no compiler found') From 13bdbe384ab921d36a71b99ba3251501301347e5 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 11:17:34 +0200 Subject: [PATCH 27/37] 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 28/37] 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 29/37] 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 30/37] 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 31/37] 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 32/37] 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 33/37] 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 34/37] 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 35/37] 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 36/37] 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 37/37] 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)