From 370abe871f9a155e75d5b65a435c98241459a2a7 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 22:23:41 +0200 Subject: [PATCH 01/22] madspace: matrix elements with a fixed channel index A non-negative channel passed to MatrixElement is given to the UMAMI channel input as a constant, which is then not an external input. Integrands and weighting steps can so evaluate one channel of a matrix element that treats its channels as independent processes. --- .../madspace/phasespace/matrix_element.hpp | 19 ++++++++++--- madspace/src/phasespace/matrix_element.cpp | 13 +++++++-- madspace/src/python/madspace.cpp | 9 +++++-- madspace/tests/test_python_function.py | 27 +++++++++++++++++++ 4 files changed, 60 insertions(+), 8 deletions(-) diff --git a/madspace/include/madspace/phasespace/matrix_element.hpp b/madspace/include/madspace/phasespace/matrix_element.hpp index 8d8a5a72f5..e30b442b14 100644 --- a/madspace/include/madspace/phasespace/matrix_element.hpp +++ b/madspace/include/madspace/phasespace/matrix_element.hpp @@ -98,6 +98,8 @@ class MatrixElement : public FunctionGenerator { * not commute, so going there straight from * the lab frame would leave a Wigner rotation * behind and rotate the polarisation axes. + * @param channel If non-negative, `channel_in` is passed this + * fixed value instead of being an input. */ MatrixElement( std::size_t matrix_element_index, @@ -107,7 +109,8 @@ class MatrixElement : public FunctionGenerator { std::size_t diagram_count = 1, bool sample_random_inputs = false, const std::vector& me_frame = {}, - std::size_t incoming_count = 2 + std::size_t incoming_count = 2, + me_int_t channel = -1 ); /** * Take @p matrix_element_index, @p particle_count and @p diagram_count from @@ -116,6 +119,9 @@ class MatrixElement : public FunctionGenerator { * @param inputs Quantities to pass in. * @param outputs Quantities to return. * @param sample_random_inputs See the other constructor. + * @param me_frame See the other constructor. + * @param incoming_count See the other constructor. + * @param channel See the other constructor. */ MatrixElement( const MatrixElementApi& matrix_element_api, @@ -123,7 +129,8 @@ class MatrixElement : public FunctionGenerator { const std::vector& outputs = {matrix_element_out}, bool sample_random_inputs = false, const std::vector& me_frame = {}, - std::size_t incoming_count = 2 + std::size_t incoming_count = 2, + me_int_t channel = -1 ) : MatrixElement( matrix_element_api.index(), @@ -133,7 +140,8 @@ class MatrixElement : public FunctionGenerator { matrix_element_api.diagram_count(), sample_random_inputs, me_frame, - incoming_count + incoming_count, + channel ) {}; /// Index of the matrix element in the plugin. std::size_t matrix_element_index() const { return _matrix_element_index; } @@ -147,8 +155,10 @@ class MatrixElement : public FunctionGenerator { const std::vector& outputs() const { return _outputs; } /// Whether the `random_*` inputs are drawn internally. bool sample_random_inputs() const { return _sample_random_inputs; } + /// The fixed channel passed as `channel_in`, or -1 if it is an input. + me_int_t channel() const { return _channel; } /// The inputs that must be supplied by the caller (excludes the ones drawn - /// internally when @p sample_random_inputs is set). + /// internally when @p sample_random_inputs is set and a fixed channel). std::vector external_inputs() const; /// Per-particle 0/1 selector built from me_frame, empty if no boost is /// applied. @@ -169,6 +179,7 @@ class MatrixElement : public FunctionGenerator { std::vector _inputs; std::vector _outputs; bool _sample_random_inputs; + me_int_t _channel; std::vector _frame_mask; std::vector _reference_mask; }; diff --git a/madspace/src/phasespace/matrix_element.cpp b/madspace/src/phasespace/matrix_element.cpp index c0ef852d4c..5f869136e8 100644 --- a/madspace/src/phasespace/matrix_element.cpp +++ b/madspace/src/phasespace/matrix_element.cpp @@ -40,7 +40,8 @@ MatrixElement::MatrixElement( std::size_t diagram_count, bool sample_random_inputs, const std::vector& me_frame, - std::size_t incoming_count + std::size_t incoming_count, + me_int_t channel ) : FunctionGenerator( "MatrixElement", @@ -81,7 +82,9 @@ MatrixElement::MatrixElement( arg_types.push_back("diagram", batch_int); break; case channel_in: - arg_types.push_back("channel", batch_int); + if (channel < 0) { + arg_types.push_back("channel", batch_int); + } break; case particle_count_in: arg_types.push_back("particle_count", batch_int); @@ -159,6 +162,7 @@ MatrixElement::MatrixElement( _particle_count(particle_count), _diagram_count(diagram_count), _sample_random_inputs(sample_random_inputs), + _channel(channel), _frame_mask( me_frame.empty() ? std::vector() : particle_mask(particle_count, me_frame, "me_frame") @@ -274,6 +278,8 @@ NamedVector MatrixElement::build_function_impl( input == random_diagram_in)) { matrix_args.push_back(random.at(random_index)); ++random_index; + } else if (input == channel_in && _channel >= 0) { + matrix_args.push_back(fb.full({_channel, fb.batch_size(args.values())})); } else if (input == momenta_in && !_frame_mask.empty()) { // Evaluate the matrix element in the frame the run card asked for // (me_frame). Only the matrix element sees the boosted momenta: @@ -346,6 +352,9 @@ NamedVector MatrixElement::build_function_impl( std::vector MatrixElement::external_inputs() const { std::vector ret; for (auto input : _inputs) { + if (input == channel_in && _channel >= 0) { + continue; + } if (!_sample_random_inputs || (input != random_color_in && input != random_helicity_in && input != random_diagram_in)) { diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 42ed9d1c83..8f39f4f12b 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1564,7 +1564,8 @@ PYBIND11_MODULE(_madspace_py, m) { std::size_t, bool, const std::vector&, - std::size_t>(), + std::size_t, + me_int_t>(), py::arg("matrix_element_index"), py::arg("particle_count"), // C++ defaults to {momenta_in} / {matrix_element_out}; kept required @@ -1576,6 +1577,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("sample_random_inputs") = false, py::arg("me_frame") = std::vector{}, py::arg("incoming_count") = 2, + py::arg("channel") = -1, pydoc::doc("MatrixElement::MatrixElement") ) .def( @@ -1585,13 +1587,15 @@ PYBIND11_MODULE(_madspace_py, m) { const std::vector&, bool, const std::vector&, - std::size_t>(), + std::size_t, + me_int_t>(), py::arg("matrix_element_api"), py::arg("inputs"), py::arg("outputs"), py::arg("sample_random_inputs") = false, py::arg("me_frame") = std::vector{}, py::arg("incoming_count") = 2, + py::arg("channel") = -1, pydoc::doc("MatrixElement::MatrixElement#2") ) .def( @@ -1611,6 +1615,7 @@ PYBIND11_MODULE(_madspace_py, m) { ) .def("inputs", &MatrixElement::inputs, pydoc::doc("MatrixElement::inputs")) .def("outputs", &MatrixElement::outputs, pydoc::doc("MatrixElement::outputs")) + .def("channel", &MatrixElement::channel, pydoc::doc("MatrixElement::channel")) .def( "external_inputs", &MatrixElement::external_inputs, diff --git a/madspace/tests/test_python_function.py b/madspace/tests/test_python_function.py index 4f290637d5..924094938a 100644 --- a/madspace/tests/test_python_function.py +++ b/madspace/tests/test_python_function.py @@ -91,6 +91,33 @@ def matrix_element(p): assert (np.asarray(helicity) == 3).all() +def test_fixed_channel(context, momenta): + def matrix_element(p, flavor, channel): + return p[:, 2, 0] * (channel + 10 * flavor) + + api = context.add_matrix_element( + ms.PythonFunction( + matrix_element, + ["momenta", "flavor_index", "channel_index"], + ["matrix_element"], + particle_count=4, + ) + ) + inputs = [ + ms.MatrixElement.momenta_in, + ms.MatrixElement.flavor_in, + ms.MatrixElement.channel_in, + ] + generator = ms.MatrixElement( + api.index(), 4, inputs, [ms.MatrixElement.matrix_element_out], channel=3 + ) + assert generator.channel() == 3 + assert generator.external_inputs() == inputs[:2] + flavor = np.full(len(momenta), 2, dtype=np.int32) + me = ms.FunctionRuntime(generator.function(), context)(momenta, flavor) + assert np.allclose(me, momenta[:, 2, 0] * 23) + + @pytest.mark.parametrize( "result, message", [ From 8b581c20fc7a1a1aba21c6881dab4e894315047b Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 22:24:58 +0200 Subject: [PATCH 02/22] madspace: integrands with discrete choices of the phase-space mapping The discrete two-solution choices of a color-ordered mapping were never passed to it by the integrand. They are now sampled uniformly. --- .../include/madspace/phasespace/integrand.hpp | 4 +- madspace/src/phasespace/integrand.cpp | 18 ++++++- madspace/tests/test_integrand_me_outputs.py | 48 +++++++++++++++++++ 3 files changed, 67 insertions(+), 3 deletions(-) diff --git a/madspace/include/madspace/phasespace/integrand.hpp b/madspace/include/madspace/phasespace/integrand.hpp index 3da4e6e9c8..1783bdca30 100644 --- a/madspace/include/madspace/phasespace/integrand.hpp +++ b/madspace/include/madspace/phasespace/integrand.hpp @@ -94,7 +94,9 @@ class Integrand : public FunctionGenerator { }; /** - * @param mapping The phase-space mapping. + * @param mapping The phase-space mapping. Its + * discrete choices, if any, are + * sampled uniformly. * @param diff_xs One differential cross section * per active flavour. * @param adaptive_map Optional continuous adaptive diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index 5e9893e8ef..868d5cfe5a 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -504,8 +504,22 @@ NamedVector Integrand::build_channel_part( chan_index_in_group = fb.full({static_cast(0), batch_size_val}); } - // Apply phase space mapping - auto mapping_result = _mapping.build_forward(fb, {latent}, mapping_conditions); + // Apply phase space mapping, its discrete choices are sampled uniformly + ValueVec mapping_inputs{latent}; + if (std::size_t discrete_dim = _mapping.discrete_dim(); discrete_dim > 0) { + ValueVec choices; + Value discrete_random = + fb.random(batch_size_val, static_cast(discrete_dim)); + for (auto& random : fb.unstack(discrete_random)) { + auto [choice, choice_det] = + fb.sample_discrete(random, static_cast(2)); + choices.push_back(choice); + weights_before_cuts.push_back(choice_det); + } + mapping_inputs.push_back(fb.stack(choices)); + } + auto mapping_result = + _mapping.build_forward(fb, mapping_inputs, mapping_conditions); weights_before_cuts.push_back(mapping_result["det"]); Value momenta = mapping_result["momenta"]; Value x0 = mapping_result["x1"]; diff --git a/madspace/tests/test_integrand_me_outputs.py b/madspace/tests/test_integrand_me_outputs.py index 81b745bdf4..82e400a5af 100644 --- a/madspace/tests/test_integrand_me_outputs.py +++ b/madspace/tests/test_integrand_me_outputs.py @@ -176,3 +176,51 @@ def test_diagram_channel_weights_need_amp2(tmp_path): first_chan_weight_remap=[[0, 1]], first_remapped_chan_count=2, ) + + +@pytest.mark.parametrize("color_order", [[0, 2, 3, 4, 1], [0, 1, 2, 3, 4]]) +def test_color_ordered_discrete_choices(tmp_path, color_order): + """The integrand samples the discrete two-solution choices of a color-ordered + mapping, so it reproduces the phase-space volume of RAMBO.""" + context, alphas_file = make_context(str(tmp_path), 2) + matrix_element = python_matrix_element( + context, lambda momenta: np.ones(len(momenta)) + ) + cuts = ms.Cuts( + [ + ms.CutItem(ms.Observable(PIDS, ms.Observable.obs_pt, [[i]]), min=20.0) + for i in (2, 3, 4) + ] + ) + results = [] + for mode, order in [ + (ms.PhaseSpaceMapping.color_ordered, color_order), + (ms.PhaseSpaceMapping.rambo, None), + ]: + mapping = ms.PhaseSpaceMapping( + [0.0] * len(PIDS), + CM_ENERGY, + leptonic=True, + mode=mode, + cuts=cuts, + color_order=order, + ) + if order is not None: + assert mapping.discrete_dim() > 0 + diff_xs = ms.DifferentialCrossSection( + matrix_element=matrix_element, + cm_energy=CM_ENERGY, + running_coupling=None, + energy_scale=ms.CachedScale(), + ) + integrand = ms.Integrand( + mapping, + [diff_xs], + running_coupling=ms.RunningCoupling(ms.AlphaSGrid(alphas_file)), + energy_scale=ms.EnergyScale(len(PIDS), 91.188), + channel_indices=[0], + ) + results.append(generate(str(tmp_path), context, integrand, str(mode))[0]) + status, rambo_status = results + error = np.hypot(status.error, rambo_status.error) + assert abs(status.mean - rambo_status.mean) < 5 * error From f59cecca997076e5798bb280c1abe4bbc7eb0191 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 22:26:53 +0200 Subject: [PATCH 03/22] madspace: LHE completion of subprocesses without diagrams A subprocess whose matrix element has no diagrams, like a color-ordered one, has no propagators to insert. Its events carry the diagram index -1. --- .../include/madspace/driver/lhe_output.hpp | 6 ++- madspace/src/driver/lhe_output.cpp | 11 +++-- madspace/tests/test_lhe.py | 44 +++++++++++++++++++ 3 files changed, 56 insertions(+), 5 deletions(-) diff --git a/madspace/include/madspace/driver/lhe_output.hpp b/madspace/include/madspace/driver/lhe_output.hpp index fde1385ea2..d580554edb 100644 --- a/madspace/include/madspace/driver/lhe_output.hpp +++ b/madspace/include/madspace/driver/lhe_output.hpp @@ -162,7 +162,8 @@ class LHECompleter { struct SubprocArgs { /// Process id shared by every event of this subprocess. int process_id; - /// Decay topology of each diagram channel. + /// Decay topology of each diagram channel. A subprocess without diagrams + /// has one topology and no permutations and diagram indices. std::vector topologies; /// Final-state particle permutation of each channel. nested_vector3 permutations; @@ -197,7 +198,8 @@ class LHECompleter { * @param event The event to complete, modified in place. * @param subprocess_index Index into the `subproc_args` passed to the * constructor. - * @param diagram_index Sampled diagram. + * @param diagram_index Sampled diagram, -1 for a subprocess without + * diagrams. * @param color_index Sampled color configuration. * @param flavor_index Sampled flavor assignment. * @param helicity_index Sampled helicity configuration. diff --git a/madspace/src/driver/lhe_output.cpp b/madspace/src/driver/lhe_output.cpp index 1e586ac138..245715e552 100644 --- a/madspace/src/driver/lhe_output.cpp +++ b/madspace/src/driver/lhe_output.cpp @@ -529,7 +529,10 @@ void LHECompleter::complete_event_data( if (event.particles.size() != subproc_data.particle_count) { throw std::runtime_error("Invalid particle number for subprocess"); } - if (diagram_index < 0 || diagram_index >= subproc_data.diagram_count) { + // a subprocess without diagrams has no propagators to insert + bool has_diagrams = subproc_data.diagram_count > 0; + if (has_diagrams ? diagram_index < 0 || diagram_index >= subproc_data.diagram_count + : diagram_index != -1) { throw std::runtime_error("Diagram index out of range"); } if (color_index < 0 || color_index >= subproc_data.color_count) { @@ -583,6 +586,9 @@ void LHECompleter::complete_event_data( ++particle_index; } + if (!has_diagrams) { + return; + } auto find_propagators = _propagator_index_and_count.find( cantor_pairing(subprocess_index, diagram_index, color_index) ); @@ -619,8 +625,7 @@ void LHECompleter::complete_event_data( // negative m_min would instead turn it into a large positive floor and // reject nearly everything, so a wider window would write *fewer* // resonances than a narrow one. - double m_min = - std::max(0., propagator.mass - _bw_cutoff * propagator.width); + double m_min = std::max(0., propagator.mass - _bw_cutoff * propagator.width); double m_max = propagator.mass + _bw_cutoff * propagator.width; if (m2 > m_min * m_min && m2 < m_max * m_max) { auto [color, anti_color] = prop_color; diff --git a/madspace/tests/test_lhe.py b/madspace/tests/test_lhe.py index c1c14fb07f..9df98bfbe8 100644 --- a/madspace/tests/test_lhe.py +++ b/madspace/tests/test_lhe.py @@ -461,3 +461,47 @@ def test_particle_momentum_kwargs_match_attributes(): for name in ("p_x", "p_y", "p_z"): with pytest.raises(TypeError): ms.LHEParticle(**{name: 1.0}) + + +def chain_topology(masses): + """t-channel chain topology without propagator information.""" + n_out = len(masses) - 2 + vertices = [["i0", "p0", "o0"]] + vertices += [[f"p{i - 1}", f"p{i}", f"o{i}"] for i in range(1, n_out - 1)] + vertices.append(["i1", f"p{n_out - 2}", f"o{n_out - 1}"]) + diagram = ms.Diagram( + masses[:2], masses[2:], [ms.Propagator() for _ in range(n_out - 1)], vertices + ) + return ms.Topology(diagram) + + +def test_subprocess_without_diagrams(): + """Without diagrams, events carry the diagram index -1 and get no + intermediate particles.""" + pdg_ids = [21, 21, 21, 21, 21] + color_flows = [ + [[502, 501], [503, 502], [503, 504], [505, 501], [505, 504]], + [[501, 502], [502, 503], [504, 503], [501, 505], [504, 505]], + ] + subproc_args = ms.SubprocArgs( + process_id=3, + topologies=[chain_topology([0.0] * 5)], + permutations=[[]], + diagram_indices=[[]], + diagram_color_indices=[[]], + color_flows=color_flows, + pdg_color_types={21: 8}, + helicities=[[1, 1, -1, -1, 1], [-1, -1, 1, 1, -1]], + pdg_ids=[[pdg_ids]], + ) + lhe_completer = ms.LHECompleter([subproc_args], bw_cutoff=BW_CUTOFF) + event = build_event(np.ones((5, 4)) * [100.0, 0.0, 0.0, 0.0]) + lhe_completer.complete_event_data(event, 0, -1, 1, 0, 1, ms.MixMaxRandom(1)) + assert event.process_id == 3 + assert [p.pdg_id for p in event.particles] == pdg_ids + assert [(p.color, p.anti_color) for p in event.particles] == [ + tuple(flow) for flow in color_flows[1] + ] + assert [p.spin for p in event.particles] == [-1, -1, 1, 1, -1] + with pytest.raises(RuntimeError, match="Diagram index out of range"): + lhe_completer.complete_event_data(event, 0, 0, 0, 0, 0, ms.MixMaxRandom(1)) From cf7f1b7e9ddef6d96e4353fa6e3b6b7430f4f55a Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:08:22 +0200 Subject: [PATCH 04/22] mg7: split the launcher into modules launch.py stays the entry point and re-exports what the launcher modules define. The code is moved without changes: common (madspace import, logging, helpers), observables, subprocess_madmatrix, madgraph_process, options, model_widths, postprocess, driver and cli. model_widths does not import madspace, so the run interface loads the model without it. --- madgraph/iolibs/template_files/mg7/cli.py | 470 ++ madgraph/iolibs/template_files/mg7/common.py | 200 + madgraph/iolibs/template_files/mg7/driver.py | 186 + madgraph/iolibs/template_files/mg7/launch.py | 4754 +---------------- .../template_files/mg7/madgraph_process.py | 1966 +++++++ .../iolibs/template_files/mg7/model_widths.py | 155 + .../iolibs/template_files/mg7/observables.py | 178 + madgraph/iolibs/template_files/mg7/options.py | 86 + .../iolibs/template_files/mg7/postprocess.py | 380 ++ .../template_files/mg7/run_interface.py | 2 +- .../mg7/subprocess_madmatrix.py | 1131 ++++ tests/unit_tests/interface/test_mg7_launch.py | 7 +- 12 files changed, 4828 insertions(+), 4687 deletions(-) create mode 100644 madgraph/iolibs/template_files/mg7/cli.py create mode 100644 madgraph/iolibs/template_files/mg7/common.py create mode 100644 madgraph/iolibs/template_files/mg7/driver.py create mode 100644 madgraph/iolibs/template_files/mg7/madgraph_process.py create mode 100644 madgraph/iolibs/template_files/mg7/model_widths.py create mode 100644 madgraph/iolibs/template_files/mg7/observables.py create mode 100644 madgraph/iolibs/template_files/mg7/options.py create mode 100644 madgraph/iolibs/template_files/mg7/postprocess.py create mode 100644 madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py diff --git a/madgraph/iolibs/template_files/mg7/cli.py b/madgraph/iolibs/template_files/mg7/cli.py new file mode 100644 index 0000000000..d096cca0d3 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/cli.py @@ -0,0 +1,470 @@ +"""The MG7> command line interface and its card editing.""" + +import argparse +import os +import re +import resource + +from madgraph.iolibs.template_files.mg7.common import logger +from madgraph.interface.extended_cmd import Cmd +from madgraph.various.banner import RunCardMG7 +from madgraph.iolibs.template_files.mg7 import npy_to_lhe +from madgraph.iolibs.template_files.mg7.driver import force_lhe_output_if_needed, run_generation +from madgraph.iolibs.template_files.mg7.model_widths import compute_auto_widths, load_process_model +from madgraph.iolibs.template_files.mg7.options import load_mg5_options +from madgraph.iolibs.template_files.mg7.postprocess import _off, _setup_logging + +class MG7Cmd(Cmd): + """The mg7 run interface: the command object that drives a launch. + + It plays two roles, and they are the same object on purpose. + + 1. It is the *mother* of the merged switch/card question (the thing + ``AskRunEditCard`` calls back into for ``keep_cards``/``do_open``/ + ``compute_widths``). + 2. It is the child command interface MG5 registers via + ``define_child_cmd_interface`` when you type ``launch`` on an mg7 output. + + Role 2 is what makes role 1 behave. ``Cmd.ask`` resolves a scripted answer + through ``self.check_answer_in_input_file``, where ``self`` is the mother -- + so once MG5 hands us its ``inputfile`` and registers itself as our + ``mother``, the question reads the very lines that follow ``launch`` in the + user's command file, and any line it does not recognise goes back up to MG5 + through ``store_line`` instead of being silently consumed. This is exactly + how a madevent output already behaves; before this class existed the mg7 + branch shelled out to ``bin/generate_events`` and had to guess which + following lines belonged to the run, which quietly swallowed every ``set``. + """ + + prompt = 'MG7> ' + + def __init__(self, me_dir='.', options=None, *args, **opts): + super().__init__(*args, **opts) + self.me_dir = os.path.abspath(me_dir) if me_dir != '.' else '.' + # The rest of the launcher is written against the process directory as + # the working directory, so do_generate_events chdirs there; but at + # construction time MG5 has not moved, hence the explicit me_dir. + self.options = load_mg5_options( + self.me_dir if self.me_dir != '.' else None) + if options: + # MG5's own resolved tool paths win: they are the ones the user + # configured in the session that is doing the launching. + self.options.update({k: v for k, v in options.items() + if k in self.options and v is not None}) + self.plugin_path = [] + self.proc_characteristics = {'grouped_matrix': False, 'limitations': []} + self._model = None # lazily imported by get_model() + + # ------------------------------------------------------------------ + # the contract the card question expects from its mother + # ------------------------------------------------------------------ + def keep_cards(self, need_card=[], ignore=[]): + from madgraph.interface.common_run_interface import CommonRunCmd + return CommonRunCmd.keep_cards(self, need_card, ignore) + + def do_open(self, line): + from madgraph.interface.common_run_interface import CommonRunCmd + CommonRunCmd.do_open(self, line) + + def check_open(self, args): + from madgraph.interface.common_run_interface import CommonRunCmd + CommonRunCmd.check_open(self, args) + + def get_model(self): + """The model this process was generated with (an ``import_ufo`` model). + + This is the hook the card question goes through for everything that + needs the model rather than the card: 'update dependent' (which resets + the masses/widths the model computes from the free parameters) and the + generic auto-width handler both call ``mother_interface.get_model()``. + Without it the launch question could only warn that it had failed to + update the dependent parameters, and MadSpin/the shower then read + whatever inconsistent values the param_card still held. + + None when the process recorded no model (an output made before + SubProcesses/model.txt existed) -- the caller warns in that case. + """ + if self._model is None: + me_dir = self.me_dir if self.me_dir != '.' else None + # False (not None) so a model that cannot be imported is not + # retried at every question. + self._model = load_process_model(me_dir) or False + return self._model or None + + def do_compute_widths(self, line): + # The interactive card editor delegates 'auto' width computation to + # the mother interface. Reuse the runtime helper (madgraph subprocess + # + the model stored at output time). ``line`` looks like + # " --path= [--nlo]"; we only need the card path. + m = re.search(r'--path=(\S+)', line or "") + path = m.group(1) if m else os.path.join("Cards", "param_card.dat") + compute_auto_widths(path) + # return an empty mapping: the caller iterates out.items() for the + # small-width treatment, which mg7 does not apply. + return {} + + # ------------------------------------------------------------------ + # running + # ------------------------------------------------------------------ + def do_generate_events(self, line): + """generate_events [-f] [-n NAME] [--laststep=parton] + + Run the mg7 generation and the selected post-processing tools.""" + + opts = self._parse_run_options(line) + cwd = os.getcwd() + target = self.me_dir if self.me_dir != '.' else cwd + try: + os.chdir(target) + _setup_logging() + if opts['name']: + self._set_run_name(opts['name']) + switch = {} + if not opts['force']: + switch = ask_edit_cards(mother=self) + switch = self._apply_laststep(switch, opts) + force_lhe_output_if_needed(switch) + _raise_open_file_limit() + run_generation(switch) + finally: + os.chdir(cwd) + + # ``launch`` is the name MG5 users type; keep it working here too. + do_launch = do_generate_events + + def do_npy_to_lhe(self, line): + """npy_to_lhe [RUN_NAME] [-o OUTPUT] [--seed N] [--no-gzip] + + Write the LHE file (events.lhe.gz) of a run that produced npy events, + e.g. to shower it. Without a run name, the most recent such run.""" + import shlex + me_dir = self.me_dir if self.me_dir != '.' else os.getcwd() + try: + npy_to_lhe.main(shlex.split(line), me_dir=me_dir) + except SystemExit as error: + if error.code not in (None, 0): + logger.error(str(error.code)) + + def do_quit(self, line): + """Leave the mg7 run interface.""" + return super().do_quit(line) + + do_exit = do_quit + + # ------------------------------------------------------------------ + # helpers + # ------------------------------------------------------------------ + def _parse_run_options(self, line): + """Understand the subset of MG5's `launch` options that mean something + for an mg7 run, and say so out loud for the ones that do not.""" + out = {'force': False, 'name': '', 'laststep': ''} + args = self.split_arg(line) + index = 0 + while index < len(args): + arg = args[index] + index += 1 + if arg in ('-f', '--force'): + out['force'] = True + elif arg in ('-n', '--name') and index < len(args): + out['name'] = args[index] + index += 1 + elif arg.startswith(('-n=', '--name=')): + out['name'] = arg.split('=', 1)[1] + elif arg in ('-s', '--laststep') and index < len(args): + out['laststep'] = args[index] + index += 1 + elif arg.startswith(('-s=', '--laststep=')): + out['laststep'] = arg.split('=', 1)[1] + elif arg.startswith('-'): + logger.warning( + "'%s' is not supported by the mg7 output and is ignored; " + "the equivalent settings live in Cards/run_card.toml", arg) + elif arg and not out['name']: + out['name'] = arg + return out + + def _set_run_name(self, name): + """Honour `launch -n NAME`: the mg7 run directory is + Events/_NN, so the name is a run_card setting.""" + path = os.path.join("Cards", "run_card.toml") + run_card = RunCardMG7(path, consistency=False) + if run_card["run"]["run_name"] != name: + run_card["run"]["run_name"] = name + run_card.write(path) + logger.info("run name set to '%s'", name) + + def _apply_laststep(self, switch, opts): + """`--laststep=parton` means: generate the events, run nothing on them. + + The mg7 tool selection is the switch dict the question returns, so + stopping at parton level is switching those entries off. + """ + if opts['laststep'] not in ('parton', 'auto', ''): + logger.warning( + "--laststep=%s is not supported by the mg7 output; " + "select the programs to run in the launch question instead", + opts['laststep']) + return switch + if opts['laststep'] != 'parton' or not switch: + return switch + switch = dict(switch) + for key in ('shower', 'detector', 'analysis', 'madspin', 'reweight'): + if not _off(switch.get(key, 'OFF')): + logger.info("--laststep=parton: not running %s", key) + switch[key] = 'OFF' + return switch + + +def build_selector_cmd(mother=None): + """Build the (monkey-patched) mother command + merged switch/card selector + used by the mg7 output. Returns the selector *class* and a mother instance + understood by AskRun/AskforEditCard. + + ``mother`` lets an in-process caller supply its own :class:`MG7Cmd` -- the + one MG5 registered as its child -- so that the question resolves scripted + answers against MG5's own command file (see :class:`MG7Cmd`). When it is + omitted a standalone instance is built, which is what ``bin/generate_events`` + gets. + """ + + from madgraph.interface.madevent_interface import AskRunEditCard + + from madgraph.various import banner as _banner_mod + from madgraph.various import misc as _misc + + class MG7Selector(AskRunEditCard): + """Merged switch/card question for the mg7 output. + + The mg7 run_card is a TOML file, so define_paths/init_run/do_set are + overridden *on this class* (not globally on AskforEditCard) to treat it + as such. Scoping them here is important: MadSpin/reweight spawn their own + legacy madevent runs in the same process, and a global patch would break + their (run_card.dat based) card editing.""" + + # param_card + the TOML run_card are always offered + always_cards = ['param_card.dat', 'run_card.toml'] + optional_cards = [] + + def define_paths(self, **opt): + super().define_paths(**opt) + self.paths["run"] = os.path.join(self.me_dir, "Cards", "run_card.toml") + self.paths["run_card.toml"] = os.path.join(self.me_dir, "Cards", "run_card.toml") + # the TOML run_card uses its own default file (concrete defaults + # written at output time); this powers "set default". + self.paths["run_default"] = os.path.join(self.me_dir, "Cards", "run_card_default.toml") + + def init_run(self, cards): + # Make sure the run_card is loaded as a RunCardMG7 (the generic + # editor does not recognise run_card.toml), else "set " fails. + out = super().init_run(cards) + if not isinstance(getattr(self, "run_card", None), RunCardMG7): + toml_path = self.paths.get("run") or os.path.join( + self.me_dir, "Cards", "run_card.toml") + if os.path.exists(toml_path): + try: + # allow_scan so a run_card holding scan:[...] values loads + with _misc.TMP_variable(_banner_mod.RunCard, "allow_scan", True): + self.run_card = RunCardMG7(toml_path, consistency="warning") + self.run_set = list(self.run_card.keys()) + except Exception as err: + logger.warning("could not load %s: %s", toml_path, err) + if isinstance(getattr(self, "run_card", None), RunCardMG7): + self.run_card.allow_scan = True + return getattr(self, "run_set", out) + + def do_set(self, line, *args, **kwargs): + # madevent-style shortcuts (lhc/lep/fixed_scale/no_parton_cut), cut + # editing, energy units and arithmetic/mass expressions on the TOML + # run_card, intercepted before delegating to the generic editor. + targs = self.split_arg(line) + run_card = getattr(self, "run_card", None) + if isinstance(run_card, RunCardMG7) and targs: + start = 1 if targs[0] == "run_card" else 0 + if len(targs) > start: + name = targs[start] + nlow = name.lower() + rest = " ".join(targs[start + 1:]).split("#")[0].strip() + masses = run_card.get_mass_shortcuts(getattr(self, "param_card", None)) + + if nlow in ("no_parton_cut", "nocut", "no_cut"): + run_card.remove_all_cut() + logger.info("removing all cuts from the run_card.toml") + self.modified_card.add("run") + return + # [histograms] is a list of observables, not a parameter, + # so "set histograms ..." can only mean the whole section. + if nlow in ("histograms", "histogram", "plots"): + value = rest.lower() + if value in ("off", "none", "no", "false", "0"): + run_card.remove_all_histograms() + logger.info( + "removing all histograms from the run_card.toml") + self.modified_card.add("run") + return + if value in ("default", "on"): + if run_card.restore_default_histograms( + self.paths.get("run_default")): + logger.info( + "restored the %d histograms written at " + "output time", + len(run_card["histograms"])) + self.modified_card.add("run") + else: + logger.warning( + "no run_card_default.toml to restore the " + "histograms from") + return + logger.warning( + "'set histograms %s': only OFF (remove them all) " + "and default (put the generated ones back) are " + "understood; edit the [histograms] section of " + "Cards/run_card.toml for anything else", + rest or "(no value)") + return + if nlow in ("lhc", "lep", "ilc", "lcc") and rest: + ecm = run_card.set_collider(nlow, rest, masses) + logger.info("set %s collider: e_cm = %s GeV", nlow, ecm) + self.modified_card.add("run") + return + if nlow == "fixed_scale" and rest: + val = run_card.set_fixed_scale(rest, masses) + logger.info("set fixed scales to %s GeV", val) + self.modified_card.add("run") + return + + if rest and run_card.is_cut_name(name): + cut, bound, val = run_card.set_cut(name, run_card.evaluate(rest, masses)) + logger.info("modify cut %s.%s of the run_card.toml to %s", cut, bound, val) + self.modified_card.add("run") + return + + # 'iseed' is the madevent spelling of the mg7 run.seed, + # but the two disagree on how to ask for a random seed: + # madevent uses 0, mg7 uses -1 (0 being a perfectly good + # fixed seed there). Translate here rather than in + # _LO_SCALAR_MAP, which also drives the LO -> MG7 run_card + # conversion and would silently pin those runs to seed 0. + if rest and nlow == 'iseed': + value = run_card.evaluate(rest, masses) + try: + value = int(value) + except (TypeError, ValueError): + value = None + if value is None: + logger.warning("ignoring 'set iseed %s': not an integer", rest) + return + seed = -1 if value == 0 else value + run_card.set('run.seed', seed, user=True) + logger.info("set iseed (mg7 run.seed) of the run_card.toml to %s%s", + seed, " (random)" if seed == -1 else "") + self.modified_card.add("run") + return + + # legacy madevent run_card names -> mg7 "section.key", so a + # madevent-style launch script ("set nevents 500", "set + # use_syst F", "set bwcutoff 10", ...) edits the mg7 + # run_card.toml verbatim. RunCardMG7._LO_SCALAR_MAP is the + # same rename table used by the LO->MG7 run_card conversion. + if rest and nlow in run_card._LO_SCALAR_MAP \ + and nlow not in [k.lower() for k in run_card.keys()]: + target = run_card._LO_SCALAR_MAP[nlow] + run_card.set(target, rest, user=True) + logger.info("set %s (mg7 %s) of the run_card.toml to %s", + nlow, target, run_card[target]) + self.modified_card.add("run") + return + + if rest and nlow in [k.lower() for k in run_card.keys()]: + current = run_card[nlow] + if isinstance(current, (int, float)) and not isinstance(current, bool): + resolved = run_card.evaluate(rest, masses) + if not isinstance(resolved, str): + prefix = "run_card " if start == 1 else "" + line = "%s%s %s" % (prefix, name, resolved) + return super().do_set(line, *args, **kwargs) + + return MG7Selector, (mother if mother is not None else MG7Cmd()) + + +def ask_edit_cards(mother=None) -> dict: + """Single (MadDM-style) question letting the user both pick which programs + to run after generation and edit the associated cards. Returns the switch + dict describing the selected tools. + + ``mother`` is the :class:`MG7Cmd` asking the question. Passing the instance + MG5 registered as its child is what lets a scripted ``launch`` answer this + question from MG5's own command file. + """ + + selector_class, mother = build_selector_cmd(mother) + # path_msg is what makes Cmd.check_answer_in_input_file accept a bare path as + # an answer (its "elif path:" branch), so that a scripted launch can hand the + # question a card/banner path -- as the question itself advertises -- and have + # it replace the corresponding card. Without it the path is rejected ("This + # answer is not valid for current question") and the default is used instead. + try: + switch, question = mother.ask('', '0', [], path_msg='enter path', + ask_class=selector_class, + mode='auto', line_args=[], force=False, + return_instance=True) + except BaseException: + # interrupted (e.g. ctrl-C) before the prune below could run: prune + # here too, using the instance ask() stashed pre-interrupt, or the + # eagerly-materialised tool cards get misread as "left ON" next time. + instance = getattr(mother, '_last_ask_instance', None) + if instance is not None and hasattr(instance, 'active_cards'): + try: + mother.keep_cards(instance.active_cards()) + except Exception as error: + logger.debug('could not prune tool cards after aborted launch question: %s', error) + raise + switch = dict(switch) + prune_unselected_tool_cards(mother, question, switch) + return switch + + +def prune_unselected_tool_cards(mother, question, switch) -> None: + """Hide the cards of the tools that were NOT selected in the question. + + The question materialises every candidate tool card (copying the + *_default.dat) so that it can offer them all for edition; madevent's + ask_run_configuration prunes them again afterwards (keep_cards on the + selected cards), but the mg7 launcher had no such cleanup. The materialised + cards therefore persisted, and a later re-launch of the same output saw them + all present and defaulted every switch to ON (set_default_ keys off + card presence). Prune here too, using the question's own switch->card map, so + a re-launch reflects the previous selection instead of turning everything on. + """ + try: + keep = list(question.always_cards) + for spec in question.switch_cards: + if spec['on'](switch): + keep.append(spec['card']) + mother.keep_cards(keep) + except Exception as error: + logger.debug('could not prune unselected tool cards: %s', error) + + +def _raise_open_file_limit() -> None: + """Remove the soft limit on the number of open files: it can be quite low + on some systems and the event-combination step opens one file per channel.""" + soft_lim, hard_lim = resource.getrlimit(resource.RLIMIT_NOFILE) + try: + resource.setrlimit(resource.RLIMIT_NOFILE, (hard_lim, hard_lim)) + except (ValueError, OSError) as error: + logger.debug("could not raise the open-file limit: %s", error) + + +def main() -> None: + """``bin/generate_events`` entry point. + + It runs the same :class:`MG7Cmd` command MG5's in-process ``launch`` uses, + so the two entry points cannot drift apart; the only difference is that + nothing here supplies a ``mother``/``inputfile``, so the question falls back + to reading this process's own stdin. + """ + parser = argparse.ArgumentParser() + parser.add_argument("-f", action="store_false", dest="ask_edit_cards") + args = parser.parse_args() + + MG7Cmd().run_cmd("generate_events" if args.ask_edit_cards else "generate_events -f") diff --git a/madgraph/iolibs/template_files/mg7/common.py b/madgraph/iolibs/template_files/mg7/common.py new file mode 100644 index 0000000000..253f4d06d9 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/common.py @@ -0,0 +1,200 @@ +"""Shared state of the mg7 launcher: the madspace import, logging and small helpers.""" + +import logging +import os +import re +import subprocess +import sys +import time + +# Install madspace on first use, then import it. The bootstrap lives in its own +# module (which must not import madspace) so that an in-process caller can run +# it *before* importing this one -- see bootstrap.ensure_madspace. +from madgraph.iolibs.template_files.mg7.bootstrap import ( + ensure_madspace as _ensure_madspace, + drop_install_path as _drop_install_path, + MADSPACE_DIR as _MADSPACE_DIR, +) + +_ensure_madspace() + +import madspace as ms + +# madspace is imported; stop its install directory from shadowing the caller's +# yaml/packaging/... for the rest of what is now MG5's own session. +_drop_install_path() + +_source_hash = subprocess.run( + [sys.executable, str(_MADSPACE_DIR / "source_hash.py")], + capture_output=True, text=True, check=True, +).stdout.strip() +if _source_hash != ms.SOURCE_HASH: + print() + print( + "\033[1m\033[31mWARNING\033[39m: madspace source and installed binaries " + "are not compatible (source hash mismatch) — consider recompiling " + "madspace (e.g. `install madspace -y`)\033[0m" + ) + print() + +logger = logging.getLogger("madgraph7") +LOG_LEVEL_MAP = { + ms.Logger.LogLevel.level_debug: logging.DEBUG, + ms.Logger.LogLevel.level_info: logging.INFO, + ms.Logger.LogLevel.level_warning: logging.WARNING, + ms.Logger.LogLevel.level_error: logging.ERROR, +} +def ms_log_handler(level: ms.Logger.LogLevel, message: str): + logger.log(LOG_LEVEL_MAP[level], message) +ms.Logger.set_log_handler(ms_log_handler) + + +def get_start_time(): + return time.time(), time.process_time() + + +def format_time(t: int, centi: bool = False): + hours, t = divmod(t, 3600) + minutes, seconds = divmod(t, 60) + if centi: + return f"{int(hours):02}:{int(minutes):02}:{seconds:02.2f}" + else: + return f"{int(hours):02}:{int(minutes):02}:{seconds:02.0f}" + + +def resolve_verbosity(verbosity: str) -> str: + """Resolve the run_card "auto" verbosity to "pretty"/"log" depending on + whether stdout is attached to a terminal; other values pass through + unchanged.""" + if verbosity == "auto": + return "pretty" if sys.stdout.isatty() else "log" + return verbosity + + +def resolve_seed(seed: int) -> int: + """Resolve the run_card "seed": -1 draws a fresh 64-bit seed via + os.urandom, any other value is used as-is.""" + if seed == -1: + return int.from_bytes(os.urandom(8), "big") + return seed + + +def device_type_of(device_name: str) -> str: + """The device type of one 'device' run_card entry (the optional ':' + suffix is stripped).""" + return device_name.split(":")[0] + + +def backend_of(device_name: str, cpu_mode: str) -> str: + """The build BACKEND that serves one 'device' run_card entry. + + ``cpu_mode`` describes the SIMD width of the CPU code and is therefore + meaningless for the cuda/hip devices: those build the backend named after + the device itself. + """ + device_type = device_type_of(device_name) + return cpu_mode if device_type == "cpu" else device_type + + +_PRECISION_2_FPTYPE = {"all64": "d", "all32": "f", "color32": "m", "denom64": "v"} + + +def fptype_of(precision: str) -> str: + return _PRECISION_2_FPTYPE.get(str(precision).lower(), str(precision).lower()) + + +def resolve_auto_backend(build_path: str) -> str: + """Ask the matrix-element Makefile to resolve the ``auto`` cpu_mode. + + Given the produced shared library will have its name taken from the resolved + backend name, we need to make sure the detection is taking place correctly + and catch any possible error. + """ + command = ["make", "-n", "BACKEND=auto", "detect-backend"] + try: + result = subprocess.run( + command, + cwd=build_path, + capture_output=True, + text=True, + check=True, + ) + except OSError as exc: + raise RuntimeError( + f"Could not run make to resolve cpu_mode='auto' in '{build_path}': {exc}" + ) from exc + except subprocess.CalledProcessError as exc: + output = "\n".join( + part.strip() for part in (exc.stdout, exc.stderr) if part.strip() + ) + detail = f"\nmake output:\n{output}" if output else "" + raise RuntimeError( + f"Could not resolve cpu_mode='auto' in '{build_path}': " + f"Exit status {exc.returncode}.{detail}" + ) from exc + + match = re.search( + r"^BACKEND=(\S+) \(was auto\)$", result.stdout, re.MULTILINE + ) + if match is None or match.group(1) == "auto": + output = "\n".join( + part.strip() for part in (result.stdout, result.stderr) if part.strip() + ) + detail = f"\nmake output:\n{output}" if output else "" + raise RuntimeError( + f"Could not resolve cpu_mode='auto' in '{build_path}': " + f"make failed to report a backend.{detail}" + ) + return match.group(1) + + +def me_parameters(run_card): + """Run-time parameters of the matrix-element libraries, passed to each + instance through umami_set_parameter by ms.Context.load_matrix_element: + the window of the $-excluded propagators is the run card bw_cutoff (the + one the phase space uses around the resonances too), and with + [generation] interference_helicity = "summed" the helicity sum as the + |M|^2 of an interference (see interference_helicities_summed). That one is + only sent when asked for: a library made before it existed refuses an + unknown parameter.""" + parameters = {"bwcutoff": float(run_card["phasespace"]["bw_cutoff"])} + if interference_helicities_summed(run_card): + parameters["interference_helicity_summed"] = 1. + return parameters + + +def interference_helicities_summed(run_card): + """Whether [generation] interference_helicity is "summed". + + For an interference (squared split orders dropping a component), a helicity + can contribute negatively. "exact", the default, draws the helicity of each + event on |T_i| and gives the event sign(T_i) * sum_j |T_j|, as madevent does: + the helicity of the event is exact, at the cost of a larger sum|w| (about 2x + the events for the same precision). "summed" keeps the helicity sum sum_j T_j + as the weight (lower variance), and the helicity then means nothing: it is + written as 9 in the LHE (lhe_helicities).""" + return run_card["generation"].get("interference_helicity", "exact") == "summed" + + +def lhe_helicities(meta, run_card): + """The helicity table the LHE completer writes for a subprocess: its own, + or 9 for every particle of an interference subprocess whose |M|^2 is the + helicity sum (interference_helicities_summed).""" + helicities = meta["helicities"] + if meta.get("interference") and interference_helicities_summed(run_card): + return [[9] * len(row) for row in helicities] + return helicities + + +def lhe_weight_info(status): + """(IDWTUP, XMAXUP) for the LHE block of a run with this generator + status. The events are unweighted to +-sigma_abs, keeping the sign of the + integrand, and XMAXUP is that unit weight rather than the signed cross + section. A sample that can hold negative weights -- an interference -- + declares it with IDWTUP = -4, as madevent does: the cross section is then + the mean of the signed event weights. With -3 a shower takes it as + |XSECUP| times the mean sign instead (Pythia8 does), which is not the + cross section, and +3 is for positive weights only. gridpack.py applies the + same rule inline, since it must not import this module.""" + negative = status.mean < 0 or status.mean_abs > abs(status.mean) * (1 + 1e-12) + return (-4 if negative else 3), status.mean_abs diff --git a/madgraph/iolibs/template_files/mg7/driver.py b/madgraph/iolibs/template_files/mg7/driver.py new file mode 100644 index 0000000000..ed1d974ce9 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/driver.py @@ -0,0 +1,186 @@ +"""Runs, scans and the entry point of the mg7 launcher.""" + +import gc +import logging +import os +import re +import shutil + +from madgraph.iolibs.template_files.mg7.common import logger +from madgraph.various.banner import RunCardMG7 +from madgraph.iolibs.template_files.mg7.model_widths import compute_auto_widths +from madgraph.iolibs.template_files.mg7.postprocess import run_lhe_postprocessing, run_selected_tools, selected_tools +from madgraph.iolibs.template_files.mg7.madgraph_process import MadgraphProcess + +def _release_channel_generators(process) -> None: + """Drop the per-channel event/weight generators so their intermediate + .npy files are deleted right away. madspace only frees them when the + owning Python objects are collected, and process.event_generator / + process.phasespaces sit in a reference cycle that plain refcounting + never breaks -- only an explicit gc.collect() does, promptly, instead of + leaving it to whenever the cyclic GC next runs on its own.""" + process.event_generator = None + process.phasespaces = None + gc.collect() + + +def run_single(switch=None) -> "MadgraphProcess": + """Run a single generation and return the process (for its result).""" + compute_auto_widths() + process = MadgraphProcess() + process.survey() + process.train_madnis() + process.generate_events() + if process.gridpack_run_mode != "regular": + return process + # run_card-driven LHE post-processing (displaced vertex + systematics) + run_lhe_postprocessing(process) + # run the post-processing tools (Pythia8/Delphes/MadSpin/reweight/analysis) + # selected in the merged question above on the generated events. + if switch: + run_selected_tools(switch, process) + return process + + +def detect_run_scan(run_card_path): + """Return a banner.RunCardIterator if the run_card contains scan:[...] + values, else None.""" + from madgraph.various import banner as banner_mod + from madgraph.various import misc as _misc + with _misc.TMP_variable(banner_mod.RunCard, 'allow_scan', True): + rc = banner_mod.RunCard(run_card_path, consistency=False) + if getattr(rc, 'scan_set', None): + return banner_mod.RunCardIterator(run_card_path) + return None + + +def detect_param_scan(param_card_path): + """Return a ParamCardIterator if the param_card contains scan:[...] values, + else None.""" + if not os.path.exists(param_card_path): + return None + from models import check_param_card as param_card_mod + it = param_card_mod.ParamCardIterator(param_card_path) + for block in it.order: + for param in block: + if isinstance(param.value, str) and param.value.strip().lower().startswith('scan'): + return it + return None + + +def run_scan(iterator, card_path, switch=None) -> None: + """Iterate over all scan points, running a full generation for each and + accumulating the results, then write the scan summary. Works for both the + run_card (RunCardIterator) and the param_card (ParamCardIterator); their + interface (__iter__/write/store_entry/get_next_name/write_summary) is the + same. The scan card is restored afterwards.""" + import tomllib + with open(os.path.join("Cards", "run_card.toml"), "rb") as f: + run_name = tomllib.load(f).get("run", {}).get("run_name", "run") + + from models import check_param_card as param_card_mod + is_param_scan = isinstance(iterator, param_card_mod.ParamCardIterator) + + backup = card_path + ".scan_bak" + shutil.copy(card_path, backup) + try: + for i, point in enumerate(iterator): + point.write(card_path) + logger.info("=== scan point %d ===", i + 1) + process = run_single(switch) + # use the run directory the process actually created, so the + # per-point params.dat written by write_summary has a home + name = os.path.basename(process.run_path) + if is_param_scan: + # pass the (possibly auto-width-updated) card so the summary + # records the model-computed width for each 'auto' entry + iterator.store_entry(name, process.get_result(), + param_card_path=card_path) + else: + iterator.store_entry(name, process.get_result()) + _release_channel_generators(process) + os.makedirs("Events", exist_ok=True) + summary = os.path.join("Events", "scan_%s.txt" % run_name) + iterator.write_summary(summary) + logger.info("scan results written to %s", summary) + finally: + shutil.move(backup, card_path) + + +def run_generation(switch=None) -> None: + """Run the generation, expanding a scan over the run_card or the param_card + when one is present (scanning both simultaneously is not allowed).""" + run_card_path = os.path.join("Cards", "run_card.toml") + param_card_path = os.path.join("Cards", "param_card.dat") + run_iter = detect_run_scan(run_card_path) + param_iter = detect_param_scan(param_card_path) + if run_iter and param_iter: + raise RuntimeError( + "Scanning simultaneously over the run_card and the param_card is " + "not allowed. Please keep the scan:[...] entries in only one card.") + if run_iter: + run_scan(run_iter, run_card_path, switch) + elif param_iter: + run_scan(param_iter, param_card_path, switch) + else: + _release_channel_generators(run_single(switch)) + + +def _lhe_needed_by(switch, card) -> "str | None": + """Why the run needs an LHE event file, or None when it does not: every + post-processing tool run_selected_tools drives and the run_card + [postprocessing] steps read the LHE file. ``card`` is the raw TOML content + of the run_card.""" + tools = selected_tools(switch) if switch else [] + if tools: + return tools[0] + cfg = card.get("postprocessing", {}) + try: + if float(cfg.get("time_of_flight", -1.0)) >= 0: + return "time_of_flight" + except (TypeError, ValueError): + pass + # the legacy systematics.py path only runs when madspace did not compute + # the weights itself (see run_lhe_postprocessing) + if cfg.get("systematics") \ + and not card.get("systematics", {}).get("enable", True): + return "systematics" + return None + + +def force_lhe_output_if_needed(switch) -> None: + """The events are written as npy by default; switch the run_card to the + LHE format when a selected post-processing needs an LHE file. The card is + rewritten rather than refused, so that asking for a shower in the launch + question is enough -- the user does not have to know that it constrains + output_format. + + The card is read as raw TOML and only its output_format line is edited: a + run_card scan ("scan:[...]" values) is not a valid RunCardMG7 and must + survive untouched for run_scan.""" + import tomllib + path = os.path.join("Cards", "run_card.toml") + if not os.path.exists(path): + return + with open(path, "rb") as f: + card = tomllib.load(f) + output_format = card.get("run", {}).get( + "output_format", RunCardMG7()["run"]["output_format"]) + if output_format == "lhe": + return + reason = _lhe_needed_by(switch, card) + if reason is None: + return + with open(path) as f: + text = f.read() + text, count = re.subn(r'(?m)^(\s*output_format\s*=\s*)("[^"]*"|\'[^\']*\'|\S+)', + r'\g<1>"lhe"', text, count=1) + if not count: + text, count = re.subn(r'(?m)^\[run\][^\n]*\n', '\\g<0>output_format = "lhe"\n', + text, count=1) + if not count: + text += '\n[run]\noutput_format = "lhe"\n' + with open(path, "w") as f: + f.write(text) + logging.getLogger("madevent").info( + "output_format set to 'lhe' (required by %s).", reason) diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index b871e59039..ceb4420a49 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -1,4685 +1,73 @@ -import argparse -import gc -import os -import sys -import time -from datetime import timedelta -from pathlib import Path -import glob -import shutil -import json -import subprocess -import re -import logging -from dataclasses import dataclass, field -from typing import Literal, NamedTuple -import resource - -# Install madspace on first use, then import it. The bootstrap lives in its own -# module (which must not import madspace) so that an in-process caller can run -# it *before* importing this one -- see bootstrap.ensure_madspace. -from madgraph.iolibs.template_files.mg7.bootstrap import ( - ensure_madspace as _ensure_madspace, - drop_install_path as _drop_install_path, - MG_ROOT as _MG_ROOT, - MADSPACE_DIR as _MADSPACE_DIR, - INSTALL_DIR as _INSTALL_DIR, +"""Entry point of the mg7 launcher, re-exporting its modules.""" + +from madgraph.iolibs.template_files.mg7.common import * # noqa: F401,F403 +from madgraph.iolibs.template_files.mg7.common import ms, logger # noqa: F401 +from madgraph.iolibs.template_files.mg7.observables import ( # noqa: F401 + CutItem, + HistItem, + _MERGED_PID_REPRESENTATIVE, + clean_pids, + _SIGNED_OBSERVABLES, + _flavor_options, + observable_pids, + build_observable, + decay_products, + pid_is_qcd, ) - -_ensure_madspace() - -import madspace as ms - -# madspace is imported; stop its install directory from shadowing the caller's -# yaml/packaging/... for the rest of what is now MG5's own session. -_drop_install_path() -from models.check_param_card import ParamCard -from madgraph.iolibs.template_files.mg7 import hwu_output -from madgraph.iolibs.template_files.mg7 import plots -from madgraph.iolibs.template_files.mg7 import systematics_summary -from madgraph.iolibs.template_files.mg7 import npy_to_lhe -from madgraph.iolibs.template_files.mg7.custom_function_loader import ( - load_custom_function, - load_surrogate, +from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import ( # noqa: F401 + Channel, + PhaseSpace, + MultiChannelData, + build_topologies, + build_multi_channel_data, + MadgraphSubprocess, +) +from madgraph.iolibs.template_files.mg7.madgraph_process import ( # noqa: F401 + MadgraphProcess, +) +from madgraph.iolibs.template_files.mg7.options import ( # noqa: F401 + _ROOTED_OPTIONS, + load_mg5_options, + lhapdf_paths, +) +from madgraph.iolibs.template_files.mg7.model_widths import ( # noqa: F401 + _model_hash_warned, + _model_cache, + read_stored_model, + _proc_characteristic, + load_process_model, + compute_auto_widths, +) +from madgraph.iolibs.template_files.mg7.postprocess import ( # noqa: F401 + _find_event_file, + _find_or_convert_event_file, + _report_failure, + _off, + _has_effective_handler, + _TOOL_LOGGING_READY, + _setup_logging, + selected_tools, + run_selected_tools, + _add_time_of_flight, + _lhapdf_config_path, + _run_systematics, + run_lhe_postprocessing, +) +from madgraph.iolibs.template_files.mg7.driver import ( # noqa: F401 + _release_channel_generators, + run_single, + detect_run_scan, + detect_param_scan, + run_scan, + run_generation, + _lhe_needed_by, + force_lhe_output_if_needed, +) +from madgraph.iolibs.template_files.mg7.cli import ( # noqa: F401 + MG7Cmd, + build_selector_cmd, + ask_edit_cards, + prune_unselected_tool_cards, + _raise_open_file_limit, + main, ) -from madgraph.various.banner import RunCardMG7 -from madgraph.various import misc -from madgraph.interface.extended_cmd import Cmd - -_source_hash = subprocess.run( - [sys.executable, str(_MADSPACE_DIR / "source_hash.py")], - capture_output=True, text=True, check=True, -).stdout.strip() -if _source_hash != ms.SOURCE_HASH: - print() - print( - "\033[1m\033[31mWARNING\033[39m: madspace source and installed binaries " - "are not compatible (source hash mismatch) — consider recompiling " - "madspace (e.g. `install madspace -y`)\033[0m" - ) - print() - -logger = logging.getLogger("madgraph7") -LOG_LEVEL_MAP = { - ms.Logger.LogLevel.level_debug: logging.DEBUG, - ms.Logger.LogLevel.level_info: logging.INFO, - ms.Logger.LogLevel.level_warning: logging.WARNING, - ms.Logger.LogLevel.level_error: logging.ERROR, -} -def ms_log_handler(level: ms.Logger.LogLevel, message: str): - logger.log(LOG_LEVEL_MAP[level], message) -ms.Logger.set_log_handler(ms_log_handler) - - -def get_start_time(): - return time.time(), time.process_time() - - -def format_time(t: int, centi: bool = False): - hours, t = divmod(t, 3600) - minutes, seconds = divmod(t, 60) - if centi: - return f"{int(hours):02}:{int(minutes):02}:{seconds:02.2f}" - else: - return f"{int(hours):02}:{int(minutes):02}:{seconds:02.0f}" - - -def resolve_verbosity(verbosity: str) -> str: - """Resolve the run_card "auto" verbosity to "pretty"/"log" depending on - whether stdout is attached to a terminal; other values pass through - unchanged.""" - if verbosity == "auto": - return "pretty" if sys.stdout.isatty() else "log" - return verbosity - - -def resolve_seed(seed: int) -> int: - """Resolve the run_card "seed": -1 draws a fresh 64-bit seed via - os.urandom, any other value is used as-is.""" - if seed == -1: - return int.from_bytes(os.urandom(8), "big") - return seed - - -def device_type_of(device_name: str) -> str: - """The device type of one 'device' run_card entry (the optional ':' - suffix is stripped).""" - return device_name.split(":")[0] - - -def backend_of(device_name: str, cpu_mode: str) -> str: - """The build BACKEND that serves one 'device' run_card entry. - - ``cpu_mode`` describes the SIMD width of the CPU code and is therefore - meaningless for the cuda/hip devices: those build the backend named after - the device itself. - """ - device_type = device_type_of(device_name) - return cpu_mode if device_type == "cpu" else device_type - - -_PRECISION_2_FPTYPE = {"all64": "d", "all32": "f", "color32": "m", "denom64": "v"} - - -def fptype_of(precision: str) -> str: - return _PRECISION_2_FPTYPE.get(str(precision).lower(), str(precision).lower()) - - -def resolve_auto_backend(build_path: str) -> str: - """Ask the matrix-element Makefile to resolve the ``auto`` cpu_mode. - - Given the produced shared library will have its name taken from the resolved - backend name, we need to make sure the detection is taking place correctly - and catch any possible error. - """ - command = ["make", "-n", "BACKEND=auto", "detect-backend"] - try: - result = subprocess.run( - command, - cwd=build_path, - capture_output=True, - text=True, - check=True, - ) - except OSError as exc: - raise RuntimeError( - f"Could not run make to resolve cpu_mode='auto' in '{build_path}': {exc}" - ) from exc - except subprocess.CalledProcessError as exc: - output = "\n".join( - part.strip() for part in (exc.stdout, exc.stderr) if part.strip() - ) - detail = f"\nmake output:\n{output}" if output else "" - raise RuntimeError( - f"Could not resolve cpu_mode='auto' in '{build_path}': " - f"Exit status {exc.returncode}.{detail}" - ) from exc - - match = re.search( - r"^BACKEND=(\S+) \(was auto\)$", result.stdout, re.MULTILINE - ) - if match is None or match.group(1) == "auto": - output = "\n".join( - part.strip() for part in (result.stdout, result.stderr) if part.strip() - ) - detail = f"\nmake output:\n{output}" if output else "" - raise RuntimeError( - f"Could not resolve cpu_mode='auto' in '{build_path}': " - f"make failed to report a backend.{detail}" - ) - return match.group(1) - - -def me_parameters(run_card): - """Run-time parameters of the matrix-element libraries, passed to each - instance through umami_set_parameter by ms.Context.load_matrix_element: - the window of the $-excluded propagators is the run card bw_cutoff (the - one the phase space uses around the resonances too), and with - [generation] interference_helicity = "summed" the helicity sum as the - |M|^2 of an interference (see interference_helicities_summed). That one is - only sent when asked for: a library made before it existed refuses an - unknown parameter.""" - parameters = {"bwcutoff": float(run_card["phasespace"]["bw_cutoff"])} - if interference_helicities_summed(run_card): - parameters["interference_helicity_summed"] = 1. - return parameters - - -def interference_helicities_summed(run_card): - """Whether [generation] interference_helicity is "summed". - - For an interference (squared split orders dropping a component), a helicity - can contribute negatively. "exact", the default, draws the helicity of each - event on |T_i| and gives the event sign(T_i) * sum_j |T_j|, as madevent does: - the helicity of the event is exact, at the cost of a larger sum|w| (about 2x - the events for the same precision). "summed" keeps the helicity sum sum_j T_j - as the weight (lower variance), and the helicity then means nothing: it is - written as 9 in the LHE (lhe_helicities).""" - return run_card["generation"].get("interference_helicity", "exact") == "summed" - - -def lhe_helicities(meta, run_card): - """The helicity table the LHE completer writes for a subprocess: its own, - or 9 for every particle of an interference subprocess whose |M|^2 is the - helicity sum (interference_helicities_summed).""" - helicities = meta["helicities"] - if meta.get("interference") and interference_helicities_summed(run_card): - return [[9] * len(row) for row in helicities] - return helicities - - -def lhe_weight_info(status): - """(IDWTUP, XMAXUP) for the LHE block of a run with this generator - status. The events are unweighted to +-sigma_abs, keeping the sign of the - integrand, and XMAXUP is that unit weight rather than the signed cross - section. A sample that can hold negative weights -- an interference -- - declares it with IDWTUP = -4, as madevent does: the cross section is then - the mean of the signed event weights. With -3 a shower takes it as - |XSECUP| times the mean sign instead (Pythia8 does), which is not the - cross section, and +3 is for positive weights only. gridpack.py applies the - same rule inline, since it must not import this module.""" - negative = status.mean < 0 or status.mean_abs > abs(status.mean) * (1 + 1e-12) - return (-4 if negative else 3), status.mean_abs - - -@dataclass -class Channel: - phasespace_mapping: ms.PhaseSpaceMapping - adaptive_mapping: ms.Flow | ms.VegasMapping - discrete_sym: ms.DiscreteSampler | ms.DiscreteFlow | None - discrete_flavor: ms.DiscreteSampler | ms.DiscreteFlow | None - channel_weight_indices: list[int] | None - name: str - active_flavors: list[int] - event_generator: ms.ChannelEventGenerator | None = None - - -@dataclass -class PhaseSpace: - mode: Literal["multichannel", "flat", "both"] - channels: list[Channel] - symfact: list[int | None] - first_chan_weight_remap: list[list[int]] = field(default_factory=list) - first_remapped_chan_count: int = 0 - second_chan_weight_remap: list[int] = field(default_factory=list) - second_remapped_chan_count: int = 0 - prop_chan_weights: ms.PropagatorChannelWeights | None = None - subchan_weights: ms.SubchannelWeights | None = None - cwnet: ms.ChannelWeightNetwork | None = None - - -class MultiChannelData(NamedTuple): - amp2_remaps: list[list[int]] - symfact: list[int | None] - topologies: list[list[ms.Topology]] - permutations: list[list[list[int]]] - channel_indices: list[list[int]] - channel_weight_indices: list[list[list[int]]] - diagram_indices: list[list[int]] - diagram_color_indices: list[list[list[int]]] - diagram_propagator_pdgs: list[list[list[int]]] - active_flavors: list[list[list[int]]] - qcd_s_channel_count: list[int] - - -@dataclass -class CutItem: - observable_kwargs: dict - min: float - max: float - mode: str - - -@dataclass -class HistItem: - observable_kwargs: dict - min: float - max: float - bin_count: int - # the distribution of the event weight itself (the reserved "weight" key - # of [histograms]); it is read off the event record, not computed from - # the momenta, so it has no observable - name: str = "" - from_weight: bool = False - - -class MadgraphProcess: - def __init__(self): - self.madnis_trained = False - self.lhe_completer = None - self.load_cards() - self.init_backend() - self.init_event_dir() - self.init_context() - self.init_custom_functions() - self.init_surrogate() - self.init_cuts() - self.init_histograms() - self.init_generator_config() - self.init_beam() - self.init_subprocesses() - - def load_cards(self) -> None: - self.run_card = RunCardMG7(os.path.join("Cards", "run_card.toml")) - # Resolved once so every generator built during this run shares the same - # seed; the concrete value (even if randomly drawn) is recorded in each - # generator's info.json. - self.run_seed = resolve_seed(self.run_card["run"]["seed"]) - self.param_card_path = os.path.join("Cards", "param_card.dat") - self.param_card = ParamCard(self.param_card_path) - with open(os.path.join("SubProcesses", "subprocesses.json")) as f: - self.subprocess_data = json.load(f) - if self.run_card["phasespace"]["merge_subprocesses"]: - with open(os.path.join("SubProcesses", "merged_subprocesses.json")) as f: - self.merged_subprocess_data = json.load(f) - else: - self.merged_subprocess_data = None - - self.init_decay_mode() - self.init_me_frame() - if interference_helicities_summed(self.run_card): - count = sum(1 for meta in self.subprocess_data if meta.get("interference")) - if count: - logger.info( - "interference_helicity = summed: the weight of the %d interference " - "subprocess(es) is the helicity sum, and their helicities are " - "written as 9 in the LHE", count) - - def init_me_frame(self) -> None: - """Resolve the run card's me_frame into the list of external particles - whose momentum sum defines the frame the matrix element is evaluated in. - - Empty (the default) means no boost: the matrix element sees the momenta - in the frame they are generated in. Unlike madevent, which hands its - matrix element momenta that are already in the partonic centre of mass, - madspace generates them in the lab frame, so asking for madevent's own - frame ([1, 2]) does mean an actual boost here. It changes nothing for a - Lorentz invariant matrix element (only the rounding), but it is what - defines the polarisation axes of a polarised one. - - The particle numbers are checked here rather than only in madspace, - which would not see them until the integrands are built -- long after - the matrix-element libraries have been compiled. Every subprocess is - checked against its own particles: an output can mix multiplicities - (p p > z j and p p > z j j, say, for MLM merging), and madspace builds - one matrix element per subprocess, each refusing a particle it does not - have. That is stricter than madevent, which ignores the particles a - process does not have and so changes the frame from one multiplicity - to the next without a word. - """ - me_frame = list(self.run_card["run"]["me_frame"]) - pdgs = [clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) - for meta in self.subprocess_data] - counts = sorted({len(pdg_list) for pdg_list in pdgs}) - for meta, pdg_list in zip(self.subprocess_data, pdgs): - for index in me_frame: - if 1 <= index <= len(pdg_list): - continue - process = "%s > %s" % ( - " ".join(str(pid) for pid in meta["incoming"]), - " ".join(str(pid) for pid in meta["outgoing"]), - ) - mixed = ( - f"; this output mixes subprocesses with " - f"{', '.join(str(count) for count in counts)} external " - "particles, and me_frame has to name particles every one " - "of them has" - if len(counts) > 1 else "" - ) - raise ValueError( - f"me_frame particle {index} out of range: subprocess " - f"{process} has {len(pdg_list)} external particles{mixed}" - ) - if len(me_frame) == 1: - # A massless particle has no rest frame: its momentum squares to - # zero, so the boost divides by it and every momentum comes back as - # a NaN, silently. Name the mistake instead. - massless = [pdg_list[me_frame[0] - 1] for pdg_list in pdgs - if self.get_mass(pdg_list[me_frame[0] - 1]) == 0.] - if massless: - raise ValueError( - f"me_frame = {me_frame} asks for the rest frame of particle " - f"{me_frame[0]}, which is massless (pdg " - f"{sorted(set(massless))}) and has none. Name a massive " - "particle, or several particles whose sum is massive." - ) - self.me_frame = me_frame - self.incoming_count = 1 if self.is_decay else 2 - - def init_decay_mode(self) -> None: - """Decide whether this directory is a decay (1 -> n) or a collision. - - Read off the exported process rather than the run card: the two cannot - then disagree. MadSpin generates its decay matrix elements this way. - """ - incoming_counts = { - len(clean_pids(meta["incoming"])) for meta in self.subprocess_data - } - if incoming_counts - {1, 2}: - raise ValueError( - f"processes with {sorted(incoming_counts)} incoming particles " - "are not supported" - ) - if len(incoming_counts) > 1: - raise ValueError( - "cannot mix decays and collisions in one output directory" - ) - self.is_decay = incoming_counts == {1} - if not self.is_decay: - self.decaying_mass = None - return - - masses = { - self.get_mass(clean_pids(meta["incoming"])[0]) - for meta in self.subprocess_data - } - if len(masses) > 1: - raise ValueError( - f"decaying particles have different masses: {sorted(masses)}" - ) - self.decaying_mass = masses.pop() - if self.decaying_mass <= 0.0: - raise ValueError("the decaying particle must have a non-zero mass") - self.drop_closed_channels() - - def drop_closed_channels(self) -> None: - """Remove subprocesses the decaying particle is too light to produce. - - A multiparticle decay definition enumerates every vertex the model - allows, closed ones included: "t > b w+, w+ > all all" yields - t > b t b~ (and b W+ Z, b W+ h, ...), which need more mass than the top - has. Their partial width is exactly zero, but the phase-space mapping - has no physical point to hand back -- the invariant's lower bound ends - up above its upper bound -- so it produces NaN momenta and poisons the - whole integral. Drop them here instead. - """ - kept, dropped = [], [] - for meta in self.subprocess_data: - total = sum( - self.get_mass(pid) for pid in clean_pids(meta["outgoing"]) - ) - if total < self.decaying_mass: - kept.append(meta) - else: - dropped.append((meta["outgoing"], total)) - if dropped and self.merged_subprocess_data is not None: - # merged_subprocesses.json indexes the *unfiltered* subprocess - # list, so dropping entries here would silently shift every index - # it holds. Refuse rather than mis-map. - raise ValueError( - "merge_subprocesses is not supported for a decay directory " - "with kinematically closed channels" - ) - if dropped: - for outgoing, total in dropped: - logger.info( - "skipping closed decay channel -> %s (needs %.4g GeV, " - "the decaying particle has %.4g GeV)", - outgoing, total, self.decaying_mass, - ) - if not kept: - raise ValueError( - "every decay channel is kinematically closed: the decaying " - f"particle's mass is {self.decaying_mass} GeV" - ) - self.subprocess_data = kept - - def init_backend(self) -> None: - ms.set_simd_vector_size(self.run_card["run"]["simd_vector_size"]) - - def init_event_dir(self) -> None: - output_dir = self.run_card["gridpack"]["output_dir"] - temp_dir = self.run_card["gridpack"]["temp_output_dir"] - if output_dir: - self.run_path = os.path.abspath(os.path.expanduser(output_dir)) - os.makedirs(self.run_path, exist_ok=True) - else: - run_name = self.run_card["run"]["run_name"] - os.makedirs("Events", exist_ok=True) - run_dir_prefix = os.path.join("Events", f"{run_name}_") - existing_run_dirs = glob.glob(f"{run_dir_prefix}*") - run_index = 1 - for run_dir in existing_run_dirs: - run_index_str = run_dir[len(run_dir_prefix):] - if run_index_str.isnumeric(): - run_index = max(run_index, int(run_index_str) + 1) - while True: - try: - self.run_path = f"{run_dir_prefix}{run_index:02d}" - os.mkdir(self.run_path) - break - except FileExistsError: - run_index += 1 - # temporary npy files of the channel generators - if temp_dir: - self.temp_path = os.path.abspath(os.path.expanduser(temp_dir)) - os.makedirs(self.temp_path, exist_ok=True) - else: - self.temp_path = self.run_path - # Absolute on purpose. StatusFile is a C++ object that finishes its - # write from its destructor (rename info.json.tmp -> info.json), and - # that destructor runs whenever Python gets round to collecting the - # process object. When the launcher runs in MG5's process, an - # interrupted run is collected *after* the launch command has restored - # the working directory, and a relative path would then resolve - # somewhere else -- the rename throws a C++ filesystem_error out of a - # destructor, which is an immediate abort. - self.status_file = ms.StatusFile( - os.path.abspath(os.path.join(self.run_path, "info.json"))) - - def init_context(self) -> None: - device_names = self.run_card["run"]["device"] - self.contexts = [] - self.device_types = [] - self.devices = [] - self.pool_sizes = [] - for device_name in device_names: - if ":" in device_name: - device_type, device_index_str = device_name.split(":") - device_index = int(device_index_str) - else: - device_type = device_name - device_index = 0 - self.device_types.append(device_type) - if device_type == "cuda": - device = ms.cuda_device(device_index) - pool_size = self.run_card["run"]["gpu_thread_pool_size"] - elif device_type == "hip": - device = ms.hip_device(device_index) - pool_size = self.run_card["run"]["gpu_thread_pool_size"] - else: - device = ms.cpu_device() - pool_size = self.run_card["run"]["cpu_thread_pool_size"] - self.devices.append(device) - self.pool_sizes.append(pool_size) - self.contexts.append(ms.Context(device=device, thread_count=pool_size)) - - CUSTOM_FUNCTION_KINDS = ("cut", "scale", "pdf", "bias") - - def init_custom_functions(self) -> None: - """Load the functions of [custom_functions], building the libraries - first. They are loaded before the matrix elements, so that they have the - same index in every context and in a gridpack made from this run.""" - self.custom_functions = {} - self.custom_function_paths = [] - for kind in self.CUSTOM_FUNCTION_KINDS: - path = self.run_card["custom_functions"][kind] - if not path: - continue - build_dir = f"custom_{kind}" - if path.endswith(".so") and os.path.isfile(os.path.join(build_dir, "Makefile")): - self.build_custom_function(build_dir) - if not os.path.isfile(path): - raise RuntimeError( - f"[custom_functions] {kind} = {path!r}: file not found. " - "Create a .so library or a .py file with bin/init_custom_func." - ) - apis = load_custom_function(ms, self.contexts, self.device_types, kind, path) - if len({api.index() for api in apis}) != 1: - raise RuntimeError(f"custom {kind} function loaded at different indices") - self.custom_functions[kind] = apis[0] - self.custom_function_paths.append([kind, path]) - logger.info("Using the custom %s function %s", kind, path) - - def init_surrogate(self) -> None: - """Check the [surrogate] settings. The surrogates are loaded per - subprocess, by MadgraphSubprocess.""" - settings = self.run_card["surrogate"] - self.surrogate = None - if not settings["function"]: - return - path = settings["function"] - if path.endswith(".py") and not os.path.isfile(path): - raise RuntimeError(f"[surrogate] function = {path!r}: file not found") - if self.run_card["phasespace"]["sde_strategy"] != "denominators": - raise RuntimeError( - "[surrogate] needs [phasespace] sde_strategy = \"denominators\", " - "as the surrogate does not return channel weights" - ) - if self.run_card["run"]["dummy_matrix_element"]: - raise RuntimeError("[surrogate] cannot be used with dummy_matrix_element") - if self.run_card["gridpack"]["save_gridpack"]: - raise RuntimeError("[surrogate] does not support gridpacks yet") - self.surrogate = dict(settings) - if not self.surrogate["param_card"]: - self.surrogate["param_card"] = self.param_card_path - logger.info( - "Using the surrogate %s, %s with the full matrix element", - path, - "unweighting" if settings["mode"] == "unweight" else "reweighting", - ) - - def weighting_groups(self, subprocess_channels: list[list[int]]) -> list: - """One ms.WeightingGroup per subprocess with a [surrogate]: the second - pass with the full matrix element over its channels, given by their - indices in the event generator.""" - if self.surrogate is None: - return [] - gen_args = self.run_card["generation"] - return [ - ms.WeightingGroup( - channel_indices, - [ms.WeightingStep( - subproc.step_matrix_elements, - mode=self.surrogate["mode"], - max_overweight_truncation=gen_args["max_overweight_truncation"], - freeze_max_weight_after=gen_args["freeze_max_weight_after"], - )], - ) - for subproc, channel_indices in zip(self.subprocesses, subprocess_channels) - if channel_indices - ] - - @staticmethod - def build_custom_function(build_dir: str) -> None: - """Run make in a custom_ directory, logging to compile.log.""" - log_path = os.path.join(build_dir, "compile.log") - with open(log_path, "w") as log: - returncode = subprocess.call( - ["make"], cwd=build_dir, stdout=log, stderr=subprocess.STDOUT - ) - if returncode != 0: - raise RuntimeError( - f"Compilation of {build_dir} failed with exit status {returncode}; " - f"see {log_path} for details" - ) - - def custom_function(self, kind: str, particle_count: int, pid_options: list): - """The ms.CustomFunction of [custom_functions] , or None.""" - api = self.custom_functions.get(kind) - if api is None: - return None - function_type = getattr(ms.CustomFunction, kind) - return ms.CustomFunction(api, function_type, particle_count, pid_options) - - def parse_observable(self, name: str, order_observable: str) -> dict: - parts = name.split("-") - sum_momenta = False - sum_observable = False - ordered = False - multiparticles = self.run_card["multiparticles"] - - if len(parts) == 0: - raise ValueError("Invalid observable name") - elif len(parts) == 1: - # event-level observables - obs_name = parts[0] - select_pids = [] - else: - if parts[-1] == "sum": - sum_observable = True - obs_name = parts[-2] - selection = parts[:-2] - elif parts[-2] == "sum": - sum_momenta = True - obs_name = parts[-1] - selection = parts[:-2] - else: - obs_name = parts[-1] - selection = parts[:-1] - select_pids = [] - order_indices = [] - for mp_name in selection: - mp_parts = mp_name.split("_") - if mp_parts[-1].isnumeric(): - order_indices.append(int(mp_parts[-1])) - select_pids.append(multiparticles["_".join(mp_parts[:-1])]) - ordered = True - else: - order_indices.append(0) - select_pids.append(multiparticles[mp_name]) - - return dict( - observable=obs_name, - select_pids=select_pids, - sum_momenta=sum_momenta, - sum_observable=sum_observable, - order_observable=order_observable if ordered else None, - order_indices=order_indices if ordered else [], - ignore_incoming=True, - name=name, - ) - - def init_cuts(self) -> None: - inf = float("inf") - order_observable = self.run_card["cuts"].get("order_by", "pt") - self.cut_data = [ - CutItem( - observable_kwargs=self.parse_observable(key, order_observable), - min=values.get("min", -inf), - max=values.get("max", inf), - mode=values.get("mode", "all"), - ) - for key, values in self.run_card["cuts"].items() - if key != "order_by" - ] - - # [histograms] key that means "the distribution of the event weight", - # normalised to the unit weight sigma_abs, rather than an observable of the - # momenta: an unweighted sample is a spike at 1 (and at -1 for its negative - # weights) and a partially unweighted one shows its spread. - weight_histogram_key = "weight" - - def init_histograms(self) -> None: - order_observable = self.run_card["histograms"].get("order_by", "pt") - self.hist_data = [] - for key, values in self.run_card["histograms"].items(): - if key == "order_by": - continue - from_weight = key == self.weight_histogram_key - self.hist_data.append(HistItem( - observable_kwargs=( - None if from_weight - else self.parse_observable(key, order_observable)), - min=values["min"], - max=values["max"], - bin_count=values["bin_count"], - name=key, - from_weight=from_weight, - )) - - def ensure_pdf_set(self, pdf_set: str) -> misc.LhapdfPaths: - """Make sure the requested LHAPDF set is available, downloading it into - the first writable PDF directory if needed, and return the (possibly - updated) LHAPDF resolution. Nothing is downloaded when the set is - already present.""" - paths = lhapdf_paths() - if paths.find_set(pdf_set): - return paths - if not paths.config: - logger.debug("no usable lhapdf-config: cannot download %s", pdf_set) - return paths - if not paths.download_path: - logger.debug("no writable PDF directory: cannot download %s", pdf_set) - return paths - try: - from madgraph.interface.common_run_interface import CommonRunCmd - os.makedirs(paths.download_path, exist_ok=True) - logger.info("PDF set %s not found; downloading into %s", - pdf_set, paths.download_path) - CommonRunCmd.install_lhapdf_pdfset_static( - paths.config, paths.download_path, pdf_set) - except Exception as err: - logger.warning("Could not download PDF set %s: %s", pdf_set, err) - return paths - global _LHAPDF - _LHAPDF = paths.with_data_path(paths.download_path) - return _LHAPDF - - def _no_pdf_message(self, pdf_set: str) -> str: - """Explain a missing PDF set in terms of the settings that fix it.""" - searched = os.pathsep.join(self.lhapdf.data_paths) or \ - "(no PDF set directory could be located)" - if self.lhapdf.config: - found = "using lhapdf-config: %s" % self.lhapdf.config - else: - found = ("no usable lhapdf-config was found, so the set could not " - "be downloaded automatically") - return ( - "LHAPDF set %r, requested by [beam] pdf in Cards/run_card.toml, " - "was not found.\n" - " searched: %s\n" - " %s\n" - "Fix this in one of the following ways:\n" - " * install LHAPDF for this MadGraph: MG7> install lhapdf6\n" - " * point MadGraph at an existing one:\n" - " MG7> set lhapdf /path/to/lhapdf-config\n" - " MG7> save options\n" - " * or set it for this directory only, by adding\n" - " lhapdf = /path/to/lhapdf-config\n" - " to %s\n" - " * or set $LHAPDF_DATA_PATH to a directory containing %r." - % (pdf_set, searched, found, - os.path.join("Cards", "me5_configuration.txt"), pdf_set)) - - def init_beam(self) -> None: - beam_args = self.run_card["beam"] - - # built lazily by build_systematics (needs the subprocess data) - self.systematics = None - self.systematics_data = None - self.systematics_context = None - self.event_histograms = None - self.event_histograms_context = None - - if self.is_decay: - # No beams: the total energy is the decaying particle's mass, and - # "leptonic" is what the mappings call "no parton luminosity". - self.e_cm = self.decaying_mass - self.leptonic = True - else: - self.e_cm = beam_args["e_cm"] - self.leptonic = beam_args["leptonic"] - - dynamical_scales = { - "transverse_energy": ms.EnergyScale.transverse_energy, - "transverse_mass": ms.EnergyScale.transverse_mass, - "half_transverse_mass": ms.EnergyScale.half_transverse_mass, - "partonic_energy": ms.EnergyScale.partonic_energy, - } - if beam_args["dynamical_scale_choice"] in dynamical_scales: - dynamical_scale_type = dynamical_scales[beam_args["dynamical_scale_choice"]] - else: - raise ValueError("Unknown dynamical scale choice") - self.scale_kwargs = dict( - dynamical_scale_type=dynamical_scale_type, - ren_scale_fixed=beam_args["fixed_ren_scale"], - fact_scale_fixed=beam_args["fixed_fact_scale"], - ren_scale=beam_args["ren_scale"], - fact_scale1=beam_args["fact_scale1"], - fact_scale2=beam_args["fact_scale2"], - scale_factor=float(beam_args.get("scale_factor", 1.)), - ) - if self.is_decay and self.custom_functions.pop("scale", None) is not None: - logger.warning("custom scale ignored: a decay uses the decaying mass") - if self.is_decay and self.custom_functions.pop("pdf", None) is not None: - logger.warning("custom pdf ignored: a decay has no parton densities") - if self.is_decay: - # One scale is available for a decay -- the decaying mass -- so use - # it, fixed, whatever the card asks for. - self.scale_kwargs.update( - ren_scale_fixed=True, - fact_scale_fixed=True, - ren_scale=self.decaying_mass, - fact_scale1=self.decaying_mass, - fact_scale2=self.decaying_mass, - scale_factor=1., - ) - self.pdf_grid = None - self.lhapdf = None - self.pdf_dir = None - self.fixed_alphas_info = self.write_fixed_alphas_info() - self.alphas_grid = ms.AlphaSGrid(self.fixed_alphas_info) - for context in self.contexts: - self.alphas_grid.initialize_globals(context) - self.running_coupling = ms.RunningCoupling(self.alphas_grid) - return - - pdf_set = beam_args["pdf"] - self.pdf_set = pdf_set - self.fixed_alphas_info = None - self.lhapdf = self.ensure_pdf_set(pdf_set) - self.pdf_dir = self.lhapdf.find_set(pdf_set) - if self.pdf_dir is None: - raise RuntimeError(self._no_pdf_message(pdf_set)) - self.alphas_grid = ms.AlphaSGrid(os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}.info")) - if "pdf" in self.custom_functions: - # alpha_s still comes from the PDF set of the run card - self.pdf_grid = ms.CustomPdf(self.custom_functions["pdf"]) - else: - self.pdf_grid = ms.PdfGrid( - os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}_0000.dat") - ) - for context in self.contexts: - if not isinstance(self.pdf_grid, ms.CustomPdf): - self.pdf_grid.initialize_globals(context) - self.alphas_grid.initialize_globals(context) - self.running_coupling = ms.RunningCoupling(self.alphas_grid) - - # ------------------------------------------------------------------ - # scale / PDF systematics - # ------------------------------------------------------------------ - def systematics_enabled(self) -> bool: - """The native scale/PDF weights are requested ([systematics] enable, or - the legacy [generation] systematics alias).""" - try: - enabled = bool(self.run_card["systematics"]["enable"]) - except Exception: - enabled = False - if not enabled: - try: - enabled = bool(self.run_card["generation"]["systematics"]) - except Exception: - enabled = False - if enabled and "pdf" in getattr(self, "custom_functions", {}): - # the variations re-evaluate the PDF set, which a custom PDF replaces - if not getattr(self, "_warned_custom_pdf_systematics", False): - logger.warning("systematics are disabled with a custom pdf") - self._warned_custom_pdf_systematics = True - return False - return enabled - - @staticmethod - def pdf_set_info(info_path: str) -> dict: - """Read SetIndex, NumMembers, ErrorType and SetDesc from an LHAPDF - .info file (missing keys get neutral defaults).""" - info = {"SetIndex": 0, "NumMembers": 1, "ErrorType": "", "SetDesc": ""} - try: - with open(info_path) as f: - for line in f: - key, sep, value = line.partition(":") - key = key.strip() - if not sep or key not in info: - continue - value = value.strip().strip('"').strip("'") - if key in ("SetIndex", "NumMembers"): - info[key] = int(value) - else: - info[key] = value - except OSError as err: - logger.warning("could not read %s: %s", info_path, err) - return info - - def resolve_pdf_set_name(self, entry: str) -> tuple[str, int | None]: - """Turn a [systematics] pdf entry ("", "", "@", - "@") into (set name, member or None) using the LHAPDF - pdfsets.index when an id is given.""" - member = None - if "@" in entry: - entry, member_str = entry.split("@", 1) - member = int(member_str) - entry = entry.strip() - if not entry.isdigit(): - return entry, member - lhaid = int(entry) - index_path = os.path.join(self.pdf_dir or "", "pdfsets.index") - sets = [] - try: - with open(index_path) as f: - for line in f: - parts = line.split() - if len(parts) >= 2 and parts[0].isdigit(): - sets.append((int(parts[0]), parts[1])) - except OSError: - raise ValueError( - "cannot resolve LHAPDF id %s: no %s" % (lhaid, index_path)) - for set_id, name in sets: - if set_id == lhaid: - return name, member - # an id inside a set's member range - sets.sort() - for (set_id, name), (next_id, _) in zip(sets, sets[1:] + [(None, None)]): - if set_id < lhaid and (next_id is None or lhaid < next_id): - return name, lhaid - set_id if member is None else member - raise ValueError("cannot resolve LHAPDF id %s" % lhaid) - - def resolve_pdf_variations(self) -> list: - """Expand the [systematics] pdf entries into ms.PdfMemberSpec objects - (downloading the sets when needed). The nominal member itself is never - listed: the calculator adds it as the central weight of its group.""" - entries = self.run_card["systematics"]["pdf"] - if isinstance(entries, str): - entries = [entries] - specs = [] - seen = set() - for raw in entries: - entry = str(raw).strip() - if not entry or entry.lower() in ("none", "no", "false"): - continue - if entry.lower() == "central": - continue - if entry.lower() == "errorset": - set_name, member = self.pdf_set, None - else: - set_name, member = self.resolve_pdf_set_name(entry) - variation_lhapdf = self.ensure_pdf_set(set_name) - variation_dir = variation_lhapdf.find_set(set_name) - if variation_dir is None: - logger.warning("PDF set %s not available; its variation is skipped", set_name) - continue - info_path = os.path.join(variation_dir, set_name, f"{set_name}.info") - if not os.path.exists(info_path): - logger.warning("PDF set %s not available; its variation is skipped", set_name) - continue - info = self.pdf_set_info(info_path) - is_nominal = set_name == self.pdf_set - if member is not None: - members = [member] - else: - members = list(range(info["NumMembers"])) - if is_nominal: - members = members[1:] - if not members: - logger.warning( - "PDF set %s has a single member: no PDF uncertainty can be " - "computed from it. Use a set with error members in " - "[systematics] pdf (e.g. a PDF4LHC or NNPDF error set).", - set_name) - for m in members: - if is_nominal and m == 0: - continue - if (set_name, m) in seen: - continue - seen.add((set_name, m)) - grid_file = os.path.join(variation_dir, set_name, f"{set_name}_{m:04d}.dat") - if not os.path.exists(grid_file): - logger.warning("PDF member file %s not found; skipped", grid_file) - continue - specs.append(ms.PdfMemberSpec( - set_name=set_name, set_lhaid=info["SetIndex"], member=m, - grid_file=grid_file, info_file=info_path, - error_type=info["ErrorType"], description=info["SetDesc"], - )) - return specs - - def build_systematics_args(self) -> list: - """One ms.SubprocessSystArgs per (unmerged) subprocess: the alpha_s power - of |M|^2 and the beam parton ids of every flavor index, indexed like the - LHECompleter (subprocess_index / flavor_index of the combined events).""" - args = [] - for meta in self.subprocess_data: - beam_pdgs = [ - [int(flavor["options"][0][0]), int(flavor["options"][0][1])] - for flavor in meta["flavors"] - ] - args.append(ms.SubprocessSystArgs( - qcd_power=int(meta.get("qcd_power", -1)), beam_pdgs=beam_pdgs)) - return args - - def build_systematics(self): - """Create the ms.SystematicsCalculator from the [systematics] section, or - None when disabled. Also keeps the JSON description needed by gridpacks.""" - if not self.systematics_enabled(): - self.systematics = None - return None - syst = self.run_card["systematics"] - config = ms.SystematicsConfig() - config.mur = [float(v) for v in syst["mur"]] - config.muf = [float(v) for v in syst["muf"]] - config.together = bool(syst["together"]) - config.dyn_scales = self.resolve_dynamical_scales() - # alternative dynamical scales get the same factor as the nominal one - config.scale_factor = float(self.run_card["beam"].get("scale_factor", 1.)) - config.write_inputs = bool(syst["write_inputs"]) - config.has_pdf = not self.leptonic - # Without parton luminosity -- a decay, or leptonic beams -- there is no - # nominal PDF *set* to describe and no PDF variation to compute, so - # has_pdf gates the nominal-set metadata and the member list below. - # The alpha_s grid is a separate matter: the mu_R variations always need - # one, so its .info file has to be recorded whatever has_pdf says. - # init_beam() skips the PDF lookup for a decay only -- leptonic beams - # still resolve a set and take alpha_s from it -- so the path follows - # self.pdf_dir, and falls back to the constant grid written for a decay. - if self.pdf_dir is not None: - info_path = os.path.join(self.pdf_dir, self.pdf_set, f"{self.pdf_set}.info") - else: - info_path = self.fixed_alphas_info - if config.has_pdf: - info = self.pdf_set_info(info_path) - config.nominal_set_name = self.pdf_set - config.nominal_lhaid = info["SetIndex"] - config.nominal_error_type = info["ErrorType"] - config.nominal_description = info["SetDesc"] - config.pdf_members = self.resolve_pdf_variations() - else: - config.pdf_members = [] - args = self.build_systematics_args() - # the PDFs, alpha_s and (for mixed-order subprocesses) the matrix - # elements are evaluated with the batched madspace functions on this - # CPU context - self.systematics_context = ms.Context(device=ms.cpu_device(), thread_count=1) - matrix_elements, flavor_remap, me_backend = self.build_systematics_matrix_elements( - args, self.systematics_context) - self.systematics = ms.SystematicsCalculator( - config, args, self.pdf_grid, self.alphas_grid, - self.systematics_context, matrix_elements, flavor_remap) - for warning in self.systematics.warnings: - logger.warning("systematics: %s", warning) - self.systematics_data = { - "config": json.loads(config.to_json()), - "subproc_args": [json.loads(a.to_json()) for a in args], - "nominal_grid_file": os.path.join( - self.pdf_dir, self.pdf_set, f"{self.pdf_set}_0000.dat") - if config.has_pdf else None, - "nominal_info_file": info_path, - # matrix elements re-evaluated for the mixed-order subprocesses - "me_paths": [meta["me_path"] for meta in self.subprocess_data], - "me_backend": me_backend, - "flavor_remap": flavor_remap, - "me_frame": self.me_frame, - "incoming_count": self.incoming_count, - } - return self.systematics - - _DYN_SCALE_CODES = { - "transverse_energy": 1, "transverse_mass": 2, - "half_transverse_mass": 3, "partonic_energy": 4, - } - - def resolve_dynamical_scales(self) -> list: - """[systematics] dynamical_scale entries (names or systematics.py codes - 1-4) as codes, without the choice the events were generated with.""" - try: - entries = self.run_card["systematics"]["dynamical_scale"] - except Exception: - return [] - if isinstance(entries, str): - entries = [entries] - beam = self.run_card["beam"] - generated = None - if not (beam["fixed_ren_scale"] and beam["fixed_fact_scale"]): - generated = self._DYN_SCALE_CODES.get(beam["dynamical_scale_choice"]) - codes = [] - for raw in entries: - entry = str(raw).strip().lower() - if not entry: - continue - if entry.isdigit(): - code = int(entry) - elif entry in self._DYN_SCALE_CODES: - code = self._DYN_SCALE_CODES[entry] - else: - raise ValueError("unknown dynamical scale choice %r in [systematics] " - "dynamical_scale" % raw) - if code not in (1, 2, 3, 4): - raise ValueError("dynamical scale code %s must be 1-4" % code) - if code == generated or code in codes: - continue - codes.append(code) - return codes - - def build_systematics_matrix_elements(self, args, context): - """For the subprocesses whose |M|^2 mixes several powers of alpha_s - (qcd_power = -1), load the matrix element into the systematics CPU - context so that the renormalisation scale variations can re-evaluate it - at the varied alpha_s. Returns ([MatrixElement or None per subprocess], - flavor remap, backend name); empty when nothing is needed or no CPU - library exists.""" - need = [i for i, a in enumerate(args) if a.qcd_power < 0] - empty = ([], [], None) - if not need or self.run_card["run"]["dummy_matrix_element"]: - return empty - cpu_backends = [b for b in self.backends if not b.startswith(("cuda", "hip"))] - if not cpu_backends: - logger.warning("systematics: no CPU matrix-element library available, the " - "renormalisation scale variations of the mixed-order " - "subprocesses cannot be computed") - return empty - backend = cpu_backends[0] - matrix_elements = [] - flavor_remap = [] - for i, meta in enumerate(self.subprocess_data): - flavor_remap.append([int(f["index"]) for f in meta["flavors"]]) - if i not in need: - matrix_elements.append(None) - continue - api = context.load_matrix_element( - meta["me_path"].format(device=backend), self.param_card_path, - me_parameters(self.run_card)) - matrix_elements.append(ms.MatrixElement( - api, - [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, - ms.MatrixElement.flavor_in], - [ms.MatrixElement.matrix_element_out], - False, - self.me_frame, - self.incoming_count, - )) - logger.info("systematics: re-evaluating the matrix element of %d mixed-order " - "subprocess(es) for the renormalisation scale variations", len(need)) - return matrix_elements, flavor_remap, backend - - def build_event_histograms(self): - """ms.EventHistograms for the [histograms] observables, filled with the - written events and all their variation weights (info.json - "event_histograms"); None without histograms.""" - if not self.hist_data \ - or not self.run_card["run"]["postprocessing_histograms"]: - self.event_histograms = None - return None - context = ms.Context(device=ms.cpu_device(), thread_count=1) - specs = [ - ms.EventHistogramSpec( - name=item.name, min=item.min, max=item.max, - bin_count=item.bin_count, from_weight=item.from_weight) - for item in self.hist_data - ] - from_momenta = [item for item in self.hist_data if not item.from_weight] - observables = [] - for meta in self.subprocess_data: - if not from_momenta: - # nothing to evaluate on the momenta (only the weight is - # histogrammed): no observable runtime to build - observables.append(None) - continue - all_pids = clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) - values = ms.ObservableValues([ - build_observable( - observable_pids(meta, item.observable_kwargs), - len(meta["incoming"]), item.observable_kwargs - ) - for item in from_momenta - ]) - observables.append(ms.SubprocessObservables(values, len(all_pids))) - # the weight histograms are drawn in units of the unit weight sigma_abs - # (the cross section of a sample without negative weights) - self.event_histograms_context = context - self.event_histograms = ms.EventHistograms( - context, specs, observables, - reference_weight=self._mean_event_weight()) - return self.event_histograms - - def _mean_event_weight(self) -> float: - """The unit weight of the unweighted events, sigma_abs, used as the - unit of the weight histograms: the cross section of a positive sample, - and for one with negative weights (an interference) what puts them at - +-1 rather than dividing by a signed, possibly small, total. 0 (the raw - weight) when it is not available.""" - try: - return float(self.event_generator.status().mean_abs) - except Exception as error: - logger.debug("no cross section for the weight histograms: %s", error) - return 0. - - def log_systematics_summary(self) -> None: - """Print the scale/PDF uncertainties on the total cross section (the - per-variation cross sections are in events.weights.json / info.json).""" - if self.systematics is None or self.systematics.weight_count == 0: - return - summary = json.loads(self.systematics.summary()) - # the percentages come from systematics_summary, which get_result() - # reads too: the box and the scan row must not disagree, so neither - # computes them itself - xsec = systematics_summary.nominal_cross_section(summary) - if xsec is None: - return - def format_variation(up, down): - # right-justify the signed numbers (not the sign alone) so the % - # values line up across rows, with extra spacing between the two - return f"{'+%.3g' % up:>6}% {'-%.3g' % down:>6}%" - - scale_count = sum(1 for v in self.systematics.variations if v.is_scale) - member_count = len(self.systematics.members) - rows = [("Variations per event:", - f"{self.systematics.weight_count} ({scale_count} scale, {member_count} PDF members)")] - for pdf in summary.get("pdf", []): - rows.append(("PDF set:", f"{pdf['pdf_set']}, {pdf['error_type']}")) - rows.append(("Original cross-section:", f"{xsec} pb")) - scale = systematics_summary.scale_percentages(summary) - if scale is not None: - rows.append(("Scale variation:", format_variation(*scale))) - for _entry, up, down in systematics_summary.pdf_percentages(summary): - rows.append(("PDF variation:", format_variation(up, down))) - if self.event_generator_config.verbosity == ms.Verbosity.pretty: - box = ms.PrettyBox("Systematics", len(rows), [24, 0]) - box.set_column(0, [label for label, _ in rows]) - box.set_column(1, [value for _, value in rows]) - box.print_first() - else: - for label, value in rows: - logger.info("systematics, %s %s", label, value) - - def write_fixed_alphas_info(self) -> str: - """Write a minimal LHAPDF ``.info`` holding a constant alpha_s. - - A decay has no beams, so there is no PDF set to take alpha_s from -- - and demanding one (possibly downloading it) just to evaluate a coupling - would be absurd. The renormalisation scale of a decay is fixed at the - decaying mass anyway, so a constant alpha_s is the right answer, not an - approximation: take it from the param card, exactly as the Fortran - decay matrix elements do. - """ - # SMINPUTS entry 3 is alpha_s(m_Z), the same value the Fortran - # parameter setup feeds to G. - alpha_s = float(self.param_card.get_value("sminputs", 3)) - path = os.path.join(self.run_path, "fixed_alphas.info") - # AlphaSGrid interpolates in log(q^2) across the grid, so give it a - # comfortable number of nodes rather than the bare minimum of three. - q_values = [10 ** (i / 4.0) for i in range(-4, 21)] - with open(path, "w") as f: - f.write("SetDesc: constant alpha_s for a decay (no beams)\n") - f.write("AlphaS_Qs: [%s]\n" % ", ".join(f"{q:g}" for q in q_values)) - f.write( - "AlphaS_Vals: [%s]\n" - % ", ".join(f"{alpha_s:g}" for _ in q_values) - ) - logger.info( - "decay mode: fixed alpha_s = %g at mu = %g GeV", alpha_s, self.e_cm - ) - return path - - def needs_systematics_matrix_elements(self) -> bool: - """The systematics will re-evaluate a matrix element per scale point. - - True when the native weights are on and some subprocess mixes alpha_s - powers (qcd_power < 0), i.e. when build_systematics_matrix_elements() - loads a matrix element. Conservative: it does not check that a CPU - backend exists, so it can be True where that call ends up loading - nothing. - """ - if not self.systematics_enabled(): - return False - if self.run_card["run"]["dummy_matrix_element"]: - return False - return any(int(meta.get("qcd_power", -1)) < 0 - for meta in self.subprocess_data) - - def init_generator_config(self) -> None: - run_args = self.run_card["run"] - gen_args = self.run_card["generation"] - vegas_args = self.run_card["vegas"] - cfg = ms.GeneratorConfig() - cfg.target_count = gen_args["events"] - cfg.vegas_damping = vegas_args["damping"] - cfg.max_overweight_truncation = gen_args["max_overweight_truncation"] - cfg.freeze_max_weight_after = gen_args["freeze_max_weight_after"] - if self.surrogate is not None and self.surrogate["mode"] == "unweight": - # [generation] applies to the second pass, which gives the events - cfg.max_overweight_truncation = ( - self.surrogate["first_pass_max_overweight_truncation"] - ) - cfg.freeze_max_weight_after = ( - self.surrogate["first_pass_freeze_max_weight_after"] - ) - cfg.start_batch_size = vegas_args["start_batch_size"] - cfg.max_batch_size = vegas_args["max_batch_size"] - cfg.survey_min_iters = gen_args["survey_min_iters"] - cfg.survey_max_iters = gen_args["survey_max_iters"] - cfg.survey_target_precision = gen_args["survey_target_precision"] - cfg.optimization_patience = vegas_args["optimization_patience"] - cfg.optimization_threshold = vegas_args["optimization_threshold"] - cfg.cpu_batch_size = gen_args["cpu_batch_size"] - cfg.gpu_batch_size = gen_args["gpu_batch_size"] - cfg.verbosity = resolve_verbosity(run_args["verbosity"]) - cfg.combine_thread_count = run_args["combine_thread_pool_size"] - if self.needs_systematics_matrix_elements() and cfg.combine_thread_count != 1: - # SystematicsCalculator::compute() runs inside the combine workers, - # and a matrix element's per-thread process instances live in a - # ThreadResource indexed by ThreadPool::thread_index() -- the - # *calling* thread's index. A parallel combine pool would therefore - # need one process instance per worker; building N of them to index - # them safely is not worth it, the more so as - # SystematicsCalculator::matrix_elements() already serialises every - # evaluation on its own mutex, so that parallelism was not being - # used anyway. Run the combine stage single-threaded instead. - logger.info( - "systematics: the combine stage runs single-threaded because a " - "subprocess needs its matrix element re-evaluated for the " - "renormalisation scale variations") - cfg.combine_thread_count = 1 - cfg.cut_efficiency_threshold = gen_args["cut_efficiency_threshold"] - cfg.max_cut_repetitions = gen_args["max_cut_repetitions"] - self.event_generator_config = cfg - self.event_generator = None - - def init_subprocesses(self) -> None: - self.backends = self.compile_matrix_elements() - self.subprocesses = [] - if self.merged_subprocess_data is None: - for subproc_id, meta in enumerate(self.subprocess_data): - self.subprocesses.append(MadgraphSubprocess(self, meta, subproc_id)) - else: - for subproc_id, meta in enumerate(self.merged_subprocess_data): - self.subprocesses.append( - MadgraphSubprocess(self, meta, subproc_id, self.subprocess_data) - ) - - def compile_matrix_elements(self) -> list[str]: - """Build the matrix-element library of every subprocess, and return the - build backend of each requested device (in the same order). - - A 'cpu' device builds the backend named by the 'cpu_mode' entry; the - cuda/hip devices ignore cpu_mode and build their own backend. A - cpu_mode of 'auto' is resolved to the actual SIMD backend chosen on - this machine. - - SubProcesses/makefile is a dispatcher over the P* directories, so a - single 'make -j N' there builds all the subprocesses at once with one - shared pool of N jobs: no subprocess is built with N jobs while the - others wait, and none is limited to N/#subprocesses jobs either. - """ - device_names = self.run_card["run"]["device"] - if not isinstance(device_names, list): - device_names = [device_names] - cpu_mode = self.run_card["run"]["cpu_mode"] - backends = [backend_of(name, cpu_mode) for name in device_names] - if not self.subprocess_data: - return backends - - first_proc_path = self.subprocess_data[0]["path"] - subproc_path = os.path.dirname(first_proc_path) - - fptype = fptype_of(self.run_card["run"]["precision"]) - - # Resolve cpu_mode='auto' once (the build rules pick the best SIMD - # backend available here), so all subprocesses agree on the library names. - auto_backend = None - if "auto" in backends: - auto_backend = resolve_auto_backend(first_proc_path) - logger.info("cpu_mode 'auto' resolved as '%s'", auto_backend) - resolved = [ - auto_backend if backend == "auto" else backend - for backend in backends - ] - - nb_core = self.run_card["run"]["cpu_thread_pool_size"] - if not nb_core or nb_core < 0: - nb_core = os.cpu_count() or 1 - - log_path = os.path.join(self.run_path, "compile_subprocesses.log") - for backend in resolved: - missing = [ - meta for meta in self.subprocess_data - if not os.path.isfile(meta["me_path"].format(device=backend)) - ] - if not missing: - continue - logger.info( - f"Start compilation of SubProcesses for device '{backend}' " - f"({len(missing)} subprocess(es), {nb_core} parallel job(s)), " - f"see log detail in {log_path}" - ) - start_time = time.time() - self.make_subprocesses( - subproc_path, - [f"BACKEND={backend}", f"FPTYPE={fptype}", "USEBUILDDIR=1"], - nb_core, log_path, - ) - logger.info( - f"Compilation of SubProcesses done in {time.time() - start_time:.1f} s" - ) - return resolved - - @staticmethod - def make_subprocesses( - subproc_path: str, args: list[str], nb_core: int, log_path: str - ) -> None: - """Run 'make -j nb_core' in SubProcesses/, appending the full build - output to log_path (one run per device, so the file is appended to).""" - command = ["make", f"-j{nb_core}"] + args - with open(log_path, "a") as log: - log.write(f"\n$ cd {subproc_path} && {' '.join(command)}\n") - log.flush() - returncode = subprocess.call( - command, cwd=subproc_path, stdout=log, stderr=subprocess.STDOUT - ) - if returncode != 0: - raise RuntimeError( - f"Compilation of the SubProcesses failed with exit status " - f"{returncode}; see {log_path} for details" - ) - - def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenerator: - channel_generators = [] - subprocess_channels = [[] for _ in self.subprocesses] - for i, (subproc, phasespace) in enumerate(zip(self.subprocesses, phasespaces)): - for integrand, channel in zip( - subproc.build_integrands(phasespace), phasespace.channels - ): - if channel.event_generator is None: - channel.event_generator = ms.ChannelEventGenerator( - contexts=self.contexts, - integrand=integrand, - event_file=os.path.join(self.temp_path, f"events.{i}.{channel.name}.npy"), - weight_file=os.path.join(self.temp_path, f"weights.{i}.{channel.name}.npy"), - config=self.event_generator_config, - subprocess_index=i, - name=f"{i}.{channel.name}", - histograms=subproc.histograms - ) - subprocess_channels[i].append(len(channel_generators)) - channel_generators.append(channel.event_generator) - - event_generator = ms.EventGenerator( - contexts=self.contexts, - channels=channel_generators, - status_file=self.status_file, - config=self.event_generator_config, - seed=self.run_seed, - groups=self.weighting_groups(subprocess_channels), - ) - unused_globals = ( - set(self.contexts[0].global_names()) - event_generator.used_globals() - ) - for context in self.contexts: - for global_name in unused_globals: - context.delete_global(global_name) - return event_generator - - def survey_phasespaces( - self, phasespaces: list[PhaseSpace | None], survey_pass: int = 0 - ) -> ms.EventGenerator | None: - ps_filtered = [ps for ps in phasespaces if ps is not None] - if len(ps_filtered) == 0: - return None - event_generator = self.build_event_generator(ps_filtered) - event_generator.survey(survey_pass) - return event_generator - - def survey(self) -> None: - # survey_pass distinguishes the survey() calls below: "both" mode can - # re-survey a channel carried over unchanged from the multichannel pass - # into the final (simplified) pass, and both passes schedule jobs on the - # same underlying ChannelEventGenerator. The explicit pass index keeps each - # pass's job seeds independent of the other passes' job counts, rather than - # depending on call history. - phasespace_mode = self.run_card["phasespace"]["mode"] - if self.is_decay and phasespace_mode != "multichannel": - # The flat mapping is built from a synthetic two-incoming diagram, - # which a decay has no counterpart for -- and 'both'/'auto' end up - # simplifying towards it. Decays have few channels, so the - # multichannel phase space is the right one regardless. - logger.info( - "decay mode: using the multichannel phase space instead of " - "'%s'", phasespace_mode - ) - phasespace_mode = "multichannel" - if phasespace_mode in ["multichannel", "both", "auto"]: - self.phasespaces = [ - subproc.build_multichannel_phasespace() - for subproc in self.subprocesses - ] - self.event_generator = self.survey_phasespaces(self.phasespaces, 0) - elif phasespace_mode == "flat": - self.phasespaces = [ - subproc.build_flat_phasespace() - for subproc in self.subprocesses - ] - self.event_generator = self.survey_phasespaces(self.phasespaces, 0) - else: - raise ValueError("Unknown phasespace mode") - - channel_status = self.event_generator.channel_status() - chan_offset = 0 - madnis_enabled = False - for subproc, ps in zip(self.subprocesses, self.phasespaces): - mean = 0. - variance = 0. - count_opt = 0 - for status in channel_status[chan_offset:chan_offset + len(ps.channels)]: - mean += status.mean_abs - variance += status.error_abs**2 - count_opt += status.count_opt - # A subprocess whose matrix element is identically zero (an FCNC - # channel with every Wilson coefficient at zero, say) integrates to - # exactly zero, so there is no relative spread to speak of. Report 0 - # rather than dividing by it: that asks for the smallest networks and - # leaves MadNIS off, which is what an empty subprocess wants. - rsd = (variance * count_opt)**0.5 / mean if mean else 0. - subproc.set_madnis_auto_settings(rsd) - chan_offset += len(ps.channels) - if subproc.madnis_settings["enable"]: - madnis_enabled = True - if not ( - phasespace_mode == "both" or (phasespace_mode == "auto" and madnis_enabled) - ): - return - - phasespaces_multi = self.phasespaces - cross_sections = [] - index = 0 - for phasespace in phasespaces_multi: - channel_count = len(phasespace.channels) - # rank the channels by the integral of |w|: a channel whose signed - # integral cancels (an interference) can still carry most of the - # variance, and folding it into the flat channel would blow it up - cross_sections.append([ - status.mean_abs - for status in channel_status[index:index + channel_count] - ]) - index += channel_count - - self.phasespaces = [ - subproc.simplify_phasespace(ps_multi, cross_secs) - for subproc, ps_multi, cross_secs in zip( - self.subprocesses, phasespaces_multi, cross_sections - ) - ] - if any( - ps_multi is not ps_both - for ps_multi, ps_both in zip(phasespaces_multi, self.phasespaces) - ): - # distinct survey_pass: a channel carried over unchanged from the - # multichannel pass (pass 0) into this resurvey must not share its - # seed stream with that earlier pass. - self.event_generator = self.survey_phasespaces(self.phasespaces, 1) - - def train_madnis(self) -> None: - madnis_args = self.run_card["madnis"] - if not any(subproc.madnis_settings["enable"] for subproc in self.subprocesses): - return - - gen_args = self.run_card["generation"] - run_args = self.run_card["run"] - - verbosity = resolve_verbosity(run_args["verbosity"]) - madnis_phasespaces = [] - training_args = [] - self.event_generator = None - for subproc, phasespace in zip(self.subprocesses, self.phasespaces): - for channel in phasespace.channels: - channel.event_generator = None - - config = ms.MadnisConfig() - config.learning_rate = subproc.madnis_settings["lr"] - config.batches = subproc.madnis_settings["train_batches"] - config.log_interval = madnis_args["log_interval"] - config.integration_history_length = madnis_args["integration_history_length"] - config.channel_dropping_interval = madnis_args["channel_dropping_interval"] - config.channel_dropping_threshold = madnis_args["channel_dropping_threshold"] - config.cpu_generator_batch_size = gen_args["cpu_batch_size"] - config.gpu_generator_batch_size = gen_args["gpu_batch_size"] - config.gpu_generator_batch_granularity = madnis_args["gpu_generator_batch_granularity"] - config.generator_target_size_factor = madnis_args["generator_target_size_factor"] - config.batch_size_offset = madnis_args["batch_size_offset"] - config.batch_size_per_channel = subproc.madnis_settings["batch_size_per_channel"] - config.uniform_channel_ratio = madnis_args["uniform_channel_ratio"] - config.lr_schedule = madnis_args["lr_scheduler"] - config.adam_beta1 = madnis_args["adam_beta1"] - config.adam_beta2 = madnis_args["adam_beta2"] - config.adam_eps = madnis_args["adam_eps"] - config.adam_weight_decay = madnis_args["adam_weight_decay"] - config.grad_clip_threshold = madnis_args["grad_clip_threshold"] - config.buffer_capacity = madnis_args["buffer_capacity"] - config.minimum_buffer_size = madnis_args["minimum_buffer_size"] - config.buffered_steps_fraction = madnis_args["buffered_steps_fraction"] - config.buffer_skip_batches = madnis_args["buffer_skip_batches"] - config.buffer_unweighting_quantile = madnis_args["buffer_unweighting_quantile"] - config.fixed_cwnet_fraction = subproc.madnis_settings["fixed_cwnet_fraction"] - config.softclip_threshold = madnis_args["softclip_threshold"] - config.compressed_channel_weight_count = madnis_args["compressed_channel_weight_count"] - phasespace = subproc.build_madnis(phasespace) - madnis_phasespaces.append(phasespace) - training_args.append( - ms.TrainingArgs( - config=config, - integrands=subproc.build_integrands( - phasespace, - madnis_training=True, - drop_cuts_and_rescale=madnis_args["drop_zero_integrands"] - ), - cwnet=phasespace.cwnet, - ) - ) - - gen_context = self.contexts[0] - opt_context = ms.Context( - device=self.devices[0], thread_count=self.pool_sizes[0] - ) - opt_context.copy_globals_from(gen_context) - - madnis_training = ms.MultiMadnisTraining( - generator_context=gen_context, - optimizer_context=opt_context, - training_args=training_args, - verbosity=verbosity, - status_file=self.status_file, - # Reuses the run's resolved seed (also used by build_event_generator()). - # Only the single-channel CPU sample-generation path is currently seeded - # -- buffered training and GPU multi-channel batches are still - # non-deterministic. - seed=self.run_seed, - ) - madnis_training.train() - for phasespace, active_channels in zip( - madnis_phasespaces, madnis_training.active_channels() - ): - phasespace.channels = [ - phasespace.channels[index] for index in active_channels - ] - del madnis_training - del opt_context - self.phasespaces = madnis_phasespaces - for context in self.contexts[1:]: - context.copy_globals_from(self.contexts[0]) - self.event_generator = self.build_event_generator(madnis_phasespaces) - self.madnis_trained = True - - def update_madnis_status_single( - self, batch: int, batch_target: int, loss: float, lr: float, channel_count: int - ) -> None: - now = time.time() - if batch + 1 < batch_target: - if now - self.last_update_time < 0.1: - return - self.last_update_time = now - progress_bar = ms.format_progress((batch + 1) / batch_target, 52) - time_diff = now - self.madnis_wall_time - time_str = f"{format_time(time_diff)}" - else: - progress_bar = "" - wall_diff = now - self.madnis_wall_time - cpu_diff = time.process_time() - self.madnis_cpu_time - time_str = ( - f"{format_time(wall_diff, centi=True)} wall, " - f"{format_time(cpu_diff, centi=True)} cpu" - ) - batch_str = f"{batch + 1} / {batch_target}" - self.madnis_box.set_column(1, [ - f"{batch_str:<15} {progress_bar}", - f"{loss:>.4f}", - f"{channel_count}", - time_str - ]) - self.madnis_box.print_update() - - def update_madnis_status_multi( - self, - subproc_id: int, - batch: int, - batch_target: int, - loss: float, - lr: float, - channel_count: int - ) -> None: - now = time.time() - subproc_count = len(self.subprocesses) - if batch + 1 < batch_target: - if now - self.last_update_time < 0.1: - return - self.last_update_time = now - progress_bar = ms.format_progress((batch + 1) / batch_target, 34) - progress_bar_all = ms.format_progress( - (subproc_id * batch_target + batch + 1) / (subproc_count * batch_target), - 52 - ) - time_str = f"{format_time(now - self.madnis_wall_time)}" - subproc_str = f"{subproc_id} / {subproc_count}" - elif subproc_id < subproc_count - 1: - progress_bar = "" - progress_bar_all = ms.format_progress( - ((subproc_id + 1) * batch_target + 1) / (subproc_count * batch_target), - 52 - ) - time_str = f"{format_time(now - self.madnis_wall_time)}" - subproc_str = f"{subproc_id} / {subproc_count}" - else: - progress_bar = "" - progress_bar_all = "" - wall_diff = now - self.madnis_wall_time - cpu_diff = time.process_time() - self.madnis_cpu_time - time_str = ( - f"{format_time(wall_diff, centi=True)} wall, " - f"{format_time(cpu_diff, centi=True)} cpu" - ) - subproc_str = f"{subproc_count} / {subproc_count}" - batch_str = f"{batch + 1} / {batch_target}" - self.madnis_upper_box.set_column(1, [ - f"{subproc_str:<15} {progress_bar_all}", - time_str, - ]) - self.madnis_lower_box.set_row(subproc_id + 1, [ - f"{subproc_id}", - f"{loss:>.4f}", - f"{channel_count}", - f"{batch_str:<15} {progress_bar}", - ]) - self.madnis_upper_box.print_update() - self.madnis_lower_box.print_update() - - @property - def gridpack_run_mode(self) -> str: - """How this run ends: a gridpack run may skip the events (see the - [gridpack] run_mode run card option); any other run is a regular one.""" - if not self.run_card["gridpack"]["save_gridpack"]: - return "regular" - return self.run_card["gridpack"]["run_mode"] - - def finish_without_events(self) -> None: - """The end of a gridpack run that writes no events: only what the - gridpack needs is prepared, then it is stored.""" - if self.gridpack_run_mode == "minimal": - # a trained MadNIS needs nothing more, a VEGAS grid has to converge - if not self.madnis_trained: - self.event_generator.optimize() - else: - self.event_generator.fix_max_weights() - self.build_systematics() - self.save_gridpack() - - def generate_events(self) -> None: - if self.gridpack_run_mode != "regular": - self.finish_without_events() - return - start_time = get_start_time() - self.event_generator.generate() - output_format = self.run_card["run"]["output_format"] - systematics = self.build_systematics() - histograms = self.build_event_histograms() - if output_format == "compact_npy": - self.lhe_completer = None - self.event_generator.combine_to_compact_npy( - os.path.join(self.run_path, "events.npy"), systematics, histograms - ) - self.write_lhe_header(systematics) - self.save_lhe_completer() - elif output_format == "lhe_npy": - self.lhe_completer = self.build_lhe_completer() - self.event_generator.combine_to_lhe_npy( - os.path.join(self.run_path, "events.npy"), self.lhe_completer, - systematics, histograms - ) - self.write_lhe_header(systematics) - elif output_format == "lhe": - self.lhe_completer = self.build_lhe_completer() - lhe_path = os.path.join(self.run_path, "events.lhe") - self.event_generator.combine_to_lhe( - lhe_path, self.lhe_completer, - self.build_lhe_meta(), systematics, histograms - ) - # Ship the LHE compressed by default. These files are large and - # very compressible, madevent has always stored its events - # gzipped, and every consumer here already accepts either form - # (see _find_event_file). misc.gzip replaces events.lhe with - # events.lhe.gz, and switches to an external multithreaded tool - # above 256 MB. - misc.gzip(lhe_path) - else: - raise ValueError("Unknown output format") - self.write_hwu(histograms) - self.make_plots(histograms) - if systematics is not None: - self.write_systematics_sidecar() - self.log_systematics_summary() - self.save_gridpack() - - def save_lhe_completer(self) -> None: - """Keep the LHE completer next to compact_npy events, so npy_to_lhe - can complete them into an LHE file later (a shower or MadSpin requested - after the run). Never fatal to the run.""" - try: - self.build_lhe_completer().save( - os.path.join(self.run_path, npy_to_lhe.COMPLETER_FILE)) - except Exception as err: - logger.warning("could not save the LHE completer; these events " - "cannot be converted to LHE later: %s", err) - - hwu_file_name = "MADatLO.HwU" - plot_dir_name = "plots" - - def make_plots(self, histograms) -> None: - """Draw the event-sample histograms into Events//plots, with the - scale and PDF bands info.json carries, and say where they landed. - - This does not go through the HwU file: the bands are already computed - per bin by madspace, so drawing them directly is both fewer steps and - more than histograms.py can do without the LHAPDF python module. - """ - if histograms is None or not self.run_card["run"]["make_plots"]: - return - out_dir = os.path.join(self.run_path, self.plot_dir_name) - try: - data = json.loads(histograms.to_json(self.systematics)) - written = plots.render(data, out_dir) - except plots.BackendMissing as error: - logger.info("matplotlib is not available (%s): no plots. The " - "distributions are in %s", - error, os.path.join(self.run_path, "info.json")) - return - except Exception as error: - logger.warning("could not draw the histograms: %s", error) - return - if not written: - return - logger.info("%d plot(s) stored in %s", len(written), out_dir) - - def write_hwu(self, histograms) -> None: - """Write the event-sample histograms next to the events in the HwU - format (MADatLO.HwU), the one an aMC@NLO run writes as MADatNLO.HwU - and madgraph/various/histograms.py reads: an mg7 distribution can then - be overlaid on an NLO one without converting anything in between. The - numbers are the ones info.json carries, bands included, so this is off - by default -- it is a second copy in another format, for the runs that - are going to be plotted.""" - if histograms is None or not self.run_card["run"]["write_hwu"]: - return - try: - data = json.loads(histograms.to_json(self.systematics)) - summary = (json.loads(self.systematics.summary()) - if self.systematics is not None else None) - text = hwu_output.to_hwu(data, summary) - except Exception as error: - logger.warning("could not write the HwU histograms: %s", error) - return - if not text: - return - with open(os.path.join(self.run_path, self.hwu_file_name), "w") as f: - f.write(text) - - def write_systematics_sidecar(self) -> None: - """Describe the variation weights next to the event file - (events.weights.json): ids, scale factors, PDF set/member and the - text, so npy consumers get the same metadata as the LHE - header, plus the per-variation cross sections.""" - if self.systematics is None: - return - data = json.loads(self.systematics.summary()) - data["initrwgt"] = self.systematics.initrwgt() - data["columns"] = ["rwgt_%d" % i for i in self.systematics.weight_ids] - with open(os.path.join(self.run_path, "events.weights.json"), "w") as f: - json.dump(data, f, indent=1) - - @staticmethod - def _histogram_mean(hist_min, hist_max, bin_values): - """Cross-section-weighted mean of a histogrammed observable over its - range. ``bin_values`` carries the underflow and overflow bins first and - last, as both madspace histogram kinds do; they have no position and - are left out.""" - values = list(bin_values)[1:-1] - total = sum(values) - if not values or total == 0: - return None - width = (hist_max - hist_min) / len(values) - return sum(v * (hist_min + (i + 0.5) * width) - for i, v in enumerate(values)) / total - - def _histogram_means(self) -> dict: - """name -> mean of every [histograms] entry, from the post-processing - histograms (the final event sample) when they were filled, otherwise - from the weighted integration histograms.""" - if self.event_histograms is not None: - data = json.loads(self.event_histograms.to_json(self.systematics)) - hists = [(h["name"], h["min"], h["max"], h["bin_values"]) - for h in data] - else: - hists = [(h.name, h.min, h.max, h.bin_values) - for h in self.event_generator.histograms()] - means = {} - for name, hist_min, hist_max, values in hists: - mean = self._histogram_mean(hist_min, hist_max, values) - if mean is not None: - means[name] = mean - return means - - def get_result(self) -> dict: - """Return the run result: cross-section (pb) with MC error, the number - of (unweighted) events, and the mean of every observable declared in the - [histograms] section. Used to build the scan summary.""" - status = self.event_generator.status() - result = {'cross(pb)': status.mean, 'error(pb)': status.error, - 'nb_event': status.count_unweighted} - try: - if self.systematics is not None and self.systematics.weight_count: - summary = json.loads(self.systematics.summary()) - scale = systematics_summary.scale_percentages(summary) - if scale is not None: - result['scale_up(%)'], result['scale_down(%)'] = scale - pdf = systematics_summary.pdf_percentages(summary) - if pdf: - # one column pair; the nominal set's entry comes first - _entry, up, down = pdf[0] - result['pdf_up(%)'], result['pdf_down(%)'] = up, down - except Exception as err: - logger.warning("could not extract the systematics summary: %s", err) - try: - for name, mean in self._histogram_means().items(): - result['<%s>' % name] = mean - except Exception as err: - logger.warning("could not extract observable means: %s", err) - return result - - def _beam_info(self): - """Return (beam_pdg_ids, beam_energies) for the LHE block. - - `incoming` holds the *partonic* initial state (e.g. gluons), not the - beam particle, so hadronic beams are protons (2212); leptonic beams are - the incoming leptons themselves.""" - half_e = float(self.e_cm) / 2. - if self.is_decay: - # No beams. LHE has no way to say that, so report the decaying - # particle at rest as a single "beam"; the second slot is empty. - data = self.subprocess_data[0] - pdg = clean_pids(data["incoming"])[0] - return [pdg, 0], [float(self.e_cm), 0.0] - if not self.leptonic: - # hadronic collider: proton beams (p-pbar is not distinguished) - return [2212, 2212], [half_e, half_e] - data = self.subprocess_data[0] - incoming = data["incoming"] - flavor0 = data["flavors"][0]["options"][0] - init_pdgs = flavor0[:len(incoming)] - beam_pdgs = [pdg if abs(code) in (81, 82) else code - for code, pdg in zip(incoming, init_pdgs)] - return beam_pdgs, [half_e, half_e] - - def _lhapdf_id(self): - """Central LHAPDF id of the beam PDF set (read from its .info SetIndex), - or -1 for a leptonic beam (no PDF).""" - if self.leptonic or not self.pdf_dir: - return -1 - pdf_set = self.run_card["beam"]["pdf"] - info = os.path.join(self.pdf_dir, pdf_set, "%s.info" % pdf_set) - try: - for line in open(info): - if line.strip().startswith("SetIndex:"): - return int(line.split(":", 1)[1].strip()) - except Exception as err: - logger.warning("could not read LHAPDF id from %s: %s", info, err) - return -1 - - def build_lhe_meta(self, systematics=None): - """Build the LHE header/ metadata: the param_card () and the - run_card.toml () headers plus the beam/PDF/cross-section info - needed by downstream tools (systematics, MadSpin, ...). - - `combine_to_lhe` injects the header itself, so callers that - go through it should leave `systematics` unset; pass it only when - building meta for a writer that bypasses that injection (e.g. - write_lhe_header).""" - beam_pdgs, energies = self._beam_info() - lhaid = self._lhapdf_id() - pdf_group = -1 if self.leptonic else 0 - status = self.event_generator.status() - xsec, err = status.mean, status.error - weight_mode, max_weight = lhe_weight_info(status) - with open(self.param_card_path) as f: - param_text = f.read() - with open(os.path.join("Cards", "run_card.toml")) as f: - run_text = f.read() - headers = [] - # generation commands (model + process): read by MadSpin/reweight/... - proc_card = os.path.join("Cards", "proc_card_mg5.dat") - if os.path.exists(proc_card): - with open(proc_card) as f: - headers.append(ms.LHEHeader(name="MG5ProcCard", content=f.read())) - headers.append(ms.LHEHeader(name="slha", content=param_text)) - headers.append(ms.LHEHeader(name="MG7RunCard", content=run_text)) - # The resolved seed (even when the run_card requested a random one via - # seed = -1), so the run can be reproduced from the LHE file alone. - headers.append(ms.LHEHeader(name="MG7Seed", content=str(self.run_seed))) - if systematics is not None and systematics.weight_ids: - headers.append(ms.LHEHeader( - name="initrwgt", content=systematics.initrwgt(), escape_content=False - )) - return ms.LHEMeta( - beam1_pdg_id=beam_pdgs[0], beam2_pdg_id=beam_pdgs[1], - beam1_energy=energies[0], beam2_energy=energies[1], - beam1_pdf_authors=pdf_group, beam2_pdf_authors=pdf_group, - beam1_pdf_id=lhaid, beam2_pdf_id=lhaid, - weight_mode=weight_mode, - # positional: the pybind arg name for max_weight is non-kwarg-safe - processes=[ms.LHEProcess(xsec, err, max_weight, 1)], - headers=headers, - ) - - def write_lhe_header(self, systematics) -> None: - """Write header.lhe next to events.npy: the
/ blocks - (run card, param card, beam/PDF, cross section) that the npy formats - otherwise drop, with no events.""" - header_path = os.path.join(self.run_path, "header.lhe") - writer = ms.LHEFileWriter(header_path, self.build_lhe_meta(systematics)) - del writer # closes the file (writes the closing tag) - - def build_lhe_completer(self): - all_mcdata = ( - [subproc.build_multi_channel_data() for subproc in self.subprocesses] - if self.merged_subprocess_data is None else - [build_multi_channel_data(meta, self) for meta in self.subprocess_data] - ) - subproc_args = [ - ms.SubprocArgs( - topologies = [topo[0] for topo in mcdata.topologies], - permutations = mcdata.permutations, - diagram_indices = mcdata.diagram_indices, - diagram_color_indices = mcdata.diagram_color_indices, - diagram_propagator_pdgs = mcdata.diagram_propagator_pdgs, - color_flows = meta["color_flows"], - pdg_color_types = { - int(key): value - for key, value in meta["pdg_color_types"].items() - }, - helicities = lhe_helicities(meta, self.run_card), - pdg_ids = [flavor["options"] for flavor in meta["flavors"]], - ) - for mcdata, meta in zip(all_mcdata, self.subprocess_data) - ] - return ms.LHECompleter( - subproc_args=subproc_args, - bw_cutoff=self.run_card["phasespace"]["bw_cutoff"] - ) - - def save_gridpack(self) -> None: - if not self.run_card["gridpack"]["save_gridpack"]: - return - - gridpack_path = os.path.join(self.run_path, "gridpack") - data_path = os.path.join(gridpack_path, "data") - events_path = os.path.join(gridpack_path, "Events") - os.mkdir(gridpack_path) - os.mkdir(data_path) - os.mkdir(events_path) - self.contexts[0].save_globals(os.path.join(data_path, "globals.tar")) - - # one json object per channel, joined by hand so that the channel - # json is not parsed and dumped again - channel_entries = [ - f"{json.dumps(channel.status().name)}:{channel.to_json(self.gridpack_run_mode == "fix_max_weight")}" - for channel in self.event_generator.channels() - ] - with open(os.path.join(data_path, "channels.json"), "w") as f: - f.write("{" + ",".join(channel_entries) + "}") - - lib_path = os.path.join(gridpack_path, "lib") - if self.run_card["gridpack"]["include_source"]: - os.mkdir(lib_path) - shutil.copytree("src", os.path.join(gridpack_path, "src")) - shutil.copytree("SubProcesses", os.path.join(gridpack_path, "SubProcesses")) - shutil.copytree("backend", os.path.join(gridpack_path, "backend")) - else: - shutil.copytree("lib", lib_path) - for _, path in self.custom_function_paths: - target = path if os.path.isabs(path) else os.path.join(gridpack_path, path) - if not os.path.exists(target): - os.makedirs(os.path.dirname(target), exist_ok=True) - shutil.copy(path, target) - - if self.run_card["gridpack"]["include_madspace_source"]: - shutil.copytree( - _MADSPACE_DIR, - os.path.join(gridpack_path, "madspace"), - ignore=shutil.ignore_patterns("build", "install"), - ) - - if self.run_card["gridpack"]["include_madspace"]: - shutil.copytree( - _INSTALL_DIR / "madspace", - os.path.join(gridpack_path, "madspace", "install", "madspace"), - ignore=shutil.ignore_patterns("include"), - ) - - matrix_elements = [] - for subproc in self.subprocess_data: - me_path = subproc["me_path"] - matrix_elements.append(me_path) - - cards_path = os.path.join(gridpack_path, "Cards") - os.mkdir(cards_path) - shutil.copy(os.path.join("Cards", "param_card.dat"), cards_path) - # Full run card with a header noting it is read-only in gridpack context. - import io as _io - _buf = _io.StringIO() - self.run_card.write(_buf) - _header = ( - "# This is the run card used to generate this gridpack.\n" - "# Modifying this file will have no effect on gridpack execution.\n" - "# To change event-generation settings, edit grid_run_card.toml.\n\n" - ) - with open(os.path.join(cards_path, "run_card.toml"), 'w') as _f: - _f.write(_header + _buf.getvalue()) - # Minimal card containing only the settings used by generate_events. - # The gridpack ships the libraries that were actually built, and their - # names carry the resolved backend, so the card it runs from has to name - # that backend as well: a cpu_mode of 'auto' would send the gridpack - # looking for a ..._auto.so that was never built. The full run_card.toml - # written just above keeps 'auto', since that is what was asked for. - device_names = self.run_card["run"]["device"] - if not isinstance(device_names, list): - device_names = [device_names] - resolved_cpu_mode = None - for name, backend in zip(device_names, getattr(self, "backends", [])): - if name.split(":")[0] == "cpu": - resolved_cpu_mode = backend - break - previous_cpu_mode = self.run_card["run"]["cpu_mode"] - if resolved_cpu_mode is not None: - self.run_card["run"]["cpu_mode"] = resolved_cpu_mode - # A gridpack is run to feed a shower or detector chain, which reads - # LHE, so it writes LHE by default whatever format this run used - # (bin/generate_events --output_format still selects an npy output). - previous_output_format = self.run_card["run"]["output_format"] - self.run_card["run"]["output_format"] = "lhe" - # the directories of this run mean nothing on the machine running the gridpack - previous_dirs = (self.run_card["gridpack"]["temp_output_dir"], - self.run_card["gridpack"]["output_dir"]) - self.run_card["gridpack"]["temp_output_dir"] = "" - self.run_card["gridpack"]["output_dir"] = "" - try: - self.run_card.write_gridpack_card( - os.path.join(cards_path, "grid_run_card.toml")) - finally: - self.run_card["run"]["cpu_mode"] = previous_cpu_mode - self.run_card["run"]["output_format"] = previous_output_format - (self.run_card["gridpack"]["temp_output_dir"], - self.run_card["gridpack"]["output_dir"]) = previous_dirs - - bin_path = os.path.join(gridpack_path, "bin") - os.mkdir(bin_path) - gen_events_file = os.path.join(bin_path, "generate_events") - template_dir = os.path.dirname(__file__) - shutil.copy(os.path.join(template_dir, "gridpack.py"), gen_events_file) - shutil.copy(os.path.join(template_dir, "custom_function_loader.py"), bin_path) - os.chmod(gen_events_file, 0o755) - shutil.copy(os.path.join(template_dir, "gridpack_setup.py"), bin_path) - # streaming needs the fixed maximum weights of this mode - if self.gridpack_run_mode == "fix_max_weight": - shutil.copy(os.path.join(template_dir, "event_stream.py"), bin_path) - - data = { - "matrix_elements": matrix_elements, - # loaded before the matrix elements, as in the run - "custom_functions": self.custom_function_paths, - "me_parameters": me_parameters(self.run_card), - "source_hash": ms.SOURCE_HASH, - } - with open(os.path.join(data_path, "data.json"), "w") as f: - json.dump(data, f) - - if self.lhe_completer is None: - self.lhe_completer = self.build_lhe_completer() - self.lhe_completer.save(os.path.join(data_path, "lhe.json")) - self.save_gridpack_lhe_meta(os.path.join(data_path, "lhe_meta.json")) - if self.systematics_data is not None: - systematics_data = dict(self.systematics_data) - # A decay takes alpha_s from a constant grid written into the run - # directory rather than from an LHAPDF set, so the gridpack cannot - # look that file up: ship it and point at the shipped copy. - if self.fixed_alphas_info and \ - systematics_data["nominal_info_file"] == self.fixed_alphas_info: - name = os.path.basename(self.fixed_alphas_info) - shutil.copy(self.fixed_alphas_info, os.path.join(data_path, name)) - systematics_data["nominal_info_file"] = os.path.join("data", name) - with open(os.path.join(data_path, "systematics.json"), "w") as f: - json.dump(systematics_data, f) - if self.run_card["gridpack"]["compress"]: - shutil.make_archive(gridpack_path, "gztar", self.run_path, "gridpack") - shutil.rmtree(gridpack_path) - - def save_gridpack_lhe_meta(self, path) -> None: - """The run-independent part of the LHE header (cards, beams, PDF), for - the gridpack to write complete LHE headers: it adds the cross section - and seed of its own run, which is why those two are left out here.""" - meta = self.build_lhe_meta() - data = {key: getattr(meta, key) for key in ( - "beam1_pdg_id", "beam2_pdg_id", "beam1_energy", "beam2_energy", - "beam1_pdf_authors", "beam2_pdf_authors", "beam1_pdf_id", - "beam2_pdf_id", "weight_mode")} - data["headers"] = [ - {"name": h.name, "content": h.content, - "escape_content": h.escape_content} - for h in meta.headers if h.name != "MG7Seed"] - with open(path, "w") as f: - json.dump(data, f) - - def get_mass(self, pid: int) -> float: - return self.param_card.get_value("mass", pid) - - def get_width(self, pid: int) -> float: - return self.param_card.get_value("width", pid) - - -# Flavor-merged legs carry a group id instead of a pdg. Every member of a group -# shares the same mass -- that is what makes them mergeable -- so any member is -# a valid representative for the kinematics. -_MERGED_PID_REPRESENTATIVE = { - 81: 1, # light quarks d u s c - 82: 11, # charged leptons e mu - 83: 12, # neutrinos ve vm vt -} - - -def clean_pids(pids: list[int]) -> list[int]: - pids_out = [] - for pid in pids: - pid = abs(pid) - if pid in _MERGED_PID_REPRESENTATIVE: - pid = _MERGED_PID_REPRESENTATIVE[pid] - elif 81 <= pid <= 99: - # Reserved for flavor merging. Anything outside this window is a - # real pdg (BSM models use codes in the millions), so let it pass. - raise ValueError( - f"unknown flavor-merged particle id {pid}; add its " - "representative to _MERGED_PID_REPRESENTATIVE" - ) - pids_out.append(pid) - return pids_out - - -# observables that tell a particle from its antiparticle, and so need the signed -# pdg ids of the actual flavours rather than clean_pids' unsigned representatives -_SIGNED_OBSERVABLES = {"sfos_pair_mass"} - - -def _flavor_options(meta: dict, unmerged_meta: dict | None = None) -> list[list[int]]: - """Every flavour assignment (signed pdg ids, incoming then outgoing) a - subprocess evaluates with the same momenta.""" - if "flavors" in meta: - return [ - option for flavor in meta["flavors"] for option in flavor["options"] - ] - return [ - option - for subproc in meta.get("subprocesses", []) - for option in _flavor_options(unmerged_meta[subproc]) - ] - - -def observable_pids( - meta: dict, observable_kwargs: dict, unmerged_meta: dict | None = None -) -> list[int]: - """The pdg ids an observable of this subprocess is built with. - - Most observables only select particles by group, which the unsigned - representatives of clean_pids do. One that pairs a particle with its - antiparticle needs the real signed flavours: those of the first flavour - assignment, provided every assignment of the subprocess forms the same - pairs -- the cut is applied once to momenta that all of them share. - """ - all_pids = clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) - if observable_kwargs.get("observable") not in _SIGNED_OBSERVABLES: - return all_pids - options = _flavor_options(meta, unmerged_meta) - if not options: - raise ValueError( - f"{observable_kwargs.get('name')}: no flavour information for " - f"subprocess {meta['incoming']} > {meta['outgoing']}" - ) - n_in = len(meta["incoming"]) - groups = [set(group) for group in observable_kwargs["select_pids"]] - if len(groups) == 1: - groups = groups * 2 - - def sfos_pairs(pids): - out = range(n_in, len(pids)) - return frozenset( - frozenset((i, j)) - for i in out if pids[i] in groups[0] - for j in out if pids[j] in groups[1] - if i != j and pids[i] != 0 and pids[i] == -pids[j] - ) - - reference = sfos_pairs(options[0]) - if any(sfos_pairs(option) != reference for option in options[1:]): - raise ValueError( - f"{observable_kwargs.get('name')}: the flavours grouped into subprocess " - f"{meta['incoming']} > {meta['outgoing']} pair the particles " - "differently, so a same-flavour opposite-sign cut cannot be applied " - "to their common momenta; regenerate the process after " - "'set apply_flavor_grouping False'" - ) - return list(options[0]) - - -def build_observable(pids: list[int], n_in: int, observable_kwargs: dict): - """ms.Observable of a subprocess with `n_in` incoming particles. - - madspace takes the first two momenta to be the beams (ignore_incoming), so - for a 1 -> N decay it would drop the first decay product from every - selection. There the selection runs over all the momenta instead, with the - incoming particle set to pdg id 0, which is in no group. - """ - if n_in == 2: - return ms.Observable(pids, **observable_kwargs) - pids = [0] * n_in + list(pids[n_in:]) - return ms.Observable(pids, **dict(observable_kwargs, ignore_incoming=False)) - - -def decay_products(meta: dict, unmerged_meta: dict | None = None) -> set[int]: - """Outgoing positions (0-based, beams excluded) of the particles that come - from an on-shell (decay-chain) propagator, as MadEvent's check_decay marks - them. Every channel of a subprocess shares its decay chains, so the first - channel tells, read through its first diagram's permutation.""" - if not meta["channels"]: - return set() - channel = meta["channels"][0] - topo_channel = channel - if "vertices" not in channel: - topo_channel = unmerged_meta[channel["subprocess"]]["channels"][channel["channel"]] - on_shell = topo_channel.get("on_shell_propagators", []) - if not on_shell: - return set() - vertices = topo_channel["vertices"] - # the closing vertex creates no propagator - created = {vertex[-1]: vertex[:-1] for vertex in vertices[:-1]} - - def legs(name): - if name in created: - return set().union(*(legs(child) for child in created[name])) - return {name} - - n_in = len(meta["incoming"]) - n_out = len(meta["outgoing"]) - products = set() - for prop in on_shell: - external = legs(f"p{prop}") - incoming = {leg for leg in external if leg.startswith("i")} - outgoing = {int(leg[1:]) for leg in external if leg.startswith("o")} - if not incoming: - products |= outgoing - elif len(incoming) == n_in: - # built from the beams: the decay is everything else - products |= set(range(n_out)) - outgoing - # event[i] = topology[permutation[i]] - permutation = channel["diagrams"][0]["permutation"] - return { - i - n_in for i in range(n_in, len(permutation)) - if permutation[i] - n_in in products - } - - -def pid_is_qcd(pid: int): - return abs(pid) in [21, 1, 2, 3, 4, 5, 6, 81] - - -def build_topologies( - incoming_masses: list[float], - outgoing_masses: list[float], - channel: dict, - process: MadgraphProcess -) -> list[ms.Topology]: - propagators = [] - for i, (pid, signed_pid) in enumerate(zip( - clean_pids(channel["propagators"]), channel["propagators"] - )): - mass = process.get_mass(pid) - width = process.get_width(pid) - bw_cutoff = process.run_card["phasespace"]["bw_cutoff"] - if i in channel["on_shell_propagators"]: - e_min = mass - bw_cutoff * width - e_max = mass + bw_cutoff * width - else: - e_min = 0 - e_max = 0 - # $-excluded propagator: the matrix element vanishes within bw_cutoff - # widths of the pole, so do not pile the channel's points up there - # (a flat density in that window instead of the Breit-Wigner peak). - # Channels written before the exporter recorded them have none. - flat_window = 0. - if i in channel.get("dollar_propagators", ()) and width > 0: - flat_window = float(bw_cutoff) - propagators.append(ms.Propagator( - mass=mass, - width=width, - integration_order=0, - e_min=e_min, - e_max=e_max, - pdg_id=signed_pid, - flat_window=flat_window, - )) - vertices = channel["vertices"] - diag = ms.Diagram( - incoming_masses, outgoing_masses, propagators, vertices - ) - return ms.Topology.topologies(diag) - - -def build_multi_channel_data( - meta: dict, process: MadgraphProcess, unmerged_meta: dict | None = None -) -> MultiChannelData: - incoming_masses = [ - process.get_mass(pid) for pid in clean_pids(meta["incoming"]) - ] - outgoing_masses = [ - process.get_mass(pid) for pid in clean_pids(meta["outgoing"]) - ] - - if unmerged_meta is None: - diagram_count = meta["diagram_count"] - amp2_remaps = [[-1] * diagram_count] - else: - amp2_remaps = [ - [-1] * unmerged_meta[subproc]["diagram_count"] - for subproc in meta["subprocesses"] - ] - symfact = [] - topologies = [] - permutations = [] - channel_indices = [] - channel_weight_indices = [] - diagram_indices = [] - diagram_color_indices = [] - diagram_propagator_pdgs = [] - active_flavors = [] - channel_index = 0 - qcd_s_channel_count = [] - - for channel in meta["channels"]: - if unmerged_meta is None: - topo_channel = channel - else: - topo_subproc = channel["subprocess"] - topo_channel_index = channel["channel"] - topo_channel = unmerged_meta[topo_subproc]["channels"][topo_channel_index] - chan_topologies = build_topologies( - incoming_masses, outgoing_masses, topo_channel, process - ) - topo_count = len(chan_topologies) - if topo_count == 0: - continue - - topo = chan_topologies[0] - s_chan_count = [0] * len(topo.decays) - non_qcd = [False] * len(topo.decays) - pdg_ids = [decay.pdg_id for decay in topo.decays] - for i, pid in zip(topo.outgoing_indices, meta["outgoing"]): - pdg_ids[i] = pid - for decay in reversed(topo.decays): - if len(decay.child_indices) == 0: - continue - if decay.index == 0 and topo.t_propagator_count > 0: - s_chan_count[0] = sum(s_chan_count[i] for i in decay.child_indices) - break - - is_qcd = pid_is_qcd(pdg_ids[decay.index]) and all( - pid_is_qcd(pdg_ids[i]) for i in decay.child_indices - ) - is_non_qcd = decay.on_shell or not is_qcd or any( - non_qcd[i] for i in decay.child_indices - ) - non_qcd[decay.index] = is_non_qcd - s_chan_count[decay.index] = ( - 0 if is_non_qcd else sum(s_chan_count[i] for i in decay.child_indices) + 1 - ) - qcd_s_channel_count.append(s_chan_count[0]) - - diagrams = channel["diagrams"] - chan_permutations = [d["permutation"] for d in diagrams] - if unmerged_meta is None: - amp2_remaps[0][diagrams[0]["diagram"]] = channel_index - else: - for amp2_remap, diag in zip(amp2_remaps, diagrams[0]["diagram"]): - if diag != -1: - amp2_remap[diag] = channel_index - - channel_index_first = channel_index - symfact_index_first = len(symfact) - channel_index += 1 - symfact.extend([None] * topo_count) - for d in diagrams[1:]: - if unmerged_meta is None: - amp2_remaps[0][d["diagram"]] = channel_index - else: - for amp2_remap, diag in zip(amp2_remaps, d["diagram"]): - if diag != -1: - amp2_remap[diag] = channel_index - channel_index += 1 - symfact.extend(range(symfact_index_first, symfact_index_first + topo_count)) - - topologies.append(chan_topologies) - permutations.append(chan_permutations) - channel_indices.append(list(range(channel_index_first, channel_index))) - channel_weight_indices.append([ - [ - symfact_index_first + topo_index + i * topo_count - for i in range(len(chan_permutations)) - ] - for topo_index in range(topo_count) - ]) - diagram_indices.append([d["diagram"] for d in diagrams]) - if unmerged_meta is None: - diagram_color_indices.append([d["active_colors"] for d in diagrams]) - diagram_propagator_pdgs.append( - [d["propagator_pdgs"] for d in diagrams] - ) - active_flavors.append([d["active_flavors"] for d in diagrams]) - - return MultiChannelData( - amp2_remaps, - symfact, - topologies, - permutations, - channel_indices, - channel_weight_indices, - diagram_indices, - diagram_color_indices, - diagram_propagator_pdgs, - active_flavors, - qcd_s_channel_count, - ) - - -class MadgraphSubprocess: - def __init__( - self, - process: MadgraphProcess, - meta: dict, - subproc_id: int, - unmerged_meta: dict | None = None - ): - self.process = process - self.meta = meta - self.subproc_id = subproc_id - self.multi_channel_data = None - - self.unmerged_meta = None - if unmerged_meta is None: - api_path_formats = [self.meta["me_path"]] - member_metas = [self.meta] - else: - api_path_formats = [] - member_metas = [] - for subproc in self.meta["subprocesses"]: - submeta = unmerged_meta[subproc] - api_path_formats.append(submeta["me_path"]) - member_metas.append(submeta) - if len(api_path_formats) == 1: - self.meta = submeta - else: - self.unmerged_meta = unmerged_meta - - # The libraries were all built up front by MadgraphProcess.compile_matrix_elements - all_api_paths = [ - [api_path_format.format(device=backend) for backend in self.process.backends] - for api_path_format in api_path_formats - ] - - self.incoming_masses = [ - self.process.get_mass(pid) for pid in clean_pids(self.meta["incoming"]) - ] - self.outgoing_masses = [ - self.process.get_mass(pid) for pid in clean_pids(self.meta["outgoing"]) - ] - self.particle_count = len(self.incoming_masses) + len(self.outgoing_masses) - # MadEvent's cut_decays = F: the decay products of on-shell propagators - # are left out of every cut selection (a pair cut then needs both - # members to be cut), while the histograms still see them - if self.process.run_card["phasespace"]["cut_decays"]: - uncut = set() - else: - uncut = decay_products(self.meta, self.unmerged_meta) - n_in = len(self.incoming_masses) - - def cut_pids(observable_kwargs): - pids = observable_pids(self.meta, observable_kwargs, self.unmerged_meta) - # pdg id 0 is in no group, so the particle is never selected - return [0 if i - n_in in uncut else pid for i, pid in enumerate(pids)] - - # custom cuts and scales see the PDG ids of the first flavor, like the - # built-in cuts, which select by flavor-independent groups - representative_pids = [self.flavor_pids()[0]] - custom_cut = self.process.custom_function( - "cut", self.particle_count, representative_pids - ) - if len(self.process.cut_data) > 0: - self.cuts = ms.Cuts([ - ms.CutItem( - observable=build_observable( - cut_pids(cut_item.observable_kwargs), - n_in, cut_item.observable_kwargs - ), - min=cut_item.min, - max=cut_item.max, - mode=cut_item.mode, - ) - for cut_item in self.process.cut_data - ], custom_cut) - elif custom_cut is not None: - self.cuts = ms.Cuts(self.particle_count, custom_cut) - else: - self.cuts = None - # the integration histograms are functions of the momenta; the weight - # distribution is a property of the final event sample, so it is only - # filled by MadgraphProcess.build_event_histograms at combine time - # They cost an observable evaluation per integration point, and the - # plots/HwU are drawn from the post-processing ones instead, so they - # are only filled on request ([run] weighted_histograms). - momentum_hists = [item for item in self.process.hist_data - if not item.from_weight] \ - if self.process.run_card["run"]["weighted_histograms"] else [] - self.histograms = ( - ms.ObservableHistograms([ - ms.HistItem( - observable=build_observable( - observable_pids( - self.meta, hist_item.observable_kwargs, self.unmerged_meta - ), - n_in, hist_item.observable_kwargs - ), - min=hist_item.min, - max=hist_item.max, - bin_count=hist_item.bin_count, - ) - for hist_item in momentum_hists - ]) - if momentum_hists - else None - ) - - self.scale = ms.EnergyScale( - particle_count=self.particle_count, - custom_scale=self.process.custom_function( - "scale", self.particle_count, representative_pids - ), - **self.process.scale_kwargs - ) - - if self.process.run_card["run"]["dummy_matrix_element"]: - self.matrix_elements = [None] * len(all_api_paths) - else: - self.matrix_elements = [] - for api_paths in all_api_paths: - for context, api_path in zip(self.process.contexts, api_paths): - mat = context.load_matrix_element( - api_path, self.process.param_card_path, - me_parameters(self.process.run_card) - ) - self.matrix_elements.append(mat) - self.surrogates = [] - if self.process.surrogate is not None: - for path_format, member_meta in zip(api_path_formats, member_metas): - self.surrogates.append(self.load_surrogate(path_format, member_meta)) - - def load_surrogate(self, me_path_format: str, member_meta: dict): - """The [surrogate] of one matrix element of the subprocess.""" - name = re.sub( - r"^libmadmatrix_(.*)_\{device\}\.so$", r"\1", os.path.basename(me_path_format) - ) - # PDG ids for each flavor index of the matrix element - flavors = member_meta["flavors"] - pdg_ids = [flavors[0]["options"][0]] * (max(f["index"] for f in flavors) + 1) - for flavor in flavors: - pdg_ids[flavor["index"]] = flavor["options"][0] - settings = self.process.surrogate - apis = load_surrogate( - ms, - self.process.contexts, - self.process.backends, - self.process.device_types, - settings["function"], - name, - pdg_ids, - settings["param_card"], - ) - if len({api.index() for api in apis}) != 1: - raise RuntimeError("surrogate loaded at different indices") - return apis[0] - - def flavor_pids(self) -> list[list[int]]: - """PDG ids of the external particles for every flavor of the - integrand, in the order of its flavor index.""" - pids = [] - for flav in self.meta["flavors"]: - if self.unmerged_meta is not None: - subproc_index = self.meta["subprocesses"][flav["subprocess"]] - flav = self.unmerged_meta[subproc_index]["flavors"][flav["flavor"]] - pids.append(flav["options"][0]) - return pids - - def build_multi_channel_data(self) -> MultiChannelData: - if self.multi_channel_data is not None: - return self.multi_channel_data - self.multi_channel_data = build_multi_channel_data( - self.meta, self.process, self.unmerged_meta - ) - return self.multi_channel_data - - def build_multichannel_phasespace(self) -> PhaseSpace: - mcdata = self.build_multi_channel_data() - channel_count = sum(len(topos) for topos in mcdata.topologies) - drop_threshold = self.process.run_card["phasespace"]["drop_qcd_s_channel"] - if drop_threshold >= 0 and channel_count > drop_threshold: - mcdata = self.drop_qcd_s_channels(mcdata) - mcdata = self.drop_empty_channels(mcdata) - - channels = [] - t_channel_mode = self.t_channel_mode( - self.process.run_card["phasespace"]["t_channel"] - ) - for channel_id, (chan_topologies, chan_permutations, chan_indices, active_flavors) in enumerate(zip( - mcdata.topologies, mcdata.permutations, mcdata.channel_weight_indices, - mcdata.active_flavors - )): - topo_count = len(chan_topologies) - for topo_index, (topo, indices) in enumerate(zip(chan_topologies, chan_indices)): - mapping = ms.PhaseSpaceMapping( - chan_topologies[0], - self.process.e_cm, - t_channel_mode=t_channel_mode, - cuts=self.cuts, - invariant_power=self.process.run_card["phasespace"]["invariant_power"], - permutations=chan_permutations, - leptonic=self.process.leptonic, - ) - prefix = f"subproc{self.subproc_id}.channel{channel_id}" - if topo_count > 1: - prefix += f".subchan{topo_index}" - discrete_sym, discrete_flavor = self.build_discrete( - len(chan_permutations), len(self.meta["flavors"]), prefix - ) - channels.append(Channel( - phasespace_mapping = mapping, - adaptive_mapping = self.build_vegas(mapping, prefix), - discrete_sym = discrete_sym, - discrete_flavor = discrete_flavor, - channel_weight_indices = indices, - name = f"{channel_id}", - active_flavors = active_flavors, - )) - - remapped_chan_count = sum( - len(indices) for indices in mcdata.channel_indices - ) - if self.process.run_card["phasespace"]["sde_strategy"] == "denominators": - prop_chan_weights = ms.PropagatorChannelWeights( - [topo[0] for topo in mcdata.topologies], mcdata.permutations, - mcdata.channel_indices - ) - chan_weight_remap = [] - else: - prop_chan_weights = None - chan_weight_remap = [ - [ - len(mcdata.symfact) if remap == -1 else remap - for remap in amp2_remap - ] - for amp2_remap in mcdata.amp2_remaps - ] - - if any(len(topos) > 1 for topos in mcdata.topologies): - subchan_weights = ms.SubchannelWeights( - mcdata.topologies, mcdata.permutations, mcdata.channel_indices - ) - else: - subchan_weights = None - - return PhaseSpace( - mode="multichannel", - channels=channels, - first_chan_weight_remap=chan_weight_remap, - first_remapped_chan_count=remapped_chan_count, - symfact=mcdata.symfact, - prop_chan_weights=prop_chan_weights, - subchan_weights=subchan_weights, - ) - - def build_flat_phasespace(self) -> PhaseSpace: - mapping = ms.PhaseSpaceMapping( - self.incoming_masses + self.outgoing_masses, - self.process.e_cm, - mode=self.t_channel_mode(self.process.run_card["phasespace"]["flat_mode"]), - cuts=self.cuts, - leptonic=self.process.leptonic, - ) - prefix = f"subproc{self.subproc_id}.flat" - discrete_sym, discrete_flavor = self.build_discrete( - 1, len(self.meta["flavors"]), prefix - ) - channel = Channel( - phasespace_mapping = mapping, - adaptive_mapping = self.build_vegas(mapping, prefix), - discrete_sym = discrete_sym, - discrete_flavor = discrete_flavor, - channel_weight_indices = [0], - name = "F", - active_flavors = [], - ) - if self.unmerged_meta is None: - remap = [list(range(self.meta["diagram_count"]))] - else: - remap = [ - list(range(self.unmerged_meta[subproc]["diagram_count"])) - for subproc in self.meta["subprocesses"] - ] - return PhaseSpace( - mode="flat", - channels=[channel], - first_chan_weight_remap=remap, - first_remapped_chan_count=1, - symfact=[None], - ) - - def simplify_phasespace( - self, - multi_phasespace: PhaseSpace, - cross_sections: list[float] - ) -> PhaseSpace | None: - assert multi_phasespace.mode == "multichannel" - - threshold = 1 - self.process.run_card["phasespace"]["combine_channel_threshold"] - kept_channels = [] - tot_cs = sum(cross_sections) - cum_cs = 0. - seen_active_flavors = set() - #seen_resonances = set() - for index, (cs, chan) in sorted( - enumerate(zip(cross_sections, multi_phasespace.channels)), - key=lambda pair: pair[1][0], - reverse=True - ): - cum_cs += cs - has_unseen_flavors = False - has_unseen_resonances = False - for flavs in chan.active_flavors: - for flav in flavs: - if flav not in seen_active_flavors: - has_unseen_flavors = True - seen_active_flavors.add(flav) - #for resonance in chan.resonances: - # if resonance not in seen_resonances: - # has_unseen_resonances = True - # seen_resonances.add(flav) - if has_unseen_flavors or has_unseen_resonances or not tot_cs \ - or cum_cs / tot_cs < threshold: - kept_channels.append(index) - if len(kept_channels) >= len(cross_sections) - 1: - return multi_phasespace - - channels = [] - channel_map = {} - symfact = [] - for old_chan_index in kept_channels: - channel = multi_phasespace.channels[old_chan_index] - perm_count = max(1, channel.phasespace_mapping.channel_count()) - channel_index = len(symfact) - symfact.append(None) - symfact.extend([channel_index] * (perm_count - 1)) - channel_map.update({ - old_index: new_index - for new_index, old_index in enumerate( - channel.channel_weight_indices, start=channel_index - ) - }) - channels.append(Channel( - phasespace_mapping = channel.phasespace_mapping, - adaptive_mapping = channel.adaptive_mapping, - discrete_sym = channel.discrete_sym, - discrete_flavor = channel.discrete_flavor, - channel_weight_indices = list(range( - channel_index, channel_index + perm_count - )), - name = channel.name, - active_flavors = channel.active_flavors, - event_generator = channel.event_generator, - )) - - flat_phasespace = self.build_flat_phasespace() - flat_channel = flat_phasespace.channels[0] - channels.append(Channel( - phasespace_mapping = flat_channel.phasespace_mapping, - adaptive_mapping = flat_channel.adaptive_mapping, - discrete_sym = flat_channel.discrete_sym, - discrete_flavor = flat_channel.discrete_flavor, - channel_weight_indices = [len(symfact)], - name = flat_channel.name, - active_flavors = flat_channel.active_flavors, - )) - flat_index = len(symfact) - symfact.append(None) - channel_map[len(multi_phasespace.symfact)] = len(symfact) - if multi_phasespace.subchan_weights is None and len(multi_phasespace.first_chan_weight_remap) > 0: - first_chan_weight_remap = [ - [ - channel_map.get(remap, flat_index) - for remap in cw_remap - ] - for cw_remap in multi_phasespace.first_chan_weight_remap - ] - first_remapped_chan_count = len(symfact) - second_chan_weight_remap = [] - second_remapped_chan_count = 0 - else: - first_chan_weight_remap = multi_phasespace.first_chan_weight_remap - first_remapped_chan_count = multi_phasespace.first_remapped_chan_count - chan_count = multi_phasespace.first_remapped_chan_count if multi_phasespace.subchan_weights is None else multi_phasespace.subchan_weights.channel_count() - second_chan_weight_remap = [ - channel_map.get(i, flat_index) - for i in range(chan_count) - ] - second_remapped_chan_count = len(symfact) - - return PhaseSpace( - mode="both", - channels=channels, - first_chan_weight_remap=first_chan_weight_remap, - first_remapped_chan_count=first_remapped_chan_count, - second_chan_weight_remap=second_chan_weight_remap, - second_remapped_chan_count=second_remapped_chan_count, - symfact=symfact, - prop_chan_weights=multi_phasespace.prop_chan_weights, - subchan_weights=multi_phasespace.subchan_weights, - ) - - def drop_qcd_s_channels(self, mcdata: MultiChannelData) -> MultiChannelData: - """Drop channels with non-resonant QCD s-channel propagators, to reduce the - channel count for processes with many diagrams. Channels are grouped by the - number of QCD s-channels, allowing for more if necessary to map out all flavor - indices. Channel weights belonging to a dropped channel are left unmapped.""" - groups_by_s_count = {} - for index, count in enumerate(mcdata.qcd_s_channel_count): - groups_by_s_count.setdefault(count, []).append(index) - - covered_flavors = set() - kept_groups = set() - for s_count in sorted(groups_by_s_count): - s_count_groups = groups_by_s_count[s_count] - if s_count == 0: - selected = s_count_groups - else: - selected = [ - index - for index in s_count_groups - if any( - flav not in covered_flavors - for flavs in mcdata.active_flavors[index] - for flav in flavs - ) - ] - kept_groups.update(selected) - for index in selected: - for flavs in mcdata.active_flavors[index]: - covered_flavors.update(flavs) - - return self.keep_channel_groups(mcdata, sorted(kept_groups)) - - def drop_empty_channels(self, mcdata: MultiChannelData) -> MultiChannelData: - """Drop the channels whose phase space the cuts exclude entirely: a pair - mass cut whose floor lies at or above the top of an on-shell window of - the channel (ms.PhaseSpaceMapping.empty). Every point such a channel - generates fails the cuts, so it could only ever fail to find one. - - If that leaves a flavour with no channel at all, its cross section is - zero -- typically a cut that excludes the window of a decay chain every - diagram shares -- and the run cannot sample it: that is reported as an - error naming the cause, rather than as the survey finding no point - passing the cuts.""" - if self.cuts is None: - return mcdata - t_channel_mode = self.t_channel_mode( - self.process.run_card["phasespace"]["t_channel"] - ) - empty = [ - ms.PhaseSpaceMapping( - chan_topologies[0], - self.process.e_cm, - t_channel_mode=t_channel_mode, - cuts=self.cuts, - invariant_power=self.process.run_card["phasespace"]["invariant_power"], - permutations=chan_permutations, - leptonic=self.process.leptonic, - ).empty() - for chan_topologies, chan_permutations in zip( - mcdata.topologies, mcdata.permutations - ) - ] - if not any(empty): - return mcdata - kept_groups = [index for index, is_empty in enumerate(empty) if not is_empty] - covered_flavors = { - flav - for index in kept_groups - for flavs in mcdata.active_flavors[index] - for flav in flavs - } - uncovered = sorted({ - flav - for index, is_empty in enumerate(empty) if is_empty - for flavs in mcdata.active_flavors[index] - for flav in flavs - } - covered_flavors) - if uncovered: - raise ValueError( - f"subprocess {self.meta['incoming']} > {self.meta['outgoing']}: " - "the pair mass cuts exclude the whole on-shell window " - f"(bw_cutoff = {self.process.run_card['phasespace']['bw_cutoff']}) " - f"of every channel of flavour(s) {uncovered}, so their cross " - "section is zero; loosen the cut or remove that decay" - ) - logger.info( - "subprocess %d: %d channel(s) dropped, the cuts exclude the " - "on-shell window they sample", self.subproc_id, len(empty) - len(kept_groups) - ) - return self.keep_channel_groups(mcdata, kept_groups) - - def keep_channel_groups( - self, mcdata: MultiChannelData, kept_groups: list[int] - ) -> MultiChannelData: - """Rebuild the multi-channel data with only the channel groups (indices - into mcdata.topologies) in kept_groups. Channel weights belonging to a - dropped channel are left unmapped.""" - if len(kept_groups) == len(mcdata.topologies): - return mcdata - - amp2_remaps = [[-1] * len(remap) for remap in mcdata.amp2_remaps] - symfact = [] - topologies = [] - permutations = [] - channel_indices = [] - channel_weight_indices = [] - diagram_indices = [] - diagram_color_indices = [] - diagram_propagator_pdgs = [] - active_flavors = [] - qcd_s_channel_count = [] - channel_index = 0 - - for group in kept_groups: - chan_topologies = mcdata.topologies[group] - chan_permutations = mcdata.permutations[group] - chan_diagram_indices = mcdata.diagram_indices[group] - topo_count = len(chan_topologies) - - channel_index_first = channel_index - symfact_index_first = len(symfact) - for i, diag in enumerate(chan_diagram_indices): - if self.unmerged_meta is None: - amp2_remaps[0][diag] = channel_index - else: - for amp2_remap, d in zip(amp2_remaps, diag): - if d != -1: - amp2_remap[d] = channel_index - if i == 0: - symfact.extend([None] * topo_count) - else: - symfact.extend(range(symfact_index_first, symfact_index_first + topo_count)) - channel_index += 1 - - topologies.append(chan_topologies) - permutations.append(chan_permutations) - channel_indices.append(list(range(channel_index_first, channel_index))) - channel_weight_indices.append([ - [ - symfact_index_first + topo_index + i * topo_count - for i in range(len(chan_permutations)) - ] - for topo_index in range(topo_count) - ]) - diagram_indices.append(chan_diagram_indices) - if mcdata.diagram_color_indices: - diagram_color_indices.append(mcdata.diagram_color_indices[group]) - if mcdata.diagram_propagator_pdgs: - diagram_propagator_pdgs.append( - mcdata.diagram_propagator_pdgs[group] - ) - active_flavors.append(mcdata.active_flavors[group]) - qcd_s_channel_count.append(mcdata.qcd_s_channel_count[group]) - - return MultiChannelData( - amp2_remaps, - symfact, - topologies, - permutations, - channel_indices, - channel_weight_indices, - diagram_indices, - diagram_color_indices, - diagram_propagator_pdgs, - active_flavors, - qcd_s_channel_count, - ) - - def set_madnis_auto_settings(self, rsd: float): - madnis_args = self.process.run_card["madnis"] - n_out = len(self.meta["outgoing"]) - n_events = self.process.run_card["generation"]["events"] - is_gridpack = self.process.run_card["gridpack"]["save_gridpack"] - train_batches = madnis_args["train_batches"] - hidden_dim = min(max(int((7 * rsd) / 32) * 32 + 64, 64), 256) - flow_layers = 3 if rsd < 32 else 4 - lr = min(max((10000 - train_batches) / 8000 * 7e-4 + 3e-4, 3e-4), 1e-3) - if n_out <= 2: - enable = False - elif n_out == 3: - enable = rsd > 10. or n_events > 1000000 or is_gridpack - else: - enable = True - fixed_cwnet_fraction = max(0.33, 1.0 - 10000. / train_batches) - batch_size_per_channel = min(max(int((7 * rsd) / 32) * 32 + 64, 128), 512) - - self.madnis_settings = { - "enable": enable, - "flow_layers": flow_layers, - "flow_hidden_dim": hidden_dim, - "discrete_hidden_dim": hidden_dim, - "cwnet_hidden_dim": hidden_dim, - "train_batches": train_batches, - "lr": lr, - "fixed_cwnet_fraction": fixed_cwnet_fraction, - "batch_size_per_channel": batch_size_per_channel, - } - for key, value in self.madnis_settings.items(): - run_card_value = madnis_args[key] - if run_card_value != "auto": - self.madnis_settings[key] = run_card_value - - def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: - madnis_args = self.process.run_card["madnis"] - # Shared across all networks below: initialize_globals() derives an - # independent, non-colliding stream per tensor from this one base seed. - seed = self.process.run_seed - channels = [] - for channel_id, channel in enumerate(phasespace.channels): - prefix = f"subproc{self.subproc_id}.channel{channel_id}" - cond_dim = 0 - - flow_dim = channel.phasespace_mapping.random_dim() - flow = ms.Flow( - input_dim=flow_dim, - condition_dim=cond_dim, - prefix=prefix, - bin_count=madnis_args["flow_spline_bins"], - subnet_hidden_dim=self.madnis_settings["flow_hidden_dim"], - subnet_layers=self.madnis_settings["flow_layers"], - subnet_activation=self.activation(madnis_args["flow_activation"]), - invert_spline=madnis_args["flow_invert_spline"], - ) - if channel.adaptive_mapping is None: - flow.initialize_globals(self.process.contexts[0], seed) - else: - flow.initialize_from_vegas( - self.process.contexts[0], channel.adaptive_mapping.grid_name(), - seed - ) - cond_dim += flow_dim - - # discrete_sym runs after the adaptive map, so it can condition on its latent. - discrete_sym = channel.discrete_sym - if discrete_sym is not None: - perm_count = channel.phasespace_mapping.channel_count() - discrete_sym = ms.DiscreteFlow( - option_counts=[perm_count], - prefix=f"{prefix}.discrete_flow_sym", - dims_with_prior=[], - condition_dim=cond_dim, - subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], - subnet_layers=madnis_args["discrete_layers"], - subnet_activation=self.activation(madnis_args["discrete_activation"]), - ) - discrete_sym.initialize_globals(self.process.contexts[0], seed) - cond_dim += perm_count - - discrete_flavor = channel.discrete_flavor - if discrete_flavor is not None: - discrete_flavor = ms.DiscreteFlow( - option_counts=[len(self.meta["flavors"])], - prefix=f"{prefix}.discrete_flow_flavor", - dims_with_prior=[0], - condition_dim=cond_dim, - subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], - subnet_layers=madnis_args["discrete_layers"], - subnet_activation=self.activation(madnis_args["discrete_activation"]), - ) - discrete_flavor.initialize_globals(self.process.contexts[0], seed) - - channels.append(Channel( - phasespace_mapping = channel.phasespace_mapping, - adaptive_mapping = flow, - discrete_sym = discrete_sym, - discrete_flavor = discrete_flavor, - channel_weight_indices = channel.channel_weight_indices, - name = channel.name, - active_flavors = channel.active_flavors, - )) - - return PhaseSpace( - mode="both", - channels=channels, - first_chan_weight_remap=phasespace.first_chan_weight_remap, - first_remapped_chan_count=phasespace.first_remapped_chan_count, - second_chan_weight_remap=phasespace.second_chan_weight_remap, - second_remapped_chan_count=phasespace.second_remapped_chan_count, - symfact=phasespace.symfact, - cwnet=self.build_cwnet(len(phasespace.symfact)), - prop_chan_weights=phasespace.prop_chan_weights, - subchan_weights=phasespace.subchan_weights, - ) - - def build_vegas(self, mapping: ms.PhaseSpaceMapping, prefix: str) -> ms.VegasMapping: - if not self.process.run_card["vegas"]["enable"]: - return None - - vegas = ms.VegasMapping( - mapping.random_dim(), - self.process.run_card["vegas"]["bins"], - prefix, - ) - for context in self.process.contexts: - vegas.initialize_globals(context) - return vegas - - def build_discrete( - self, permutation_count: int, flavor_count: int, prefix: str - ) -> tuple[ms.DiscreteSampler | None, ms.DiscreteSampler | None]: - is_adaptive = self.process.run_card["phasespace"]["adaptive_symmetry_sampling"] - if is_adaptive and permutation_count > 1: - discrete_sym = ms.DiscreteSampler( - [permutation_count], f"{prefix}.discrete_sym" - ) - for context in self.process.contexts: - discrete_sym.initialize_globals(context) - else: - discrete_sym = None - - if flavor_count > 1: - # the flavor probabilities are conditioned on the PDF prior; a - # leptonic process (or a decay) has none, and the Integrand then - # passes no condition: declaring one anyway fails to build - # ("keys and values must have the same size"), e.g. z > q q~ with - # both u- and d-type rows - discrete_flavor = ms.DiscreteSampler( - [flavor_count], f"{prefix}.discrete_flavor", - [] if self.process.leptonic else [0] - ) - for context in self.process.contexts: - discrete_flavor.initialize_globals(context) - else: - discrete_flavor = None - - return discrete_sym, discrete_flavor - - def build_cwnet(self, channel_count: int) -> ms.ChannelWeightNetwork: - #if channel_count == 1: - # return None - madnis_args = self.process.run_card["madnis"] - cwnet = ms.ChannelWeightNetwork( - channel_count=channel_count, - particle_count=self.particle_count, - hidden_dim=self.madnis_settings["cwnet_hidden_dim"], - layers=madnis_args["cwnet_layers"], - activation=self.activation(madnis_args["cwnet_activation"]), - prefix=f"subproc{self.subproc_id}.cwnet", - ) - cwnet.initialize_globals( - self.process.contexts[0], self.process.run_seed - ) - return cwnet - - def t_channel_mode(self, name: str) -> ms.PhaseSpaceMapping.TChannelMode: - modes = { - "propagator": ms.PhaseSpaceMapping.propagator, - "rambo": ms.PhaseSpaceMapping.rambo, - "chili": ms.PhaseSpaceMapping.chili, - } - if name in modes: - return modes[name] - else: - raise ValueError(f"Invalid t-channel mode '{name}'") - - def activation(self, name: str) -> ms.MLP.Activation: - activations = { - "relu": ms.MLP.relu, - "leaky_relu": ms.MLP.leaky_relu, - "elu": ms.MLP.elu, - "gelu": ms.MLP.gelu, - "sigmoid": ms.MLP.sigmoid, - "softplus": ms.MLP.softplus, - } - if name in activations: - return activations[name] - else: - raise ValueError(f"Invalid activation function '{name}'") - - def build_integrands( - self, - phasespace: PhaseSpace, - madnis_training: bool = False, - drop_cuts_and_rescale: bool = False - ) -> list[ms.Integrand]: - flavors = [] - flavor_remap = [] - flavor_factors = [] - flavor_mirror = [] - flavor_diff_xs_indices = [] - flavor_subproc_indices = [] - flavor_per_subproc_remap = [] - - for flav in self.meta["flavors"]: - if self.unmerged_meta is not None: - diff_xs_index = flav["subprocess"] - subproc_index = self.meta["subprocesses"][diff_xs_index] - ps_flavor = flav["flavor"] - flavor_diff_xs_indices.append(diff_xs_index) - flavor_subproc_indices.append(subproc_index) - flavor_per_subproc_remap.append(ps_flavor) - flav = self.unmerged_meta[subproc_index]["flavors"][ps_flavor] - flavors.append(flav["options"][0]) - flavor_remap.append(flav["index"]) - flavor_factors.append(len(flav["options"])) - flavor_mirror.append(flav["mirror"]) - - cross_sections = [] - self.step_matrix_elements = [] - for me_index, matrix_element in enumerate(self.matrix_elements): - if self.surrogates: - # the surrogate in the integrand, the full matrix element in - # the second pass, with the random numbers passed in - mat = ms.MatrixElement( - self.surrogates[me_index].index(), - self.particle_count, - [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, - ms.MatrixElement.flavor_in], - [ms.MatrixElement.matrix_element_out], - 1, - False, - self.process.me_frame, - self.process.incoming_count, - ) - self.step_matrix_elements.append(ms.MatrixElement( - matrix_element, - [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, - ms.MatrixElement.flavor_in, ms.MatrixElement.random_color_in, - ms.MatrixElement.random_helicity_in, - ms.MatrixElement.random_diagram_in], - [ms.MatrixElement.matrix_element_out, - ms.MatrixElement.color_index_out, - ms.MatrixElement.helicity_index_out, - ms.MatrixElement.diagram_index_out], - False, - self.process.me_frame, - self.process.incoming_count, - )) - elif matrix_element: - mat = ms.MatrixElement( - matrix_element, - ms.Integrand.matrix_element_inputs, - ms.Integrand.matrix_element_outputs, - True, - self.process.me_frame, - self.process.incoming_count, - ) - else: - #TODO: not working in merged mode - mat = ms.MatrixElement( - 0xBADCAFE, - self.particle_count, - ms.Integrand.matrix_element_inputs, - ms.Integrand.matrix_element_outputs, - self.meta["diagram_count"], - True, - self.process.me_frame, - self.process.incoming_count, - ) - pdf_grid = None if self.process.leptonic else self.process.pdf_grid - pdf_arg = None if self.process.leptonic else ms.CachedPdf() - cross_sections.append( - ms.DifferentialCrossSection( - matrix_element=mat, - # For a decay this is the decaying particle's mass, and the - # flux becomes 1/(2M): the result is a partial width in - # GeV, not a cross section in pb. - cm_energy=self.process.e_cm, - running_coupling=None, - energy_scale=ms.CachedScale(), - pid_options=[], - pdf1=pdf_arg, - pdf2=pdf_arg, - input_momentum_fraction=True, - decay=self.process.is_decay, - return_matrix_element=bool(self.surrogates), - ) - ) - partial_weights = self.process.systematics_enabled() - madnis_args = self.process.run_card["madnis"] - custom_bias = self.process.custom_function("bias", self.particle_count, flavors) - integrands = [] - for channel in phasespace.channels: - integrands.append(ms.Integrand( - channel.phasespace_mapping, - cross_sections, - channel.adaptive_mapping, - channel.discrete_sym, - channel.discrete_flavor, - flavors, - pdf_grid, - self.process.running_coupling, - self.scale, - phasespace.prop_chan_weights, - phasespace.subchan_weights, - phasespace.cwnet, - phasespace.first_chan_weight_remap, - phasespace.first_remapped_chan_count, - phasespace.second_chan_weight_remap, - phasespace.second_remapped_chan_count, - madnis_training, - drop_cuts_and_rescale, - partial_weights, - channel.channel_weight_indices, - channel.active_flavors, - flavor_remap, - flavor_factors, - flavor_mirror, - flavor_diff_xs_indices, - flavor_subproc_indices, - flavor_per_subproc_remap, - madnis_args["compressed_channel_weight_count"], - custom_bias, - )) - #print(integrands[1].function()) - #for i in integrands: print(i.function()) - return integrands - - -_ROOTED_OPTIONS = ('lhapdf', 'lhapdf_py3', 'lhapdf_py2', 'heptools_install_dir') - -def load_mg5_options(me_dir=None) -> dict: - """Read the tool paths from the MadGraph configuration, so the launcher - knows where LHAPDF and the optional programs (Pythia8/Delphes/MadSpin/ - reweight/analysis) live. - - Files are read least- to most-specific, each overriding the previous, which - is the same layering CommonRunCmd.set_configuration uses: the card inside - this process directory has the last word, then this installation, then the - per-user file. Relative values are resolved against the root of the file - they came from. - - ``me_dir`` is the process directory whose Cards/me5_configuration.txt has - the last word. It defaults to the working directory, which is right for - bin/generate_events (it chdirs there first) but not for an in-process - caller, which has not moved yet. - """ - - import madgraph - mg5dir = os.path.dirname(os.path.dirname(os.path.abspath(madgraph.__file__))) - me_dir = os.getcwd() if me_dir is None else me_dir - - options = { - 'pythia-pgs_path': None, 'pythia8_path': None, - 'madanalysis5_path': None, 'exrootanalysis_path': None, 'delphes_path': None, - 'rivet_path': None, 'contur_path': None, 'f2py_compiler': None, - 'lhapdf': None, 'lhapdf_py3': None, 'lhapdf_py2': None, 'timeout': 0, - 'heptools_install_dir': None, 'hepmc3_path': None, - } - config_files = [] - base_config = misc.base_config_file() - if base_config: - config_files.append((base_config, mg5dir)) - user_config = misc.user_config_file() - if user_config: - config_files.append((user_config, mg5dir)) - config_files.append((misc.install_config_file(mg5dir), mg5dir)) - config_files.append((os.path.join(me_dir, 'Cards', 'me5_configuration.txt'), - me_dir)) - for cfg, root in config_files: - if not os.path.exists(cfg): - continue - with open(cfg) as fsock: - for line in fsock: - line = line.split('#', 1)[0] - if '=' not in line: - continue - name, value = (x.strip() for x in line.split('=', 1)) - if name not in options or value in ('', 'None'): - continue - if (name.endswith('_path') or name in _ROOTED_OPTIONS) and \ - not os.path.isabs(value) and \ - (os.sep in value or value.startswith('.')): - value = os.path.normpath(os.path.join(root, value)) - options[name] = value - options['mg5_path'] = mg5dir # enables MadSpin/reweight - return options - - -_LHAPDF = None - -def lhapdf_paths(refresh: bool = False) -> misc.LhapdfPaths: - """The LHAPDF installation this run should use, resolved once. - - Uses the same helper as the 'launch' command, so driving a process - directory straight through bin/generate_events behaves identically to - launching it from MadGraph. - """ - - global _LHAPDF - if _LHAPDF is None or refresh: - _LHAPDF = misc.resolve_lhapdf(load_mg5_options()) - # Export what we found: the real LHAPDF library, systematics and the - # MadSpin/reweight subprocesses read these from the environment. - if _LHAPDF.data_paths: - os.environ["LHAPDF_DATA_PATH"] = os.pathsep.join(_LHAPDF.data_paths) - if _LHAPDF.config: - os.environ["MADGRAPH_LHAPDF_CONFIG"] = _LHAPDF.config - return _LHAPDF - - -class MG7Cmd(Cmd): - """The mg7 run interface: the command object that drives a launch. - - It plays two roles, and they are the same object on purpose. - - 1. It is the *mother* of the merged switch/card question (the thing - ``AskRunEditCard`` calls back into for ``keep_cards``/``do_open``/ - ``compute_widths``). - 2. It is the child command interface MG5 registers via - ``define_child_cmd_interface`` when you type ``launch`` on an mg7 output. - - Role 2 is what makes role 1 behave. ``Cmd.ask`` resolves a scripted answer - through ``self.check_answer_in_input_file``, where ``self`` is the mother -- - so once MG5 hands us its ``inputfile`` and registers itself as our - ``mother``, the question reads the very lines that follow ``launch`` in the - user's command file, and any line it does not recognise goes back up to MG5 - through ``store_line`` instead of being silently consumed. This is exactly - how a madevent output already behaves; before this class existed the mg7 - branch shelled out to ``bin/generate_events`` and had to guess which - following lines belonged to the run, which quietly swallowed every ``set``. - """ - - prompt = 'MG7> ' - - def __init__(self, me_dir='.', options=None, *args, **opts): - super().__init__(*args, **opts) - self.me_dir = os.path.abspath(me_dir) if me_dir != '.' else '.' - # The rest of the launcher is written against the process directory as - # the working directory, so do_generate_events chdirs there; but at - # construction time MG5 has not moved, hence the explicit me_dir. - self.options = load_mg5_options( - self.me_dir if self.me_dir != '.' else None) - if options: - # MG5's own resolved tool paths win: they are the ones the user - # configured in the session that is doing the launching. - self.options.update({k: v for k, v in options.items() - if k in self.options and v is not None}) - self.plugin_path = [] - self.proc_characteristics = {'grouped_matrix': False, 'limitations': []} - self._model = None # lazily imported by get_model() - - # ------------------------------------------------------------------ - # the contract the card question expects from its mother - # ------------------------------------------------------------------ - def keep_cards(self, need_card=[], ignore=[]): - from madgraph.interface.common_run_interface import CommonRunCmd - return CommonRunCmd.keep_cards(self, need_card, ignore) - - def do_open(self, line): - from madgraph.interface.common_run_interface import CommonRunCmd - CommonRunCmd.do_open(self, line) - - def check_open(self, args): - from madgraph.interface.common_run_interface import CommonRunCmd - CommonRunCmd.check_open(self, args) - - def get_model(self): - """The model this process was generated with (an ``import_ufo`` model). - - This is the hook the card question goes through for everything that - needs the model rather than the card: 'update dependent' (which resets - the masses/widths the model computes from the free parameters) and the - generic auto-width handler both call ``mother_interface.get_model()``. - Without it the launch question could only warn that it had failed to - update the dependent parameters, and MadSpin/the shower then read - whatever inconsistent values the param_card still held. - - None when the process recorded no model (an output made before - SubProcesses/model.txt existed) -- the caller warns in that case. - """ - if self._model is None: - me_dir = self.me_dir if self.me_dir != '.' else None - # False (not None) so a model that cannot be imported is not - # retried at every question. - self._model = load_process_model(me_dir) or False - return self._model or None - - def do_compute_widths(self, line): - # The interactive card editor delegates 'auto' width computation to - # the mother interface. Reuse the runtime helper (madgraph subprocess - # + the model stored at output time). ``line`` looks like - # " --path= [--nlo]"; we only need the card path. - m = re.search(r'--path=(\S+)', line or "") - path = m.group(1) if m else os.path.join("Cards", "param_card.dat") - compute_auto_widths(path) - # return an empty mapping: the caller iterates out.items() for the - # small-width treatment, which mg7 does not apply. - return {} - - # ------------------------------------------------------------------ - # running - # ------------------------------------------------------------------ - def do_generate_events(self, line): - """generate_events [-f] [-n NAME] [--laststep=parton] - - Run the mg7 generation and the selected post-processing tools.""" - - opts = self._parse_run_options(line) - cwd = os.getcwd() - target = self.me_dir if self.me_dir != '.' else cwd - try: - os.chdir(target) - _setup_logging() - if opts['name']: - self._set_run_name(opts['name']) - switch = {} - if not opts['force']: - switch = ask_edit_cards(mother=self) - switch = self._apply_laststep(switch, opts) - force_lhe_output_if_needed(switch) - _raise_open_file_limit() - run_generation(switch) - finally: - os.chdir(cwd) - - # ``launch`` is the name MG5 users type; keep it working here too. - do_launch = do_generate_events - - def do_npy_to_lhe(self, line): - """npy_to_lhe [RUN_NAME] [-o OUTPUT] [--seed N] [--no-gzip] - - Write the LHE file (events.lhe.gz) of a run that produced npy events, - e.g. to shower it. Without a run name, the most recent such run.""" - import shlex - me_dir = self.me_dir if self.me_dir != '.' else os.getcwd() - try: - npy_to_lhe.main(shlex.split(line), me_dir=me_dir) - except SystemExit as error: - if error.code not in (None, 0): - logger.error(str(error.code)) - - def do_quit(self, line): - """Leave the mg7 run interface.""" - return super().do_quit(line) - - do_exit = do_quit - - # ------------------------------------------------------------------ - # helpers - # ------------------------------------------------------------------ - def _parse_run_options(self, line): - """Understand the subset of MG5's `launch` options that mean something - for an mg7 run, and say so out loud for the ones that do not.""" - out = {'force': False, 'name': '', 'laststep': ''} - args = self.split_arg(line) - index = 0 - while index < len(args): - arg = args[index] - index += 1 - if arg in ('-f', '--force'): - out['force'] = True - elif arg in ('-n', '--name') and index < len(args): - out['name'] = args[index] - index += 1 - elif arg.startswith(('-n=', '--name=')): - out['name'] = arg.split('=', 1)[1] - elif arg in ('-s', '--laststep') and index < len(args): - out['laststep'] = args[index] - index += 1 - elif arg.startswith(('-s=', '--laststep=')): - out['laststep'] = arg.split('=', 1)[1] - elif arg.startswith('-'): - logger.warning( - "'%s' is not supported by the mg7 output and is ignored; " - "the equivalent settings live in Cards/run_card.toml", arg) - elif arg and not out['name']: - out['name'] = arg - return out - - def _set_run_name(self, name): - """Honour `launch -n NAME`: the mg7 run directory is - Events/_NN, so the name is a run_card setting.""" - path = os.path.join("Cards", "run_card.toml") - run_card = RunCardMG7(path, consistency=False) - if run_card["run"]["run_name"] != name: - run_card["run"]["run_name"] = name - run_card.write(path) - logger.info("run name set to '%s'", name) - - def _apply_laststep(self, switch, opts): - """`--laststep=parton` means: generate the events, run nothing on them. - - The mg7 tool selection is the switch dict the question returns, so - stopping at parton level is switching those entries off. - """ - if opts['laststep'] not in ('parton', 'auto', ''): - logger.warning( - "--laststep=%s is not supported by the mg7 output; " - "select the programs to run in the launch question instead", - opts['laststep']) - return switch - if opts['laststep'] != 'parton' or not switch: - return switch - switch = dict(switch) - for key in ('shower', 'detector', 'analysis', 'madspin', 'reweight'): - if not _off(switch.get(key, 'OFF')): - logger.info("--laststep=parton: not running %s", key) - switch[key] = 'OFF' - return switch - - -def build_selector_cmd(mother=None): - """Build the (monkey-patched) mother command + merged switch/card selector - used by the mg7 output. Returns the selector *class* and a mother instance - understood by AskRun/AskforEditCard. - - ``mother`` lets an in-process caller supply its own :class:`MG7Cmd` -- the - one MG5 registered as its child -- so that the question resolves scripted - answers against MG5's own command file (see :class:`MG7Cmd`). When it is - omitted a standalone instance is built, which is what ``bin/generate_events`` - gets. - """ - - from madgraph.interface.madevent_interface import AskRunEditCard - - from madgraph.various import banner as _banner_mod - from madgraph.various import misc as _misc - - class MG7Selector(AskRunEditCard): - """Merged switch/card question for the mg7 output. - - The mg7 run_card is a TOML file, so define_paths/init_run/do_set are - overridden *on this class* (not globally on AskforEditCard) to treat it - as such. Scoping them here is important: MadSpin/reweight spawn their own - legacy madevent runs in the same process, and a global patch would break - their (run_card.dat based) card editing.""" - - # param_card + the TOML run_card are always offered - always_cards = ['param_card.dat', 'run_card.toml'] - optional_cards = [] - - def define_paths(self, **opt): - super().define_paths(**opt) - self.paths["run"] = os.path.join(self.me_dir, "Cards", "run_card.toml") - self.paths["run_card.toml"] = os.path.join(self.me_dir, "Cards", "run_card.toml") - # the TOML run_card uses its own default file (concrete defaults - # written at output time); this powers "set default". - self.paths["run_default"] = os.path.join(self.me_dir, "Cards", "run_card_default.toml") - - def init_run(self, cards): - # Make sure the run_card is loaded as a RunCardMG7 (the generic - # editor does not recognise run_card.toml), else "set " fails. - out = super().init_run(cards) - if not isinstance(getattr(self, "run_card", None), RunCardMG7): - toml_path = self.paths.get("run") or os.path.join( - self.me_dir, "Cards", "run_card.toml") - if os.path.exists(toml_path): - try: - # allow_scan so a run_card holding scan:[...] values loads - with _misc.TMP_variable(_banner_mod.RunCard, "allow_scan", True): - self.run_card = RunCardMG7(toml_path, consistency="warning") - self.run_set = list(self.run_card.keys()) - except Exception as err: - logger.warning("could not load %s: %s", toml_path, err) - if isinstance(getattr(self, "run_card", None), RunCardMG7): - self.run_card.allow_scan = True - return getattr(self, "run_set", out) - - def do_set(self, line, *args, **kwargs): - # madevent-style shortcuts (lhc/lep/fixed_scale/no_parton_cut), cut - # editing, energy units and arithmetic/mass expressions on the TOML - # run_card, intercepted before delegating to the generic editor. - targs = self.split_arg(line) - run_card = getattr(self, "run_card", None) - if isinstance(run_card, RunCardMG7) and targs: - start = 1 if targs[0] == "run_card" else 0 - if len(targs) > start: - name = targs[start] - nlow = name.lower() - rest = " ".join(targs[start + 1:]).split("#")[0].strip() - masses = run_card.get_mass_shortcuts(getattr(self, "param_card", None)) - - if nlow in ("no_parton_cut", "nocut", "no_cut"): - run_card.remove_all_cut() - logger.info("removing all cuts from the run_card.toml") - self.modified_card.add("run") - return - # [histograms] is a list of observables, not a parameter, - # so "set histograms ..." can only mean the whole section. - if nlow in ("histograms", "histogram", "plots"): - value = rest.lower() - if value in ("off", "none", "no", "false", "0"): - run_card.remove_all_histograms() - logger.info( - "removing all histograms from the run_card.toml") - self.modified_card.add("run") - return - if value in ("default", "on"): - if run_card.restore_default_histograms( - self.paths.get("run_default")): - logger.info( - "restored the %d histograms written at " - "output time", - len(run_card["histograms"])) - self.modified_card.add("run") - else: - logger.warning( - "no run_card_default.toml to restore the " - "histograms from") - return - logger.warning( - "'set histograms %s': only OFF (remove them all) " - "and default (put the generated ones back) are " - "understood; edit the [histograms] section of " - "Cards/run_card.toml for anything else", - rest or "(no value)") - return - if nlow in ("lhc", "lep", "ilc", "lcc") and rest: - ecm = run_card.set_collider(nlow, rest, masses) - logger.info("set %s collider: e_cm = %s GeV", nlow, ecm) - self.modified_card.add("run") - return - if nlow == "fixed_scale" and rest: - val = run_card.set_fixed_scale(rest, masses) - logger.info("set fixed scales to %s GeV", val) - self.modified_card.add("run") - return - - if rest and run_card.is_cut_name(name): - cut, bound, val = run_card.set_cut(name, run_card.evaluate(rest, masses)) - logger.info("modify cut %s.%s of the run_card.toml to %s", cut, bound, val) - self.modified_card.add("run") - return - - # 'iseed' is the madevent spelling of the mg7 run.seed, - # but the two disagree on how to ask for a random seed: - # madevent uses 0, mg7 uses -1 (0 being a perfectly good - # fixed seed there). Translate here rather than in - # _LO_SCALAR_MAP, which also drives the LO -> MG7 run_card - # conversion and would silently pin those runs to seed 0. - if rest and nlow == 'iseed': - value = run_card.evaluate(rest, masses) - try: - value = int(value) - except (TypeError, ValueError): - value = None - if value is None: - logger.warning("ignoring 'set iseed %s': not an integer", rest) - return - seed = -1 if value == 0 else value - run_card.set('run.seed', seed, user=True) - logger.info("set iseed (mg7 run.seed) of the run_card.toml to %s%s", - seed, " (random)" if seed == -1 else "") - self.modified_card.add("run") - return - - # legacy madevent run_card names -> mg7 "section.key", so a - # madevent-style launch script ("set nevents 500", "set - # use_syst F", "set bwcutoff 10", ...) edits the mg7 - # run_card.toml verbatim. RunCardMG7._LO_SCALAR_MAP is the - # same rename table used by the LO->MG7 run_card conversion. - if rest and nlow in run_card._LO_SCALAR_MAP \ - and nlow not in [k.lower() for k in run_card.keys()]: - target = run_card._LO_SCALAR_MAP[nlow] - run_card.set(target, rest, user=True) - logger.info("set %s (mg7 %s) of the run_card.toml to %s", - nlow, target, run_card[target]) - self.modified_card.add("run") - return - - if rest and nlow in [k.lower() for k in run_card.keys()]: - current = run_card[nlow] - if isinstance(current, (int, float)) and not isinstance(current, bool): - resolved = run_card.evaluate(rest, masses) - if not isinstance(resolved, str): - prefix = "run_card " if start == 1 else "" - line = "%s%s %s" % (prefix, name, resolved) - return super().do_set(line, *args, **kwargs) - - return MG7Selector, (mother if mother is not None else MG7Cmd()) - - -def ask_edit_cards(mother=None) -> dict: - """Single (MadDM-style) question letting the user both pick which programs - to run after generation and edit the associated cards. Returns the switch - dict describing the selected tools. - - ``mother`` is the :class:`MG7Cmd` asking the question. Passing the instance - MG5 registered as its child is what lets a scripted ``launch`` answer this - question from MG5's own command file. - """ - - selector_class, mother = build_selector_cmd(mother) - # path_msg is what makes Cmd.check_answer_in_input_file accept a bare path as - # an answer (its "elif path:" branch), so that a scripted launch can hand the - # question a card/banner path -- as the question itself advertises -- and have - # it replace the corresponding card. Without it the path is rejected ("This - # answer is not valid for current question") and the default is used instead. - try: - switch, question = mother.ask('', '0', [], path_msg='enter path', - ask_class=selector_class, - mode='auto', line_args=[], force=False, - return_instance=True) - except BaseException: - # interrupted (e.g. ctrl-C) before the prune below could run: prune - # here too, using the instance ask() stashed pre-interrupt, or the - # eagerly-materialised tool cards get misread as "left ON" next time. - instance = getattr(mother, '_last_ask_instance', None) - if instance is not None and hasattr(instance, 'active_cards'): - try: - mother.keep_cards(instance.active_cards()) - except Exception as error: - logger.debug('could not prune tool cards after aborted launch question: %s', error) - raise - switch = dict(switch) - prune_unselected_tool_cards(mother, question, switch) - return switch - - -def prune_unselected_tool_cards(mother, question, switch) -> None: - """Hide the cards of the tools that were NOT selected in the question. - - The question materialises every candidate tool card (copying the - *_default.dat) so that it can offer them all for edition; madevent's - ask_run_configuration prunes them again afterwards (keep_cards on the - selected cards), but the mg7 launcher had no such cleanup. The materialised - cards therefore persisted, and a later re-launch of the same output saw them - all present and defaulted every switch to ON (set_default_ keys off - card presence). Prune here too, using the question's own switch->card map, so - a re-launch reflects the previous selection instead of turning everything on. - """ - try: - keep = list(question.always_cards) - for spec in question.switch_cards: - if spec['on'](switch): - keep.append(spec['card']) - mother.keep_cards(keep) - except Exception as error: - logger.debug('could not prune unselected tool cards: %s', error) - - -def _find_event_file(run_path): - """Locate the LHE event file produced by generate_events (if any).""" - for name in ("events.lhe", "events.lhe.gz"): - path = os.path.join(run_path, name) - if os.path.exists(path): - return path - return None - - -def _find_or_convert_event_file(run_path, log): - """The run's LHE event file, converting its npy events when that is all - the run wrote (see npy_to_lhe).""" - lhe_path = _find_event_file(run_path) - if lhe_path is None and npy_to_lhe.can_convert(run_path): - log.info("No LHE event file in %s: converting the npy events.", run_path) - try: - lhe_path = npy_to_lhe.convert(run_path) - except Exception as error: - _report_failure(log, "npy->LHE conversion", error, run_path) - return lhe_path - - -def _report_failure(log, what, error, directory=None): - """Log a post-processing failure and write the full traceback to a file - (whose path is printed) so the problem can be investigated.""" - import traceback - if not directory or not os.path.isdir(directory): - directory = os.getcwd() - path = os.path.join(directory, "%s_crash.log" % what.replace(" ", "_")) - try: - with open(path, "w") as fsock: - fsock.write(traceback.format_exc()) - except Exception: - path = None - log.warning("%s failed: %s", what, error) - if path: - log.warning("full traceback written to: %s", path) - - -def _off(value) -> bool: - """True when a switch value means the tool was not selected.""" - return value in (None, "OFF", "Not Avail.", "Not Avail. (numpy missing)") - - -def _has_effective_handler(lg) -> bool: - """True when ``lg`` (or an ancestor it propagates to) already has a handler, - i.e. some other party is already displaying these records.""" - while lg: - if lg.handlers: - return True - if not lg.propagate: - return False - lg = lg.parent - return False - - -_TOOL_LOGGING_READY = False - - -def _setup_logging(): - """Make the reused madevent tool drivers' progress visible on screen. - - The drivers already narrate what they do through logger.info / - update_status ("Running Pythia8 [arXiv:...]", "Splitting .lhe event - file...", the live Idle/Running/Completed job counters, "Running - MadSpin", ...), but the standalone mg7 launcher never attaches a handler - to the 'madgraph'/'madevent' loggers, so all of it is swallowed. Attach a - colored INFO console handler to them (idempotent).""" - global _TOOL_LOGGING_READY - if _TOOL_LOGGING_READY: - return - if _has_effective_handler(logging.getLogger("madgraph")): - # Someone already configured logging for us -- in practice MG5, which - # now runs the launch in its own process. Its configuration is - # authoritative (it owns stdout_level and the tutorial logger); adding - # a second handler here would print every line twice. - _TOOL_LOGGING_READY = True - return - try: - import madgraph.interface.coloring_logging # registers ColorFormatter - formatter = logging.ColorFormatter("%(message)s") - except Exception: - formatter = logging.Formatter("%(message)s") - handler = logging.StreamHandler(sys.stdout) - handler.setFormatter(formatter) - handler._mg7_tool_handler = True - for name in ("madgraph", "madevent", "cmdprint", "madgraph7"): - lg = logging.getLogger(name) - if not any(getattr(h, "_mg7_tool_handler", False) for h in lg.handlers): - lg.addHandler(handler) - lg.setLevel(logging.INFO) - lg.propagate = False - _TOOL_LOGGING_READY = True - - -def selected_tools(switch) -> list: - """The post-processing drivers run_selected_tools runs for this switch.""" - # MadAnalysis5 hadron level analyses the shower/detector output, so it only - # makes sense when a shower ran (mirrors madevent's card gating: - # analysis == 'MadAnalysis5' and shower != 'OFF'). - ma5 = switch.get("analysis") == "MadAnalysis5" - showered = not _off(switch.get("shower")) - return [t for t, on in ( - ("reweighting", not _off(switch.get("reweight"))), - ("MadSpin", not _off(switch.get("madspin"))), - ("MadAnalysis5 (parton level)", ma5), - ("Pythia8 shower", switch.get("shower") == "Pythia8"), - ("Delphes", switch.get("detector") == "Delphes"), - ("MadAnalysis5 (hadron level)", ma5 and showered), - ("Rivet", switch.get("analysis") == "Rivet"), - ) if on] - - -def run_selected_tools(switch, process) -> None: - """Run the optional post-processing programs selected in the merged - question on the generated events. - - All the tools are driven through :class:`MG7RunCmd`, a thin adapter around - the standard madevent run interface (``MadEventCmd``). This means the mg7 - output reuses the *exact* madevent tool drivers -- in particular Pythia8 is - showered in parallel over the cluster/multicore backend, identically to a - madevent run -- instead of re-implementing each tool here. The command - sequence mirrors ``MadEventCmd.do_launch`` (reweight, MadSpin, MA5 parton, - shower, Delphes, MA5 hadron, Rivet); each tool checks its own card, so a - tool is invoked only when its switch is on. - """ - log = logging.getLogger("madevent") - - tools = selected_tools(switch) - if not tools: - # A switch that is not off does not necessarily select a driver here: - # "Not Avail." is not off, analysis = ExRoot has no mg7 driver, and a - # shower switch set to something other than Pythia8 selects nothing - # either. Without this a plain generate/output/launch announced a - # post-processing step with nothing after the colon, and paid for - # building the run interface (or an npy->LHE conversion) to do nothing. - return - - lhe_path = _find_or_convert_event_file(process.run_path, log) - if lhe_path is None: - log.warning("No LHE event file in %s; cannot run %s.", - process.run_path, ", ".join(tools)) - return - - run_name = os.path.basename(os.path.dirname(os.path.abspath(lhe_path))) - run_dir = os.path.dirname(os.path.abspath(lhe_path)) - - ma5 = switch.get("analysis") == "MadAnalysis5" - showered = not _off(switch.get("shower")) - - log.info("") - log.info("Post-processing the generated events with: %s", ", ".join(tools)) - - try: - from madgraph.iolibs.template_files.mg7.run_interface import MG7RunCmd - cmd = MG7RunCmd(os.getcwd(), load_mg5_options(), run_name, lhe_path) - cmd.set_run_name(run_name, None, "parton") - except Exception as error: - _report_failure(log, "post-processing setup", error, run_dir) - return - - def run(tool, command): - bar = "=" * 60 - log.info("") - log.info(bar) - log.info(" post-processing step: %s", tool) - log.info(bar) - start = time.time() - try: - cmd.exec_cmd(command, postcmd=False, printcmd=False) - except Exception as error: - _report_failure(log, tool, error, run_dir) - else: - log.info(" -> %s done (%.1fs)", tool, time.time() - start) - - # Same order as MadEventCmd.do_launch. The commands take no run name and so - # act on cmd.run_name: reweight runs on the generated events, then - # decay_events (MadSpin) repoints cmd.run_name to the "_decayed_i" run, - # so the shower and the analyses that follow act on the decayed events -- - # exactly as a madevent run does. Each driver still guards on its own card. - if not _off(switch.get("reweight")): - run("reweight", "reweight -from_cards") - if not _off(switch.get("madspin")): - run("MadSpin", "decay_events -from_cards") - if ma5: - run("MadAnalysis5 (parton)", "madanalysis5_parton --no_default") - if switch.get("shower") == "Pythia8": - run("Pythia8 shower", "shower --no_default") - if switch.get("detector") == "Delphes": - run("Delphes", "delphes --no_default") - # hadron-level MA5 needs the shower/detector output: only run it if a shower - # was requested (otherwise there is no hadron-level event file to analyse). - if ma5 and showered: - run("MadAnalysis5 (hadron)", "madanalysis5_hadron --no_default") - if switch.get("analysis") == "Rivet": - run("Rivet", "rivet --no_default") - - # Finalize the run exactly as MadEventCmd.do_launch does after the - # shower/detector/analysis tools: store_result() processes the deferred - # "to_store" actions -- in particular it gzips the Pythia8 HepMC to - # _pythia8_events.hepmc.gz (the pythia8_card default is HEPMCoutput:file - # = hepmc.gz). The deferred Rivet job reads exactly that .gz path, so this - # must run before the Rivet/Contur postprocessor below. - try: - cmd.store_result() - except Exception as error: - _report_failure(log, "store_result", error, run_dir) - - if switch.get("analysis") == "Rivet": - # do_rivet only *prepared* the run (wrote Events//run_rivet.sh) and - # deferred execution to the postprocessor -- the rivet_card default has - # run_rivet_later = True, so with --no_default it logged "Skipping Rivet - # for now, passing it to postprocessor" and appended the run to - # cmd.postprocessing_dirs. MadEventCmd.do_launch runs that postprocessor - # at the end of a run; do the same here so the .yoda (and, when - # run_contur = True, the Contur limits) are actually produced. - bar = "=" * 60 - log.info("") - log.info(bar) - log.info(" post-processing step: Rivet/Contur postprocessing") - log.info(bar) - start = time.time() - try: - cmd.postprocessing() - except Exception as error: - _report_failure(log, "Rivet-Contur postprocessing", error, run_dir) - else: - log.info(" -> Rivet/Contur postprocessing done (%.1fs)", - time.time() - start) - - -def _add_time_of_flight(lhe_path, threshold, param_card_path, log): - """Add invariant-lifetime (vtim) information to the LHE events, drawing each - unstable particle's decay length from its width (in mm). Mirrors - common_run_interface.CommonRunCmd.do_add_time_of_flight but reads the widths - from the local param_card instead of the LHE banner.""" - import random - try: - import madgraph.various.lhe_parser as lhe_parser - except ImportError: - import internal.lhe_parser as lhe_parser - from models import check_param_card as param_card_mod - from madgraph.various import misc as _misc - import madgraph.iolibs.files as files - - need_zip = lhe_path.endswith('.gz') - if need_zip: - _misc.gunzip(lhe_path) - lhe_path = lhe_path[:-3] - - param_card = param_card_mod.ParamCard(param_card_path) - cst = 6.58211915e-25 # hbar in GeV s - c = 299792458000 # speed of light in mm/s - log.info('Adding time of flight information on %s', lhe_path) - lhe = lhe_parser.EventFile(lhe_path) - out = open('%s_2vertex.lhe' % lhe_path, 'w') - out.write(lhe.banner) - for event in lhe: - for particle in event: - # default=0 -> particles without a decay entry are treated as stable - width = param_card['decay'].get((abs(particle.pid),), 0.).value - if width: - vtim = c * random.expovariate(width / cst) - if vtim > threshold: - particle.vtim = vtim - out.write(str(event)) - out.write('\n') - out.close() - lhe.close() - files.mv('%s_2vertex.lhe' % lhe_path, lhe_path) - if need_zip: - _misc.gzip(lhe_path) - - -def _lhapdf_config_path(): - """Path to lhapdf-config so systematics can import the python lhapdf - module (required to compute PDF/scale variations).""" - return lhapdf_paths().config - - -def _run_systematics(lhe_path, cfg, log): - """Run systematics.py (scale/PDF variations) on the LHE file, reusing the - exact worker madevent uses (systematics.call_systematics).""" - try: - import madgraph.various.systematics as systematics - except ImportError: - import internal.systematics as systematics - - def fmt(vals): - return ','.join(str(v) for v in vals) - - opts = [] - if cfg.get('systematics_mur'): - opts.append('--mur=%s' % fmt(cfg['systematics_mur'])) - if cfg.get('systematics_muf'): - opts.append('--muf=%s' % fmt(cfg['systematics_muf'])) - if cfg.get('systematics_pdf'): - opts.append('--pdf=%s' % fmt(cfg['systematics_pdf'])) - extra = cfg.get('systematics_str_options', '') - if extra: - opts.extend(extra.split()) - # The mg7 LHE has no block, so systematics cannot read the per-event - # LO reweighting info. For a process with a single alpha_s power (stored in - # SubProcesses/proc_characteristics at output time) it can be reconstructed - # from the events: pass that power as --lo_nqcd. -1 means the QCD power is - # not uniform, so the reconstruction is not applicable. - if not any(o.startswith('--lo_nqcd') for o in opts): - try: - from madgraph.various import banner as _banner_mod - except ImportError: - import internal.banner as _banner_mod - pc_path = os.path.join('SubProcesses', 'proc_characteristics') - if os.path.exists(pc_path): - try: - nqcd = int(_banner_mod.ProcCharacteristic(pc_path)['single_qcd_order']) - except Exception: - nqcd = -1 - if nqcd >= 0: - opts.append('--lo_nqcd=%d' % nqcd) - # tell systematics where to find lhapdf (so it can link the python module) - if not any(o.startswith('--lhapdf_config') for o in opts): - lhapdf_config = _lhapdf_config_path() - if lhapdf_config: - opts.append('--lhapdf_config=%s' % lhapdf_config) - - log.info('Running systematics on %s %s', lhe_path, ' '.join(opts)) - systematics.call_systematics([lhe_path, lhe_path] + opts, - log=lambda x: log.info(str(x))) - - -def run_lhe_postprocessing(process) -> None: - """Run the LHE-level post-processings configured in the run_card - [postprocessing] section on the generated event file (displaced-vertex - time-of-flight and systematics). Only applies when an LHE file exists.""" - lhe_path = _find_event_file(process.run_path) - if lhe_path is None: - return - # process.run_card is a RunCardMG7; the section view exposes .get(key, def) - try: - cfg = process.run_card["postprocessing"] - except Exception: - return - log = logging.getLogger('madevent') - - if cfg.get('systematics') and getattr(process, 'is_decay', False): - # Scale and PDF variations are a beam quantity. A decay has neither, so - # systematics can only fail here ("not supported for pdlabel=none") -- - # and it is not free: MadSpin reruns the launcher for every pool refill. - log.info("decay mode: skipping systematics (no beams to vary)") - elif cfg.get('systematics'): - native = getattr(process, 'systematics', None) - if native is not None and native.weight_count > 0: - log.warning("[postprocessing] systematics is ignored: the scale/PDF " - "weights were already computed by madspace ([systematics] " - "section); set systematics.enable = false to use the legacy " - "systematics.py path instead.") - else: - try: - _run_systematics(lhe_path, cfg, log) - except Exception as error: - _report_failure(log, "systematics computation", error, - os.path.dirname(lhe_path)) - - tof = cfg.get('time_of_flight', -1.0) - try: - if tof is not None and float(tof) >= 0: - _add_time_of_flight(lhe_path, float(tof), process.param_card_path, log) - except Exception as error: - _report_failure(log, "add_time_of_flight", error, - os.path.dirname(lhe_path)) - - -_model_hash_warned = set() -_model_cache = {} - - -def read_stored_model(me_dir=None) -> "str | None": - """The model this process was generated with, as recorded at output time in - ``SubProcesses/model.txt``: a UFO path carrying its restriction (e.g. - ``.../models/sm-no_b_mass``) or a model name -- in either case something - ``import model`` understands. Returns None when the process was written - without that information. - - Warns (once per directory) when the model on disk no longer matches the - hash stored on the second line, i.e. it changed since the output was made. - """ - base = me_dir or os.getcwd() - model_file = os.path.join(base, "SubProcesses", "model.txt") - if not os.path.exists(model_file): - return None - with open(model_file) as f: - lines = f.read().splitlines() - model = lines[0].strip() if lines else "" - stored_hash = lines[1].strip() if len(lines) > 1 else "" - if not model: - logger.warning("SubProcesses/model.txt is empty; the model of this " - "process is unknown.") - return None - - # verify the model on disk still matches the one used at output time. The - # reference may carry a restriction suffix ('sm-no_b_mass'); the hash was - # taken on the UFO directory itself ('sm'). - model_dir = model if os.path.isdir(model) else model.rsplit('-', 1)[0] - if stored_hash and os.path.isdir(model_dir) \ - and model_file not in _model_hash_warned: - current_hash = misc.hash_model_files(model_dir) - if current_hash and current_hash != stored_hash: - _model_hash_warned.add(model_file) - logger.warning( - "The model at %s has changed since this process was generated " - "(hash mismatch); anything computed from it now (the 'auto' " - "widths, the dependent masses/widths of the param_card) may be " - "inconsistent with the matrix element.", model_dir) - return model - - -def _proc_characteristic(me_dir=None) -> dict: - """SubProcesses/proc_characteristics as a plain dict (empty when the file is - missing or unreadable). Plain, because ProcCharacteristic is a ConfigFile, - whose get() is the parameter accessor and takes no default.""" - from madgraph.various import banner as _banner_mod - path = os.path.join(me_dir or os.getcwd(), 'SubProcesses', - 'proc_characteristics') - if os.path.exists(path): - try: - return dict(_banner_mod.ProcCharacteristic(path)) - except Exception as error: - logger.debug("could not read %s: %s", path, error) - return {} - - -def load_process_model(me_dir=None): - """Import the UFO model this process was generated with, as - :func:`models.import_ufo.import_model` builds it -- the object every - model-level consumer of the param_card expects (see - :meth:`MG7Cmd.get_model`). None when the process recorded no model, or when - that model no longer imports. - - Cached per (reference, complex-mass-scheme): importing a model is slow and - the card question asks for it again at every 'update dependent'.""" - ref = read_stored_model(me_dir) - if not ref: - return None - cms = _proc_characteristic(me_dir).get('complex_mass_scheme', False) - key = (ref, bool(cms)) - if key not in _model_cache: - import models.import_ufo as import_ufo - try: - with misc.MuteLogger(['madgraph.model'], [50]): - _model_cache[key] = import_ufo.import_model( - ref, complex_mass_scheme=bool(cms)) - except Exception as error: - # 'update dependent' calls this under a SIGALRM whose handler - # raises a TimeOutError of its own (a class defined inside the - # caller, so it can only be recognised by name). Let it through: - # the caller then says the model took too long to load and how to - # force it, instead of reporting a model that does not import. - if type(error).__name__ == 'TimeOutError': - raise - logger.debug("could not import the model %s: %s", ref, error) - _model_cache[key] = None - return _model_cache[key] - - -def compute_auto_widths(param_card_path=os.path.join("Cards", "param_card.dat")) -> None: - """Fill any width set to ``auto`` in the param_card, using madgraph and the - model stored at output time (``SubProcesses/model.txt``), and write the - result back into the card. A no-op when the card has no ``auto`` width. - - Called before every generation, so a single run and each scan point (whose - param_card is rewritten just before ``run_single``) both get model-computed - widths.""" - try: - with open(param_card_path) as f: - text = f.read() - except OSError: - return - # matches "DECAY auto" (optionally "auto@NLO"), as in - # common_run_interface.static_check_param_card - pdgs = re.findall(r"(?im)^\s*decay\s+([+-]?\d+)\s+auto", text) - if not pdgs: - return - pdgs = list(dict.fromkeys(pdgs)) # de-duplicate, keep order - - model = read_stored_model() - if not model: - logger.warning( - "The param_card requests 'auto' width(s) for %s but the model of " - "this process is unknown; leaving them as-is.", " ".join(pdgs)) - return - - mg5 = str(_MG_ROOT / "bin" / "madgraph") - if not os.path.exists(mg5): - logger.warning("Cannot find madgraph at %s; 'auto' widths not computed.", mg5) - return - - import tempfile - cmds = "import model %s\ncompute_widths %s --path=%s\n" % ( - model, " ".join(pdgs), os.path.abspath(param_card_path)) - with tempfile.NamedTemporaryFile("w", suffix=".mg5", delete=False) as fh: - fh.write(cmds) - cmdfile = fh.name - try: - logger.info("Computing 'auto' width(s) for %s ...", " ".join(pdgs)) - proc = subprocess.run([mg5, cmdfile]) - if proc.returncode != 0: - logger.warning( - "compute_widths returned a non-zero exit code; the param_card " - "may still contain 'auto' entries.") - finally: - try: - os.remove(cmdfile) - except OSError: - pass - - -def _release_channel_generators(process) -> None: - """Drop the per-channel event/weight generators so their intermediate - .npy files are deleted right away. madspace only frees them when the - owning Python objects are collected, and process.event_generator / - process.phasespaces sit in a reference cycle that plain refcounting - never breaks -- only an explicit gc.collect() does, promptly, instead of - leaving it to whenever the cyclic GC next runs on its own.""" - process.event_generator = None - process.phasespaces = None - gc.collect() - - -def run_single(switch=None) -> "MadgraphProcess": - """Run a single generation and return the process (for its result).""" - compute_auto_widths() - process = MadgraphProcess() - process.survey() - process.train_madnis() - process.generate_events() - if process.gridpack_run_mode != "regular": - return process - # run_card-driven LHE post-processing (displaced vertex + systematics) - run_lhe_postprocessing(process) - # run the post-processing tools (Pythia8/Delphes/MadSpin/reweight/analysis) - # selected in the merged question above on the generated events. - if switch: - run_selected_tools(switch, process) - return process - - -def detect_run_scan(run_card_path): - """Return a banner.RunCardIterator if the run_card contains scan:[...] - values, else None.""" - from madgraph.various import banner as banner_mod - from madgraph.various import misc as _misc - with _misc.TMP_variable(banner_mod.RunCard, 'allow_scan', True): - rc = banner_mod.RunCard(run_card_path, consistency=False) - if getattr(rc, 'scan_set', None): - return banner_mod.RunCardIterator(run_card_path) - return None - - -def detect_param_scan(param_card_path): - """Return a ParamCardIterator if the param_card contains scan:[...] values, - else None.""" - if not os.path.exists(param_card_path): - return None - from models import check_param_card as param_card_mod - it = param_card_mod.ParamCardIterator(param_card_path) - for block in it.order: - for param in block: - if isinstance(param.value, str) and param.value.strip().lower().startswith('scan'): - return it - return None - - -def run_scan(iterator, card_path, switch=None) -> None: - """Iterate over all scan points, running a full generation for each and - accumulating the results, then write the scan summary. Works for both the - run_card (RunCardIterator) and the param_card (ParamCardIterator); their - interface (__iter__/write/store_entry/get_next_name/write_summary) is the - same. The scan card is restored afterwards.""" - import tomllib - with open(os.path.join("Cards", "run_card.toml"), "rb") as f: - run_name = tomllib.load(f).get("run", {}).get("run_name", "run") - - from models import check_param_card as param_card_mod - is_param_scan = isinstance(iterator, param_card_mod.ParamCardIterator) - - backup = card_path + ".scan_bak" - shutil.copy(card_path, backup) - try: - for i, point in enumerate(iterator): - point.write(card_path) - logger.info("=== scan point %d ===", i + 1) - process = run_single(switch) - # use the run directory the process actually created, so the - # per-point params.dat written by write_summary has a home - name = os.path.basename(process.run_path) - if is_param_scan: - # pass the (possibly auto-width-updated) card so the summary - # records the model-computed width for each 'auto' entry - iterator.store_entry(name, process.get_result(), - param_card_path=card_path) - else: - iterator.store_entry(name, process.get_result()) - _release_channel_generators(process) - os.makedirs("Events", exist_ok=True) - summary = os.path.join("Events", "scan_%s.txt" % run_name) - iterator.write_summary(summary) - logger.info("scan results written to %s", summary) - finally: - shutil.move(backup, card_path) - - -def run_generation(switch=None) -> None: - """Run the generation, expanding a scan over the run_card or the param_card - when one is present (scanning both simultaneously is not allowed).""" - run_card_path = os.path.join("Cards", "run_card.toml") - param_card_path = os.path.join("Cards", "param_card.dat") - run_iter = detect_run_scan(run_card_path) - param_iter = detect_param_scan(param_card_path) - if run_iter and param_iter: - raise RuntimeError( - "Scanning simultaneously over the run_card and the param_card is " - "not allowed. Please keep the scan:[...] entries in only one card.") - if run_iter: - run_scan(run_iter, run_card_path, switch) - elif param_iter: - run_scan(param_iter, param_card_path, switch) - else: - _release_channel_generators(run_single(switch)) - - -def _lhe_needed_by(switch, card) -> "str | None": - """Why the run needs an LHE event file, or None when it does not: every - post-processing tool run_selected_tools drives and the run_card - [postprocessing] steps read the LHE file. ``card`` is the raw TOML content - of the run_card.""" - tools = selected_tools(switch) if switch else [] - if tools: - return tools[0] - cfg = card.get("postprocessing", {}) - try: - if float(cfg.get("time_of_flight", -1.0)) >= 0: - return "time_of_flight" - except (TypeError, ValueError): - pass - # the legacy systematics.py path only runs when madspace did not compute - # the weights itself (see run_lhe_postprocessing) - if cfg.get("systematics") \ - and not card.get("systematics", {}).get("enable", True): - return "systematics" - return None - - -def force_lhe_output_if_needed(switch) -> None: - """The events are written as npy by default; switch the run_card to the - LHE format when a selected post-processing needs an LHE file. The card is - rewritten rather than refused, so that asking for a shower in the launch - question is enough -- the user does not have to know that it constrains - output_format. - - The card is read as raw TOML and only its output_format line is edited: a - run_card scan ("scan:[...]" values) is not a valid RunCardMG7 and must - survive untouched for run_scan.""" - import tomllib - path = os.path.join("Cards", "run_card.toml") - if not os.path.exists(path): - return - with open(path, "rb") as f: - card = tomllib.load(f) - output_format = card.get("run", {}).get( - "output_format", RunCardMG7()["run"]["output_format"]) - if output_format == "lhe": - return - reason = _lhe_needed_by(switch, card) - if reason is None: - return - with open(path) as f: - text = f.read() - text, count = re.subn(r'(?m)^(\s*output_format\s*=\s*)("[^"]*"|\'[^\']*\'|\S+)', - r'\g<1>"lhe"', text, count=1) - if not count: - text, count = re.subn(r'(?m)^\[run\][^\n]*\n', '\\g<0>output_format = "lhe"\n', - text, count=1) - if not count: - text += '\n[run]\noutput_format = "lhe"\n' - with open(path, "w") as f: - f.write(text) - logging.getLogger("madevent").info( - "output_format set to 'lhe' (required by %s).", reason) - - -def _raise_open_file_limit() -> None: - """Remove the soft limit on the number of open files: it can be quite low - on some systems and the event-combination step opens one file per channel.""" - soft_lim, hard_lim = resource.getrlimit(resource.RLIMIT_NOFILE) - try: - resource.setrlimit(resource.RLIMIT_NOFILE, (hard_lim, hard_lim)) - except (ValueError, OSError) as error: - logger.debug("could not raise the open-file limit: %s", error) - - -def main() -> None: - """``bin/generate_events`` entry point. - - It runs the same :class:`MG7Cmd` command MG5's in-process ``launch`` uses, - so the two entry points cannot drift apart; the only difference is that - nothing here supplies a ``mother``/``inputfile``, so the question falls back - to reading this process's own stdin. - """ - parser = argparse.ArgumentParser() - parser.add_argument("-f", action="store_false", dest="ask_edit_cards") - args = parser.parse_args() - - MG7Cmd().run_cmd("generate_events" if args.ask_edit_cards else "generate_events -f") diff --git a/madgraph/iolibs/template_files/mg7/madgraph_process.py b/madgraph/iolibs/template_files/mg7/madgraph_process.py new file mode 100644 index 0000000000..7321b86638 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/madgraph_process.py @@ -0,0 +1,1966 @@ +"""The run of one process directory: cards, integration, events and output.""" + +import glob +import json +import os +import shutil +import subprocess +import time + +from madgraph.iolibs.template_files.mg7.common import backend_of, format_time, fptype_of, get_start_time, interference_helicities_summed, lhe_helicities, lhe_weight_info, logger, me_parameters, ms, resolve_auto_backend, resolve_seed, resolve_verbosity +from madgraph.various import misc +from madgraph.various.banner import RunCardMG7 +from models.check_param_card import ParamCard +from madgraph.iolibs.template_files.mg7 import hwu_output +from madgraph.iolibs.template_files.mg7 import npy_to_lhe +from madgraph.iolibs.template_files.mg7 import options +from madgraph.iolibs.template_files.mg7 import plots +from madgraph.iolibs.template_files.mg7 import systematics_summary +from madgraph.iolibs.template_files.mg7.bootstrap import INSTALL_DIR as _INSTALL_DIR +from madgraph.iolibs.template_files.mg7.bootstrap import MADSPACE_DIR as _MADSPACE_DIR +from madgraph.iolibs.template_files.mg7.custom_function_loader import load_custom_function +from madgraph.iolibs.template_files.mg7.observables import CutItem, HistItem, build_observable, clean_pids, observable_pids +from madgraph.iolibs.template_files.mg7.options import lhapdf_paths +from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import MadgraphSubprocess, PhaseSpace, build_multi_channel_data + +class MadgraphProcess: + def __init__(self): + self.madnis_trained = False + self.lhe_completer = None + self.load_cards() + self.init_backend() + self.init_event_dir() + self.init_context() + self.init_custom_functions() + self.init_surrogate() + self.init_cuts() + self.init_histograms() + self.init_generator_config() + self.init_beam() + self.init_subprocesses() + + def load_cards(self) -> None: + self.run_card = RunCardMG7(os.path.join("Cards", "run_card.toml")) + # Resolved once so every generator built during this run shares the same + # seed; the concrete value (even if randomly drawn) is recorded in each + # generator's info.json. + self.run_seed = resolve_seed(self.run_card["run"]["seed"]) + self.param_card_path = os.path.join("Cards", "param_card.dat") + self.param_card = ParamCard(self.param_card_path) + with open(os.path.join("SubProcesses", "subprocesses.json")) as f: + self.subprocess_data = json.load(f) + if self.run_card["phasespace"]["merge_subprocesses"]: + with open(os.path.join("SubProcesses", "merged_subprocesses.json")) as f: + self.merged_subprocess_data = json.load(f) + else: + self.merged_subprocess_data = None + + self.init_decay_mode() + self.init_me_frame() + if interference_helicities_summed(self.run_card): + count = sum(1 for meta in self.subprocess_data if meta.get("interference")) + if count: + logger.info( + "interference_helicity = summed: the weight of the %d interference " + "subprocess(es) is the helicity sum, and their helicities are " + "written as 9 in the LHE", count) + + def init_me_frame(self) -> None: + """Resolve the run card's me_frame into the list of external particles + whose momentum sum defines the frame the matrix element is evaluated in. + + Empty (the default) means no boost: the matrix element sees the momenta + in the frame they are generated in. Unlike madevent, which hands its + matrix element momenta that are already in the partonic centre of mass, + madspace generates them in the lab frame, so asking for madevent's own + frame ([1, 2]) does mean an actual boost here. It changes nothing for a + Lorentz invariant matrix element (only the rounding), but it is what + defines the polarisation axes of a polarised one. + + The particle numbers are checked here rather than only in madspace, + which would not see them until the integrands are built -- long after + the matrix-element libraries have been compiled. Every subprocess is + checked against its own particles: an output can mix multiplicities + (p p > z j and p p > z j j, say, for MLM merging), and madspace builds + one matrix element per subprocess, each refusing a particle it does not + have. That is stricter than madevent, which ignores the particles a + process does not have and so changes the frame from one multiplicity + to the next without a word. + """ + me_frame = list(self.run_card["run"]["me_frame"]) + pdgs = [clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) + for meta in self.subprocess_data] + counts = sorted({len(pdg_list) for pdg_list in pdgs}) + for meta, pdg_list in zip(self.subprocess_data, pdgs): + for index in me_frame: + if 1 <= index <= len(pdg_list): + continue + process = "%s > %s" % ( + " ".join(str(pid) for pid in meta["incoming"]), + " ".join(str(pid) for pid in meta["outgoing"]), + ) + mixed = ( + f"; this output mixes subprocesses with " + f"{', '.join(str(count) for count in counts)} external " + "particles, and me_frame has to name particles every one " + "of them has" + if len(counts) > 1 else "" + ) + raise ValueError( + f"me_frame particle {index} out of range: subprocess " + f"{process} has {len(pdg_list)} external particles{mixed}" + ) + if len(me_frame) == 1: + # A massless particle has no rest frame: its momentum squares to + # zero, so the boost divides by it and every momentum comes back as + # a NaN, silently. Name the mistake instead. + massless = [pdg_list[me_frame[0] - 1] for pdg_list in pdgs + if self.get_mass(pdg_list[me_frame[0] - 1]) == 0.] + if massless: + raise ValueError( + f"me_frame = {me_frame} asks for the rest frame of particle " + f"{me_frame[0]}, which is massless (pdg " + f"{sorted(set(massless))}) and has none. Name a massive " + "particle, or several particles whose sum is massive." + ) + self.me_frame = me_frame + self.incoming_count = 1 if self.is_decay else 2 + + def init_decay_mode(self) -> None: + """Decide whether this directory is a decay (1 -> n) or a collision. + + Read off the exported process rather than the run card: the two cannot + then disagree. MadSpin generates its decay matrix elements this way. + """ + incoming_counts = { + len(clean_pids(meta["incoming"])) for meta in self.subprocess_data + } + if incoming_counts - {1, 2}: + raise ValueError( + f"processes with {sorted(incoming_counts)} incoming particles " + "are not supported" + ) + if len(incoming_counts) > 1: + raise ValueError( + "cannot mix decays and collisions in one output directory" + ) + self.is_decay = incoming_counts == {1} + if not self.is_decay: + self.decaying_mass = None + return + + masses = { + self.get_mass(clean_pids(meta["incoming"])[0]) + for meta in self.subprocess_data + } + if len(masses) > 1: + raise ValueError( + f"decaying particles have different masses: {sorted(masses)}" + ) + self.decaying_mass = masses.pop() + if self.decaying_mass <= 0.0: + raise ValueError("the decaying particle must have a non-zero mass") + self.drop_closed_channels() + + def drop_closed_channels(self) -> None: + """Remove subprocesses the decaying particle is too light to produce. + + A multiparticle decay definition enumerates every vertex the model + allows, closed ones included: "t > b w+, w+ > all all" yields + t > b t b~ (and b W+ Z, b W+ h, ...), which need more mass than the top + has. Their partial width is exactly zero, but the phase-space mapping + has no physical point to hand back -- the invariant's lower bound ends + up above its upper bound -- so it produces NaN momenta and poisons the + whole integral. Drop them here instead. + """ + kept, dropped = [], [] + for meta in self.subprocess_data: + total = sum( + self.get_mass(pid) for pid in clean_pids(meta["outgoing"]) + ) + if total < self.decaying_mass: + kept.append(meta) + else: + dropped.append((meta["outgoing"], total)) + if dropped and self.merged_subprocess_data is not None: + # merged_subprocesses.json indexes the *unfiltered* subprocess + # list, so dropping entries here would silently shift every index + # it holds. Refuse rather than mis-map. + raise ValueError( + "merge_subprocesses is not supported for a decay directory " + "with kinematically closed channels" + ) + if dropped: + for outgoing, total in dropped: + logger.info( + "skipping closed decay channel -> %s (needs %.4g GeV, " + "the decaying particle has %.4g GeV)", + outgoing, total, self.decaying_mass, + ) + if not kept: + raise ValueError( + "every decay channel is kinematically closed: the decaying " + f"particle's mass is {self.decaying_mass} GeV" + ) + self.subprocess_data = kept + + def init_backend(self) -> None: + ms.set_simd_vector_size(self.run_card["run"]["simd_vector_size"]) + + def init_event_dir(self) -> None: + output_dir = self.run_card["gridpack"]["output_dir"] + temp_dir = self.run_card["gridpack"]["temp_output_dir"] + if output_dir: + self.run_path = os.path.abspath(os.path.expanduser(output_dir)) + os.makedirs(self.run_path, exist_ok=True) + else: + run_name = self.run_card["run"]["run_name"] + os.makedirs("Events", exist_ok=True) + run_dir_prefix = os.path.join("Events", f"{run_name}_") + existing_run_dirs = glob.glob(f"{run_dir_prefix}*") + run_index = 1 + for run_dir in existing_run_dirs: + run_index_str = run_dir[len(run_dir_prefix):] + if run_index_str.isnumeric(): + run_index = max(run_index, int(run_index_str) + 1) + while True: + try: + self.run_path = f"{run_dir_prefix}{run_index:02d}" + os.mkdir(self.run_path) + break + except FileExistsError: + run_index += 1 + # temporary npy files of the channel generators + if temp_dir: + self.temp_path = os.path.abspath(os.path.expanduser(temp_dir)) + os.makedirs(self.temp_path, exist_ok=True) + else: + self.temp_path = self.run_path + # Absolute on purpose. StatusFile is a C++ object that finishes its + # write from its destructor (rename info.json.tmp -> info.json), and + # that destructor runs whenever Python gets round to collecting the + # process object. When the launcher runs in MG5's process, an + # interrupted run is collected *after* the launch command has restored + # the working directory, and a relative path would then resolve + # somewhere else -- the rename throws a C++ filesystem_error out of a + # destructor, which is an immediate abort. + self.status_file = ms.StatusFile( + os.path.abspath(os.path.join(self.run_path, "info.json"))) + + def init_context(self) -> None: + device_names = self.run_card["run"]["device"] + self.contexts = [] + self.device_types = [] + self.devices = [] + self.pool_sizes = [] + for device_name in device_names: + if ":" in device_name: + device_type, device_index_str = device_name.split(":") + device_index = int(device_index_str) + else: + device_type = device_name + device_index = 0 + self.device_types.append(device_type) + if device_type == "cuda": + device = ms.cuda_device(device_index) + pool_size = self.run_card["run"]["gpu_thread_pool_size"] + elif device_type == "hip": + device = ms.hip_device(device_index) + pool_size = self.run_card["run"]["gpu_thread_pool_size"] + else: + device = ms.cpu_device() + pool_size = self.run_card["run"]["cpu_thread_pool_size"] + self.devices.append(device) + self.pool_sizes.append(pool_size) + self.contexts.append(ms.Context(device=device, thread_count=pool_size)) + + CUSTOM_FUNCTION_KINDS = ("cut", "scale", "pdf", "bias") + + def init_custom_functions(self) -> None: + """Load the functions of [custom_functions], building the libraries + first. They are loaded before the matrix elements, so that they have the + same index in every context and in a gridpack made from this run.""" + self.custom_functions = {} + self.custom_function_paths = [] + for kind in self.CUSTOM_FUNCTION_KINDS: + path = self.run_card["custom_functions"][kind] + if not path: + continue + build_dir = f"custom_{kind}" + if path.endswith(".so") and os.path.isfile(os.path.join(build_dir, "Makefile")): + self.build_custom_function(build_dir) + if not os.path.isfile(path): + raise RuntimeError( + f"[custom_functions] {kind} = {path!r}: file not found. " + "Create a .so library or a .py file with bin/init_custom_func." + ) + apis = load_custom_function(ms, self.contexts, self.device_types, kind, path) + if len({api.index() for api in apis}) != 1: + raise RuntimeError(f"custom {kind} function loaded at different indices") + self.custom_functions[kind] = apis[0] + self.custom_function_paths.append([kind, path]) + logger.info("Using the custom %s function %s", kind, path) + + def init_surrogate(self) -> None: + """Check the [surrogate] settings. The surrogates are loaded per + subprocess, by MadgraphSubprocess.""" + settings = self.run_card["surrogate"] + self.surrogate = None + if not settings["function"]: + return + path = settings["function"] + if path.endswith(".py") and not os.path.isfile(path): + raise RuntimeError(f"[surrogate] function = {path!r}: file not found") + if self.run_card["phasespace"]["sde_strategy"] != "denominators": + raise RuntimeError( + "[surrogate] needs [phasespace] sde_strategy = \"denominators\", " + "as the surrogate does not return channel weights" + ) + if self.run_card["run"]["dummy_matrix_element"]: + raise RuntimeError("[surrogate] cannot be used with dummy_matrix_element") + if self.run_card["gridpack"]["save_gridpack"]: + raise RuntimeError("[surrogate] does not support gridpacks yet") + self.surrogate = dict(settings) + if not self.surrogate["param_card"]: + self.surrogate["param_card"] = self.param_card_path + logger.info( + "Using the surrogate %s, %s with the full matrix element", + path, + "unweighting" if settings["mode"] == "unweight" else "reweighting", + ) + + def weighting_groups(self, subprocess_channels: list[list[int]]) -> list: + """One ms.WeightingGroup per subprocess with a [surrogate]: the second + pass with the full matrix element over its channels, given by their + indices in the event generator.""" + if self.surrogate is None: + return [] + gen_args = self.run_card["generation"] + return [ + ms.WeightingGroup( + channel_indices, + [ms.WeightingStep( + subproc.step_matrix_elements, + mode=self.surrogate["mode"], + max_overweight_truncation=gen_args["max_overweight_truncation"], + freeze_max_weight_after=gen_args["freeze_max_weight_after"], + )], + ) + for subproc, channel_indices in zip(self.subprocesses, subprocess_channels) + if channel_indices + ] + + @staticmethod + def build_custom_function(build_dir: str) -> None: + """Run make in a custom_ directory, logging to compile.log.""" + log_path = os.path.join(build_dir, "compile.log") + with open(log_path, "w") as log: + returncode = subprocess.call( + ["make"], cwd=build_dir, stdout=log, stderr=subprocess.STDOUT + ) + if returncode != 0: + raise RuntimeError( + f"Compilation of {build_dir} failed with exit status {returncode}; " + f"see {log_path} for details" + ) + + def custom_function(self, kind: str, particle_count: int, pid_options: list): + """The ms.CustomFunction of [custom_functions] , or None.""" + api = self.custom_functions.get(kind) + if api is None: + return None + function_type = getattr(ms.CustomFunction, kind) + return ms.CustomFunction(api, function_type, particle_count, pid_options) + + def parse_observable(self, name: str, order_observable: str) -> dict: + parts = name.split("-") + sum_momenta = False + sum_observable = False + ordered = False + multiparticles = self.run_card["multiparticles"] + + if len(parts) == 0: + raise ValueError("Invalid observable name") + elif len(parts) == 1: + # event-level observables + obs_name = parts[0] + select_pids = [] + else: + if parts[-1] == "sum": + sum_observable = True + obs_name = parts[-2] + selection = parts[:-2] + elif parts[-2] == "sum": + sum_momenta = True + obs_name = parts[-1] + selection = parts[:-2] + else: + obs_name = parts[-1] + selection = parts[:-1] + select_pids = [] + order_indices = [] + for mp_name in selection: + mp_parts = mp_name.split("_") + if mp_parts[-1].isnumeric(): + order_indices.append(int(mp_parts[-1])) + select_pids.append(multiparticles["_".join(mp_parts[:-1])]) + ordered = True + else: + order_indices.append(0) + select_pids.append(multiparticles[mp_name]) + + return dict( + observable=obs_name, + select_pids=select_pids, + sum_momenta=sum_momenta, + sum_observable=sum_observable, + order_observable=order_observable if ordered else None, + order_indices=order_indices if ordered else [], + ignore_incoming=True, + name=name, + ) + + def init_cuts(self) -> None: + inf = float("inf") + order_observable = self.run_card["cuts"].get("order_by", "pt") + self.cut_data = [ + CutItem( + observable_kwargs=self.parse_observable(key, order_observable), + min=values.get("min", -inf), + max=values.get("max", inf), + mode=values.get("mode", "all"), + ) + for key, values in self.run_card["cuts"].items() + if key != "order_by" + ] + + # [histograms] key that means "the distribution of the event weight", + # normalised to the unit weight sigma_abs, rather than an observable of the + # momenta: an unweighted sample is a spike at 1 (and at -1 for its negative + # weights) and a partially unweighted one shows its spread. + weight_histogram_key = "weight" + + def init_histograms(self) -> None: + order_observable = self.run_card["histograms"].get("order_by", "pt") + self.hist_data = [] + for key, values in self.run_card["histograms"].items(): + if key == "order_by": + continue + from_weight = key == self.weight_histogram_key + self.hist_data.append(HistItem( + observable_kwargs=( + None if from_weight + else self.parse_observable(key, order_observable)), + min=values["min"], + max=values["max"], + bin_count=values["bin_count"], + name=key, + from_weight=from_weight, + )) + + def ensure_pdf_set(self, pdf_set: str) -> misc.LhapdfPaths: + """Make sure the requested LHAPDF set is available, downloading it into + the first writable PDF directory if needed, and return the (possibly + updated) LHAPDF resolution. Nothing is downloaded when the set is + already present.""" + paths = lhapdf_paths() + if paths.find_set(pdf_set): + return paths + if not paths.config: + logger.debug("no usable lhapdf-config: cannot download %s", pdf_set) + return paths + if not paths.download_path: + logger.debug("no writable PDF directory: cannot download %s", pdf_set) + return paths + try: + from madgraph.interface.common_run_interface import CommonRunCmd + os.makedirs(paths.download_path, exist_ok=True) + logger.info("PDF set %s not found; downloading into %s", + pdf_set, paths.download_path) + CommonRunCmd.install_lhapdf_pdfset_static( + paths.config, paths.download_path, pdf_set) + except Exception as err: + logger.warning("Could not download PDF set %s: %s", pdf_set, err) + return paths + options._LHAPDF = paths.with_data_path(paths.download_path) + return options._LHAPDF + + def _no_pdf_message(self, pdf_set: str) -> str: + """Explain a missing PDF set in terms of the settings that fix it.""" + searched = os.pathsep.join(self.lhapdf.data_paths) or \ + "(no PDF set directory could be located)" + if self.lhapdf.config: + found = "using lhapdf-config: %s" % self.lhapdf.config + else: + found = ("no usable lhapdf-config was found, so the set could not " + "be downloaded automatically") + return ( + "LHAPDF set %r, requested by [beam] pdf in Cards/run_card.toml, " + "was not found.\n" + " searched: %s\n" + " %s\n" + "Fix this in one of the following ways:\n" + " * install LHAPDF for this MadGraph: MG7> install lhapdf6\n" + " * point MadGraph at an existing one:\n" + " MG7> set lhapdf /path/to/lhapdf-config\n" + " MG7> save options\n" + " * or set it for this directory only, by adding\n" + " lhapdf = /path/to/lhapdf-config\n" + " to %s\n" + " * or set $LHAPDF_DATA_PATH to a directory containing %r." + % (pdf_set, searched, found, + os.path.join("Cards", "me5_configuration.txt"), pdf_set)) + + def init_beam(self) -> None: + beam_args = self.run_card["beam"] + + # built lazily by build_systematics (needs the subprocess data) + self.systematics = None + self.systematics_data = None + self.systematics_context = None + self.event_histograms = None + self.event_histograms_context = None + + if self.is_decay: + # No beams: the total energy is the decaying particle's mass, and + # "leptonic" is what the mappings call "no parton luminosity". + self.e_cm = self.decaying_mass + self.leptonic = True + else: + self.e_cm = beam_args["e_cm"] + self.leptonic = beam_args["leptonic"] + + dynamical_scales = { + "transverse_energy": ms.EnergyScale.transverse_energy, + "transverse_mass": ms.EnergyScale.transverse_mass, + "half_transverse_mass": ms.EnergyScale.half_transverse_mass, + "partonic_energy": ms.EnergyScale.partonic_energy, + } + if beam_args["dynamical_scale_choice"] in dynamical_scales: + dynamical_scale_type = dynamical_scales[beam_args["dynamical_scale_choice"]] + else: + raise ValueError("Unknown dynamical scale choice") + self.scale_kwargs = dict( + dynamical_scale_type=dynamical_scale_type, + ren_scale_fixed=beam_args["fixed_ren_scale"], + fact_scale_fixed=beam_args["fixed_fact_scale"], + ren_scale=beam_args["ren_scale"], + fact_scale1=beam_args["fact_scale1"], + fact_scale2=beam_args["fact_scale2"], + scale_factor=float(beam_args.get("scale_factor", 1.)), + ) + if self.is_decay and self.custom_functions.pop("scale", None) is not None: + logger.warning("custom scale ignored: a decay uses the decaying mass") + if self.is_decay and self.custom_functions.pop("pdf", None) is not None: + logger.warning("custom pdf ignored: a decay has no parton densities") + if self.is_decay: + # One scale is available for a decay -- the decaying mass -- so use + # it, fixed, whatever the card asks for. + self.scale_kwargs.update( + ren_scale_fixed=True, + fact_scale_fixed=True, + ren_scale=self.decaying_mass, + fact_scale1=self.decaying_mass, + fact_scale2=self.decaying_mass, + scale_factor=1., + ) + self.pdf_grid = None + self.lhapdf = None + self.pdf_dir = None + self.fixed_alphas_info = self.write_fixed_alphas_info() + self.alphas_grid = ms.AlphaSGrid(self.fixed_alphas_info) + for context in self.contexts: + self.alphas_grid.initialize_globals(context) + self.running_coupling = ms.RunningCoupling(self.alphas_grid) + return + + pdf_set = beam_args["pdf"] + self.pdf_set = pdf_set + self.fixed_alphas_info = None + self.lhapdf = self.ensure_pdf_set(pdf_set) + self.pdf_dir = self.lhapdf.find_set(pdf_set) + if self.pdf_dir is None: + raise RuntimeError(self._no_pdf_message(pdf_set)) + self.alphas_grid = ms.AlphaSGrid(os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}.info")) + if "pdf" in self.custom_functions: + # alpha_s still comes from the PDF set of the run card + self.pdf_grid = ms.CustomPdf(self.custom_functions["pdf"]) + else: + self.pdf_grid = ms.PdfGrid( + os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}_0000.dat") + ) + for context in self.contexts: + if not isinstance(self.pdf_grid, ms.CustomPdf): + self.pdf_grid.initialize_globals(context) + self.alphas_grid.initialize_globals(context) + self.running_coupling = ms.RunningCoupling(self.alphas_grid) + + # ------------------------------------------------------------------ + # scale / PDF systematics + # ------------------------------------------------------------------ + def systematics_enabled(self) -> bool: + """The native scale/PDF weights are requested ([systematics] enable, or + the legacy [generation] systematics alias).""" + try: + enabled = bool(self.run_card["systematics"]["enable"]) + except Exception: + enabled = False + if not enabled: + try: + enabled = bool(self.run_card["generation"]["systematics"]) + except Exception: + enabled = False + if enabled and "pdf" in getattr(self, "custom_functions", {}): + # the variations re-evaluate the PDF set, which a custom PDF replaces + if not getattr(self, "_warned_custom_pdf_systematics", False): + logger.warning("systematics are disabled with a custom pdf") + self._warned_custom_pdf_systematics = True + return False + return enabled + + @staticmethod + def pdf_set_info(info_path: str) -> dict: + """Read SetIndex, NumMembers, ErrorType and SetDesc from an LHAPDF + .info file (missing keys get neutral defaults).""" + info = {"SetIndex": 0, "NumMembers": 1, "ErrorType": "", "SetDesc": ""} + try: + with open(info_path) as f: + for line in f: + key, sep, value = line.partition(":") + key = key.strip() + if not sep or key not in info: + continue + value = value.strip().strip('"').strip("'") + if key in ("SetIndex", "NumMembers"): + info[key] = int(value) + else: + info[key] = value + except OSError as err: + logger.warning("could not read %s: %s", info_path, err) + return info + + def resolve_pdf_set_name(self, entry: str) -> tuple[str, int | None]: + """Turn a [systematics] pdf entry ("", "", "@", + "@") into (set name, member or None) using the LHAPDF + pdfsets.index when an id is given.""" + member = None + if "@" in entry: + entry, member_str = entry.split("@", 1) + member = int(member_str) + entry = entry.strip() + if not entry.isdigit(): + return entry, member + lhaid = int(entry) + index_path = os.path.join(self.pdf_dir or "", "pdfsets.index") + sets = [] + try: + with open(index_path) as f: + for line in f: + parts = line.split() + if len(parts) >= 2 and parts[0].isdigit(): + sets.append((int(parts[0]), parts[1])) + except OSError: + raise ValueError( + "cannot resolve LHAPDF id %s: no %s" % (lhaid, index_path)) + for set_id, name in sets: + if set_id == lhaid: + return name, member + # an id inside a set's member range + sets.sort() + for (set_id, name), (next_id, _) in zip(sets, sets[1:] + [(None, None)]): + if set_id < lhaid and (next_id is None or lhaid < next_id): + return name, lhaid - set_id if member is None else member + raise ValueError("cannot resolve LHAPDF id %s" % lhaid) + + def resolve_pdf_variations(self) -> list: + """Expand the [systematics] pdf entries into ms.PdfMemberSpec objects + (downloading the sets when needed). The nominal member itself is never + listed: the calculator adds it as the central weight of its group.""" + entries = self.run_card["systematics"]["pdf"] + if isinstance(entries, str): + entries = [entries] + specs = [] + seen = set() + for raw in entries: + entry = str(raw).strip() + if not entry or entry.lower() in ("none", "no", "false"): + continue + if entry.lower() == "central": + continue + if entry.lower() == "errorset": + set_name, member = self.pdf_set, None + else: + set_name, member = self.resolve_pdf_set_name(entry) + variation_lhapdf = self.ensure_pdf_set(set_name) + variation_dir = variation_lhapdf.find_set(set_name) + if variation_dir is None: + logger.warning("PDF set %s not available; its variation is skipped", set_name) + continue + info_path = os.path.join(variation_dir, set_name, f"{set_name}.info") + if not os.path.exists(info_path): + logger.warning("PDF set %s not available; its variation is skipped", set_name) + continue + info = self.pdf_set_info(info_path) + is_nominal = set_name == self.pdf_set + if member is not None: + members = [member] + else: + members = list(range(info["NumMembers"])) + if is_nominal: + members = members[1:] + if not members: + logger.warning( + "PDF set %s has a single member: no PDF uncertainty can be " + "computed from it. Use a set with error members in " + "[systematics] pdf (e.g. a PDF4LHC or NNPDF error set).", + set_name) + for m in members: + if is_nominal and m == 0: + continue + if (set_name, m) in seen: + continue + seen.add((set_name, m)) + grid_file = os.path.join(variation_dir, set_name, f"{set_name}_{m:04d}.dat") + if not os.path.exists(grid_file): + logger.warning("PDF member file %s not found; skipped", grid_file) + continue + specs.append(ms.PdfMemberSpec( + set_name=set_name, set_lhaid=info["SetIndex"], member=m, + grid_file=grid_file, info_file=info_path, + error_type=info["ErrorType"], description=info["SetDesc"], + )) + return specs + + def build_systematics_args(self) -> list: + """One ms.SubprocessSystArgs per (unmerged) subprocess: the alpha_s power + of |M|^2 and the beam parton ids of every flavor index, indexed like the + LHECompleter (subprocess_index / flavor_index of the combined events).""" + args = [] + for meta in self.subprocess_data: + beam_pdgs = [ + [int(flavor["options"][0][0]), int(flavor["options"][0][1])] + for flavor in meta["flavors"] + ] + args.append(ms.SubprocessSystArgs( + qcd_power=int(meta.get("qcd_power", -1)), beam_pdgs=beam_pdgs)) + return args + + def build_systematics(self): + """Create the ms.SystematicsCalculator from the [systematics] section, or + None when disabled. Also keeps the JSON description needed by gridpacks.""" + if not self.systematics_enabled(): + self.systematics = None + return None + syst = self.run_card["systematics"] + config = ms.SystematicsConfig() + config.mur = [float(v) for v in syst["mur"]] + config.muf = [float(v) for v in syst["muf"]] + config.together = bool(syst["together"]) + config.dyn_scales = self.resolve_dynamical_scales() + # alternative dynamical scales get the same factor as the nominal one + config.scale_factor = float(self.run_card["beam"].get("scale_factor", 1.)) + config.write_inputs = bool(syst["write_inputs"]) + config.has_pdf = not self.leptonic + # Without parton luminosity -- a decay, or leptonic beams -- there is no + # nominal PDF *set* to describe and no PDF variation to compute, so + # has_pdf gates the nominal-set metadata and the member list below. + # The alpha_s grid is a separate matter: the mu_R variations always need + # one, so its .info file has to be recorded whatever has_pdf says. + # init_beam() skips the PDF lookup for a decay only -- leptonic beams + # still resolve a set and take alpha_s from it -- so the path follows + # self.pdf_dir, and falls back to the constant grid written for a decay. + if self.pdf_dir is not None: + info_path = os.path.join(self.pdf_dir, self.pdf_set, f"{self.pdf_set}.info") + else: + info_path = self.fixed_alphas_info + if config.has_pdf: + info = self.pdf_set_info(info_path) + config.nominal_set_name = self.pdf_set + config.nominal_lhaid = info["SetIndex"] + config.nominal_error_type = info["ErrorType"] + config.nominal_description = info["SetDesc"] + config.pdf_members = self.resolve_pdf_variations() + else: + config.pdf_members = [] + args = self.build_systematics_args() + # the PDFs, alpha_s and (for mixed-order subprocesses) the matrix + # elements are evaluated with the batched madspace functions on this + # CPU context + self.systematics_context = ms.Context(device=ms.cpu_device(), thread_count=1) + matrix_elements, flavor_remap, me_backend = self.build_systematics_matrix_elements( + args, self.systematics_context) + self.systematics = ms.SystematicsCalculator( + config, args, self.pdf_grid, self.alphas_grid, + self.systematics_context, matrix_elements, flavor_remap) + for warning in self.systematics.warnings: + logger.warning("systematics: %s", warning) + self.systematics_data = { + "config": json.loads(config.to_json()), + "subproc_args": [json.loads(a.to_json()) for a in args], + "nominal_grid_file": os.path.join( + self.pdf_dir, self.pdf_set, f"{self.pdf_set}_0000.dat") + if config.has_pdf else None, + "nominal_info_file": info_path, + # matrix elements re-evaluated for the mixed-order subprocesses + "me_paths": [meta["me_path"] for meta in self.subprocess_data], + "me_backend": me_backend, + "flavor_remap": flavor_remap, + "me_frame": self.me_frame, + "incoming_count": self.incoming_count, + } + return self.systematics + + _DYN_SCALE_CODES = { + "transverse_energy": 1, "transverse_mass": 2, + "half_transverse_mass": 3, "partonic_energy": 4, + } + + def resolve_dynamical_scales(self) -> list: + """[systematics] dynamical_scale entries (names or systematics.py codes + 1-4) as codes, without the choice the events were generated with.""" + try: + entries = self.run_card["systematics"]["dynamical_scale"] + except Exception: + return [] + if isinstance(entries, str): + entries = [entries] + beam = self.run_card["beam"] + generated = None + if not (beam["fixed_ren_scale"] and beam["fixed_fact_scale"]): + generated = self._DYN_SCALE_CODES.get(beam["dynamical_scale_choice"]) + codes = [] + for raw in entries: + entry = str(raw).strip().lower() + if not entry: + continue + if entry.isdigit(): + code = int(entry) + elif entry in self._DYN_SCALE_CODES: + code = self._DYN_SCALE_CODES[entry] + else: + raise ValueError("unknown dynamical scale choice %r in [systematics] " + "dynamical_scale" % raw) + if code not in (1, 2, 3, 4): + raise ValueError("dynamical scale code %s must be 1-4" % code) + if code == generated or code in codes: + continue + codes.append(code) + return codes + + def build_systematics_matrix_elements(self, args, context): + """For the subprocesses whose |M|^2 mixes several powers of alpha_s + (qcd_power = -1), load the matrix element into the systematics CPU + context so that the renormalisation scale variations can re-evaluate it + at the varied alpha_s. Returns ([MatrixElement or None per subprocess], + flavor remap, backend name); empty when nothing is needed or no CPU + library exists.""" + need = [i for i, a in enumerate(args) if a.qcd_power < 0] + empty = ([], [], None) + if not need or self.run_card["run"]["dummy_matrix_element"]: + return empty + cpu_backends = [b for b in self.backends if not b.startswith(("cuda", "hip"))] + if not cpu_backends: + logger.warning("systematics: no CPU matrix-element library available, the " + "renormalisation scale variations of the mixed-order " + "subprocesses cannot be computed") + return empty + backend = cpu_backends[0] + matrix_elements = [] + flavor_remap = [] + for i, meta in enumerate(self.subprocess_data): + flavor_remap.append([int(f["index"]) for f in meta["flavors"]]) + if i not in need: + matrix_elements.append(None) + continue + api = context.load_matrix_element( + meta["me_path"].format(device=backend), self.param_card_path, + me_parameters(self.run_card)) + matrix_elements.append(ms.MatrixElement( + api, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in], + [ms.MatrixElement.matrix_element_out], + False, + self.me_frame, + self.incoming_count, + )) + logger.info("systematics: re-evaluating the matrix element of %d mixed-order " + "subprocess(es) for the renormalisation scale variations", len(need)) + return matrix_elements, flavor_remap, backend + + def build_event_histograms(self): + """ms.EventHistograms for the [histograms] observables, filled with the + written events and all their variation weights (info.json + "event_histograms"); None without histograms.""" + if not self.hist_data \ + or not self.run_card["run"]["postprocessing_histograms"]: + self.event_histograms = None + return None + context = ms.Context(device=ms.cpu_device(), thread_count=1) + specs = [ + ms.EventHistogramSpec( + name=item.name, min=item.min, max=item.max, + bin_count=item.bin_count, from_weight=item.from_weight) + for item in self.hist_data + ] + from_momenta = [item for item in self.hist_data if not item.from_weight] + observables = [] + for meta in self.subprocess_data: + if not from_momenta: + # nothing to evaluate on the momenta (only the weight is + # histogrammed): no observable runtime to build + observables.append(None) + continue + all_pids = clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) + values = ms.ObservableValues([ + build_observable( + observable_pids(meta, item.observable_kwargs), + len(meta["incoming"]), item.observable_kwargs + ) + for item in from_momenta + ]) + observables.append(ms.SubprocessObservables(values, len(all_pids))) + # the weight histograms are drawn in units of the unit weight sigma_abs + # (the cross section of a sample without negative weights) + self.event_histograms_context = context + self.event_histograms = ms.EventHistograms( + context, specs, observables, + reference_weight=self._mean_event_weight()) + return self.event_histograms + + def _mean_event_weight(self) -> float: + """The unit weight of the unweighted events, sigma_abs, used as the + unit of the weight histograms: the cross section of a positive sample, + and for one with negative weights (an interference) what puts them at + +-1 rather than dividing by a signed, possibly small, total. 0 (the raw + weight) when it is not available.""" + try: + return float(self.event_generator.status().mean_abs) + except Exception as error: + logger.debug("no cross section for the weight histograms: %s", error) + return 0. + + def log_systematics_summary(self) -> None: + """Print the scale/PDF uncertainties on the total cross section (the + per-variation cross sections are in events.weights.json / info.json).""" + if self.systematics is None or self.systematics.weight_count == 0: + return + summary = json.loads(self.systematics.summary()) + # the percentages come from systematics_summary, which get_result() + # reads too: the box and the scan row must not disagree, so neither + # computes them itself + xsec = systematics_summary.nominal_cross_section(summary) + if xsec is None: + return + def format_variation(up, down): + # right-justify the signed numbers (not the sign alone) so the % + # values line up across rows, with extra spacing between the two + return f"{'+%.3g' % up:>6}% {'-%.3g' % down:>6}%" + + scale_count = sum(1 for v in self.systematics.variations if v.is_scale) + member_count = len(self.systematics.members) + rows = [("Variations per event:", + f"{self.systematics.weight_count} ({scale_count} scale, {member_count} PDF members)")] + for pdf in summary.get("pdf", []): + rows.append(("PDF set:", f"{pdf['pdf_set']}, {pdf['error_type']}")) + rows.append(("Original cross-section:", f"{xsec} pb")) + scale = systematics_summary.scale_percentages(summary) + if scale is not None: + rows.append(("Scale variation:", format_variation(*scale))) + for _entry, up, down in systematics_summary.pdf_percentages(summary): + rows.append(("PDF variation:", format_variation(up, down))) + if self.event_generator_config.verbosity == ms.Verbosity.pretty: + box = ms.PrettyBox("Systematics", len(rows), [24, 0]) + box.set_column(0, [label for label, _ in rows]) + box.set_column(1, [value for _, value in rows]) + box.print_first() + else: + for label, value in rows: + logger.info("systematics, %s %s", label, value) + + def write_fixed_alphas_info(self) -> str: + """Write a minimal LHAPDF ``.info`` holding a constant alpha_s. + + A decay has no beams, so there is no PDF set to take alpha_s from -- + and demanding one (possibly downloading it) just to evaluate a coupling + would be absurd. The renormalisation scale of a decay is fixed at the + decaying mass anyway, so a constant alpha_s is the right answer, not an + approximation: take it from the param card, exactly as the Fortran + decay matrix elements do. + """ + # SMINPUTS entry 3 is alpha_s(m_Z), the same value the Fortran + # parameter setup feeds to G. + alpha_s = float(self.param_card.get_value("sminputs", 3)) + path = os.path.join(self.run_path, "fixed_alphas.info") + # AlphaSGrid interpolates in log(q^2) across the grid, so give it a + # comfortable number of nodes rather than the bare minimum of three. + q_values = [10 ** (i / 4.0) for i in range(-4, 21)] + with open(path, "w") as f: + f.write("SetDesc: constant alpha_s for a decay (no beams)\n") + f.write("AlphaS_Qs: [%s]\n" % ", ".join(f"{q:g}" for q in q_values)) + f.write( + "AlphaS_Vals: [%s]\n" + % ", ".join(f"{alpha_s:g}" for _ in q_values) + ) + logger.info( + "decay mode: fixed alpha_s = %g at mu = %g GeV", alpha_s, self.e_cm + ) + return path + + def needs_systematics_matrix_elements(self) -> bool: + """The systematics will re-evaluate a matrix element per scale point. + + True when the native weights are on and some subprocess mixes alpha_s + powers (qcd_power < 0), i.e. when build_systematics_matrix_elements() + loads a matrix element. Conservative: it does not check that a CPU + backend exists, so it can be True where that call ends up loading + nothing. + """ + if not self.systematics_enabled(): + return False + if self.run_card["run"]["dummy_matrix_element"]: + return False + return any(int(meta.get("qcd_power", -1)) < 0 + for meta in self.subprocess_data) + + def init_generator_config(self) -> None: + run_args = self.run_card["run"] + gen_args = self.run_card["generation"] + vegas_args = self.run_card["vegas"] + cfg = ms.GeneratorConfig() + cfg.target_count = gen_args["events"] + cfg.vegas_damping = vegas_args["damping"] + cfg.max_overweight_truncation = gen_args["max_overweight_truncation"] + cfg.freeze_max_weight_after = gen_args["freeze_max_weight_after"] + if self.surrogate is not None and self.surrogate["mode"] == "unweight": + # [generation] applies to the second pass, which gives the events + cfg.max_overweight_truncation = ( + self.surrogate["first_pass_max_overweight_truncation"] + ) + cfg.freeze_max_weight_after = ( + self.surrogate["first_pass_freeze_max_weight_after"] + ) + cfg.start_batch_size = vegas_args["start_batch_size"] + cfg.max_batch_size = vegas_args["max_batch_size"] + cfg.survey_min_iters = gen_args["survey_min_iters"] + cfg.survey_max_iters = gen_args["survey_max_iters"] + cfg.survey_target_precision = gen_args["survey_target_precision"] + cfg.optimization_patience = vegas_args["optimization_patience"] + cfg.optimization_threshold = vegas_args["optimization_threshold"] + cfg.cpu_batch_size = gen_args["cpu_batch_size"] + cfg.gpu_batch_size = gen_args["gpu_batch_size"] + cfg.verbosity = resolve_verbosity(run_args["verbosity"]) + cfg.combine_thread_count = run_args["combine_thread_pool_size"] + if self.needs_systematics_matrix_elements() and cfg.combine_thread_count != 1: + # SystematicsCalculator::compute() runs inside the combine workers, + # and a matrix element's per-thread process instances live in a + # ThreadResource indexed by ThreadPool::thread_index() -- the + # *calling* thread's index. A parallel combine pool would therefore + # need one process instance per worker; building N of them to index + # them safely is not worth it, the more so as + # SystematicsCalculator::matrix_elements() already serialises every + # evaluation on its own mutex, so that parallelism was not being + # used anyway. Run the combine stage single-threaded instead. + logger.info( + "systematics: the combine stage runs single-threaded because a " + "subprocess needs its matrix element re-evaluated for the " + "renormalisation scale variations") + cfg.combine_thread_count = 1 + cfg.cut_efficiency_threshold = gen_args["cut_efficiency_threshold"] + cfg.max_cut_repetitions = gen_args["max_cut_repetitions"] + self.event_generator_config = cfg + self.event_generator = None + + def init_subprocesses(self) -> None: + self.backends = self.compile_matrix_elements() + self.subprocesses = [] + if self.merged_subprocess_data is None: + for subproc_id, meta in enumerate(self.subprocess_data): + self.subprocesses.append(MadgraphSubprocess(self, meta, subproc_id)) + else: + for subproc_id, meta in enumerate(self.merged_subprocess_data): + self.subprocesses.append( + MadgraphSubprocess(self, meta, subproc_id, self.subprocess_data) + ) + + def compile_matrix_elements(self) -> list[str]: + """Build the matrix-element library of every subprocess, and return the + build backend of each requested device (in the same order). + + A 'cpu' device builds the backend named by the 'cpu_mode' entry; the + cuda/hip devices ignore cpu_mode and build their own backend. A + cpu_mode of 'auto' is resolved to the actual SIMD backend chosen on + this machine. + + SubProcesses/makefile is a dispatcher over the P* directories, so a + single 'make -j N' there builds all the subprocesses at once with one + shared pool of N jobs: no subprocess is built with N jobs while the + others wait, and none is limited to N/#subprocesses jobs either. + """ + device_names = self.run_card["run"]["device"] + if not isinstance(device_names, list): + device_names = [device_names] + cpu_mode = self.run_card["run"]["cpu_mode"] + backends = [backend_of(name, cpu_mode) for name in device_names] + if not self.subprocess_data: + return backends + + first_proc_path = self.subprocess_data[0]["path"] + subproc_path = os.path.dirname(first_proc_path) + + fptype = fptype_of(self.run_card["run"]["precision"]) + + # Resolve cpu_mode='auto' once (the build rules pick the best SIMD + # backend available here), so all subprocesses agree on the library names. + auto_backend = None + if "auto" in backends: + auto_backend = resolve_auto_backend(first_proc_path) + logger.info("cpu_mode 'auto' resolved as '%s'", auto_backend) + resolved = [ + auto_backend if backend == "auto" else backend + for backend in backends + ] + + nb_core = self.run_card["run"]["cpu_thread_pool_size"] + if not nb_core or nb_core < 0: + nb_core = os.cpu_count() or 1 + + log_path = os.path.join(self.run_path, "compile_subprocesses.log") + for backend in resolved: + missing = [ + meta for meta in self.subprocess_data + if not os.path.isfile(meta["me_path"].format(device=backend)) + ] + if not missing: + continue + logger.info( + f"Start compilation of SubProcesses for device '{backend}' " + f"({len(missing)} subprocess(es), {nb_core} parallel job(s)), " + f"see log detail in {log_path}" + ) + start_time = time.time() + self.make_subprocesses( + subproc_path, + [f"BACKEND={backend}", f"FPTYPE={fptype}", "USEBUILDDIR=1"], + nb_core, log_path, + ) + logger.info( + f"Compilation of SubProcesses done in {time.time() - start_time:.1f} s" + ) + return resolved + + @staticmethod + def make_subprocesses( + subproc_path: str, args: list[str], nb_core: int, log_path: str + ) -> None: + """Run 'make -j nb_core' in SubProcesses/, appending the full build + output to log_path (one run per device, so the file is appended to).""" + command = ["make", f"-j{nb_core}"] + args + with open(log_path, "a") as log: + log.write(f"\n$ cd {subproc_path} && {' '.join(command)}\n") + log.flush() + returncode = subprocess.call( + command, cwd=subproc_path, stdout=log, stderr=subprocess.STDOUT + ) + if returncode != 0: + raise RuntimeError( + f"Compilation of the SubProcesses failed with exit status " + f"{returncode}; see {log_path} for details" + ) + + def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenerator: + channel_generators = [] + subprocess_channels = [[] for _ in self.subprocesses] + for i, (subproc, phasespace) in enumerate(zip(self.subprocesses, phasespaces)): + for integrand, channel in zip( + subproc.build_integrands(phasespace), phasespace.channels + ): + if channel.event_generator is None: + channel.event_generator = ms.ChannelEventGenerator( + contexts=self.contexts, + integrand=integrand, + event_file=os.path.join(self.temp_path, f"events.{i}.{channel.name}.npy"), + weight_file=os.path.join(self.temp_path, f"weights.{i}.{channel.name}.npy"), + config=self.event_generator_config, + subprocess_index=i, + name=f"{i}.{channel.name}", + histograms=subproc.histograms + ) + subprocess_channels[i].append(len(channel_generators)) + channel_generators.append(channel.event_generator) + + event_generator = ms.EventGenerator( + contexts=self.contexts, + channels=channel_generators, + status_file=self.status_file, + config=self.event_generator_config, + seed=self.run_seed, + groups=self.weighting_groups(subprocess_channels), + ) + unused_globals = ( + set(self.contexts[0].global_names()) - event_generator.used_globals() + ) + for context in self.contexts: + for global_name in unused_globals: + context.delete_global(global_name) + return event_generator + + def survey_phasespaces( + self, phasespaces: list[PhaseSpace | None], survey_pass: int = 0 + ) -> ms.EventGenerator | None: + ps_filtered = [ps for ps in phasespaces if ps is not None] + if len(ps_filtered) == 0: + return None + event_generator = self.build_event_generator(ps_filtered) + event_generator.survey(survey_pass) + return event_generator + + def survey(self) -> None: + # survey_pass distinguishes the survey() calls below: "both" mode can + # re-survey a channel carried over unchanged from the multichannel pass + # into the final (simplified) pass, and both passes schedule jobs on the + # same underlying ChannelEventGenerator. The explicit pass index keeps each + # pass's job seeds independent of the other passes' job counts, rather than + # depending on call history. + phasespace_mode = self.run_card["phasespace"]["mode"] + if self.is_decay and phasespace_mode != "multichannel": + # The flat mapping is built from a synthetic two-incoming diagram, + # which a decay has no counterpart for -- and 'both'/'auto' end up + # simplifying towards it. Decays have few channels, so the + # multichannel phase space is the right one regardless. + logger.info( + "decay mode: using the multichannel phase space instead of " + "'%s'", phasespace_mode + ) + phasespace_mode = "multichannel" + if phasespace_mode in ["multichannel", "both", "auto"]: + self.phasespaces = [ + subproc.build_multichannel_phasespace() + for subproc in self.subprocesses + ] + self.event_generator = self.survey_phasespaces(self.phasespaces, 0) + elif phasespace_mode == "flat": + self.phasespaces = [ + subproc.build_flat_phasespace() + for subproc in self.subprocesses + ] + self.event_generator = self.survey_phasespaces(self.phasespaces, 0) + else: + raise ValueError("Unknown phasespace mode") + + channel_status = self.event_generator.channel_status() + chan_offset = 0 + madnis_enabled = False + for subproc, ps in zip(self.subprocesses, self.phasespaces): + mean = 0. + variance = 0. + count_opt = 0 + for status in channel_status[chan_offset:chan_offset + len(ps.channels)]: + mean += status.mean_abs + variance += status.error_abs**2 + count_opt += status.count_opt + # A subprocess whose matrix element is identically zero (an FCNC + # channel with every Wilson coefficient at zero, say) integrates to + # exactly zero, so there is no relative spread to speak of. Report 0 + # rather than dividing by it: that asks for the smallest networks and + # leaves MadNIS off, which is what an empty subprocess wants. + rsd = (variance * count_opt)**0.5 / mean if mean else 0. + subproc.set_madnis_auto_settings(rsd) + chan_offset += len(ps.channels) + if subproc.madnis_settings["enable"]: + madnis_enabled = True + if not ( + phasespace_mode == "both" or (phasespace_mode == "auto" and madnis_enabled) + ): + return + + phasespaces_multi = self.phasespaces + cross_sections = [] + index = 0 + for phasespace in phasespaces_multi: + channel_count = len(phasespace.channels) + # rank the channels by the integral of |w|: a channel whose signed + # integral cancels (an interference) can still carry most of the + # variance, and folding it into the flat channel would blow it up + cross_sections.append([ + status.mean_abs + for status in channel_status[index:index + channel_count] + ]) + index += channel_count + + self.phasespaces = [ + subproc.simplify_phasespace(ps_multi, cross_secs) + for subproc, ps_multi, cross_secs in zip( + self.subprocesses, phasespaces_multi, cross_sections + ) + ] + if any( + ps_multi is not ps_both + for ps_multi, ps_both in zip(phasespaces_multi, self.phasespaces) + ): + # distinct survey_pass: a channel carried over unchanged from the + # multichannel pass (pass 0) into this resurvey must not share its + # seed stream with that earlier pass. + self.event_generator = self.survey_phasespaces(self.phasespaces, 1) + + def train_madnis(self) -> None: + madnis_args = self.run_card["madnis"] + if not any(subproc.madnis_settings["enable"] for subproc in self.subprocesses): + return + + gen_args = self.run_card["generation"] + run_args = self.run_card["run"] + + verbosity = resolve_verbosity(run_args["verbosity"]) + madnis_phasespaces = [] + training_args = [] + self.event_generator = None + for subproc, phasespace in zip(self.subprocesses, self.phasespaces): + for channel in phasespace.channels: + channel.event_generator = None + + config = ms.MadnisConfig() + config.learning_rate = subproc.madnis_settings["lr"] + config.batches = subproc.madnis_settings["train_batches"] + config.log_interval = madnis_args["log_interval"] + config.integration_history_length = madnis_args["integration_history_length"] + config.channel_dropping_interval = madnis_args["channel_dropping_interval"] + config.channel_dropping_threshold = madnis_args["channel_dropping_threshold"] + config.cpu_generator_batch_size = gen_args["cpu_batch_size"] + config.gpu_generator_batch_size = gen_args["gpu_batch_size"] + config.gpu_generator_batch_granularity = madnis_args["gpu_generator_batch_granularity"] + config.generator_target_size_factor = madnis_args["generator_target_size_factor"] + config.batch_size_offset = madnis_args["batch_size_offset"] + config.batch_size_per_channel = subproc.madnis_settings["batch_size_per_channel"] + config.uniform_channel_ratio = madnis_args["uniform_channel_ratio"] + config.lr_schedule = madnis_args["lr_scheduler"] + config.adam_beta1 = madnis_args["adam_beta1"] + config.adam_beta2 = madnis_args["adam_beta2"] + config.adam_eps = madnis_args["adam_eps"] + config.adam_weight_decay = madnis_args["adam_weight_decay"] + config.grad_clip_threshold = madnis_args["grad_clip_threshold"] + config.buffer_capacity = madnis_args["buffer_capacity"] + config.minimum_buffer_size = madnis_args["minimum_buffer_size"] + config.buffered_steps_fraction = madnis_args["buffered_steps_fraction"] + config.buffer_skip_batches = madnis_args["buffer_skip_batches"] + config.buffer_unweighting_quantile = madnis_args["buffer_unweighting_quantile"] + config.fixed_cwnet_fraction = subproc.madnis_settings["fixed_cwnet_fraction"] + config.softclip_threshold = madnis_args["softclip_threshold"] + config.compressed_channel_weight_count = madnis_args["compressed_channel_weight_count"] + phasespace = subproc.build_madnis(phasespace) + madnis_phasespaces.append(phasespace) + training_args.append( + ms.TrainingArgs( + config=config, + integrands=subproc.build_integrands( + phasespace, + madnis_training=True, + drop_cuts_and_rescale=madnis_args["drop_zero_integrands"] + ), + cwnet=phasespace.cwnet, + ) + ) + + gen_context = self.contexts[0] + opt_context = ms.Context( + device=self.devices[0], thread_count=self.pool_sizes[0] + ) + opt_context.copy_globals_from(gen_context) + + madnis_training = ms.MultiMadnisTraining( + generator_context=gen_context, + optimizer_context=opt_context, + training_args=training_args, + verbosity=verbosity, + status_file=self.status_file, + # Reuses the run's resolved seed (also used by build_event_generator()). + # Only the single-channel CPU sample-generation path is currently seeded + # -- buffered training and GPU multi-channel batches are still + # non-deterministic. + seed=self.run_seed, + ) + madnis_training.train() + for phasespace, active_channels in zip( + madnis_phasespaces, madnis_training.active_channels() + ): + phasespace.channels = [ + phasespace.channels[index] for index in active_channels + ] + del madnis_training + del opt_context + self.phasespaces = madnis_phasespaces + for context in self.contexts[1:]: + context.copy_globals_from(self.contexts[0]) + self.event_generator = self.build_event_generator(madnis_phasespaces) + self.madnis_trained = True + + def update_madnis_status_single( + self, batch: int, batch_target: int, loss: float, lr: float, channel_count: int + ) -> None: + now = time.time() + if batch + 1 < batch_target: + if now - self.last_update_time < 0.1: + return + self.last_update_time = now + progress_bar = ms.format_progress((batch + 1) / batch_target, 52) + time_diff = now - self.madnis_wall_time + time_str = f"{format_time(time_diff)}" + else: + progress_bar = "" + wall_diff = now - self.madnis_wall_time + cpu_diff = time.process_time() - self.madnis_cpu_time + time_str = ( + f"{format_time(wall_diff, centi=True)} wall, " + f"{format_time(cpu_diff, centi=True)} cpu" + ) + batch_str = f"{batch + 1} / {batch_target}" + self.madnis_box.set_column(1, [ + f"{batch_str:<15} {progress_bar}", + f"{loss:>.4f}", + f"{channel_count}", + time_str + ]) + self.madnis_box.print_update() + + def update_madnis_status_multi( + self, + subproc_id: int, + batch: int, + batch_target: int, + loss: float, + lr: float, + channel_count: int + ) -> None: + now = time.time() + subproc_count = len(self.subprocesses) + if batch + 1 < batch_target: + if now - self.last_update_time < 0.1: + return + self.last_update_time = now + progress_bar = ms.format_progress((batch + 1) / batch_target, 34) + progress_bar_all = ms.format_progress( + (subproc_id * batch_target + batch + 1) / (subproc_count * batch_target), + 52 + ) + time_str = f"{format_time(now - self.madnis_wall_time)}" + subproc_str = f"{subproc_id} / {subproc_count}" + elif subproc_id < subproc_count - 1: + progress_bar = "" + progress_bar_all = ms.format_progress( + ((subproc_id + 1) * batch_target + 1) / (subproc_count * batch_target), + 52 + ) + time_str = f"{format_time(now - self.madnis_wall_time)}" + subproc_str = f"{subproc_id} / {subproc_count}" + else: + progress_bar = "" + progress_bar_all = "" + wall_diff = now - self.madnis_wall_time + cpu_diff = time.process_time() - self.madnis_cpu_time + time_str = ( + f"{format_time(wall_diff, centi=True)} wall, " + f"{format_time(cpu_diff, centi=True)} cpu" + ) + subproc_str = f"{subproc_count} / {subproc_count}" + batch_str = f"{batch + 1} / {batch_target}" + self.madnis_upper_box.set_column(1, [ + f"{subproc_str:<15} {progress_bar_all}", + time_str, + ]) + self.madnis_lower_box.set_row(subproc_id + 1, [ + f"{subproc_id}", + f"{loss:>.4f}", + f"{channel_count}", + f"{batch_str:<15} {progress_bar}", + ]) + self.madnis_upper_box.print_update() + self.madnis_lower_box.print_update() + + @property + def gridpack_run_mode(self) -> str: + """How this run ends: a gridpack run may skip the events (see the + [gridpack] run_mode run card option); any other run is a regular one.""" + if not self.run_card["gridpack"]["save_gridpack"]: + return "regular" + return self.run_card["gridpack"]["run_mode"] + + def finish_without_events(self) -> None: + """The end of a gridpack run that writes no events: only what the + gridpack needs is prepared, then it is stored.""" + if self.gridpack_run_mode == "minimal": + # a trained MadNIS needs nothing more, a VEGAS grid has to converge + if not self.madnis_trained: + self.event_generator.optimize() + else: + self.event_generator.fix_max_weights() + self.build_systematics() + self.save_gridpack() + + def generate_events(self) -> None: + if self.gridpack_run_mode != "regular": + self.finish_without_events() + return + start_time = get_start_time() + self.event_generator.generate() + output_format = self.run_card["run"]["output_format"] + systematics = self.build_systematics() + histograms = self.build_event_histograms() + if output_format == "compact_npy": + self.lhe_completer = None + self.event_generator.combine_to_compact_npy( + os.path.join(self.run_path, "events.npy"), systematics, histograms + ) + self.write_lhe_header(systematics) + self.save_lhe_completer() + elif output_format == "lhe_npy": + self.lhe_completer = self.build_lhe_completer() + self.event_generator.combine_to_lhe_npy( + os.path.join(self.run_path, "events.npy"), self.lhe_completer, + systematics, histograms + ) + self.write_lhe_header(systematics) + elif output_format == "lhe": + self.lhe_completer = self.build_lhe_completer() + lhe_path = os.path.join(self.run_path, "events.lhe") + self.event_generator.combine_to_lhe( + lhe_path, self.lhe_completer, + self.build_lhe_meta(), systematics, histograms + ) + # Ship the LHE compressed by default. These files are large and + # very compressible, madevent has always stored its events + # gzipped, and every consumer here already accepts either form + # (see _find_event_file). misc.gzip replaces events.lhe with + # events.lhe.gz, and switches to an external multithreaded tool + # above 256 MB. + misc.gzip(lhe_path) + else: + raise ValueError("Unknown output format") + self.write_hwu(histograms) + self.make_plots(histograms) + if systematics is not None: + self.write_systematics_sidecar() + self.log_systematics_summary() + self.save_gridpack() + + def save_lhe_completer(self) -> None: + """Keep the LHE completer next to compact_npy events, so npy_to_lhe + can complete them into an LHE file later (a shower or MadSpin requested + after the run). Never fatal to the run.""" + try: + self.build_lhe_completer().save( + os.path.join(self.run_path, npy_to_lhe.COMPLETER_FILE)) + except Exception as err: + logger.warning("could not save the LHE completer; these events " + "cannot be converted to LHE later: %s", err) + + hwu_file_name = "MADatLO.HwU" + plot_dir_name = "plots" + + def make_plots(self, histograms) -> None: + """Draw the event-sample histograms into Events//plots, with the + scale and PDF bands info.json carries, and say where they landed. + + This does not go through the HwU file: the bands are already computed + per bin by madspace, so drawing them directly is both fewer steps and + more than histograms.py can do without the LHAPDF python module. + """ + if histograms is None or not self.run_card["run"]["make_plots"]: + return + out_dir = os.path.join(self.run_path, self.plot_dir_name) + try: + data = json.loads(histograms.to_json(self.systematics)) + written = plots.render(data, out_dir) + except plots.BackendMissing as error: + logger.info("matplotlib is not available (%s): no plots. The " + "distributions are in %s", + error, os.path.join(self.run_path, "info.json")) + return + except Exception as error: + logger.warning("could not draw the histograms: %s", error) + return + if not written: + return + logger.info("%d plot(s) stored in %s", len(written), out_dir) + + def write_hwu(self, histograms) -> None: + """Write the event-sample histograms next to the events in the HwU + format (MADatLO.HwU), the one an aMC@NLO run writes as MADatNLO.HwU + and madgraph/various/histograms.py reads: an mg7 distribution can then + be overlaid on an NLO one without converting anything in between. The + numbers are the ones info.json carries, bands included, so this is off + by default -- it is a second copy in another format, for the runs that + are going to be plotted.""" + if histograms is None or not self.run_card["run"]["write_hwu"]: + return + try: + data = json.loads(histograms.to_json(self.systematics)) + summary = (json.loads(self.systematics.summary()) + if self.systematics is not None else None) + text = hwu_output.to_hwu(data, summary) + except Exception as error: + logger.warning("could not write the HwU histograms: %s", error) + return + if not text: + return + with open(os.path.join(self.run_path, self.hwu_file_name), "w") as f: + f.write(text) + + def write_systematics_sidecar(self) -> None: + """Describe the variation weights next to the event file + (events.weights.json): ids, scale factors, PDF set/member and the + text, so npy consumers get the same metadata as the LHE + header, plus the per-variation cross sections.""" + if self.systematics is None: + return + data = json.loads(self.systematics.summary()) + data["initrwgt"] = self.systematics.initrwgt() + data["columns"] = ["rwgt_%d" % i for i in self.systematics.weight_ids] + with open(os.path.join(self.run_path, "events.weights.json"), "w") as f: + json.dump(data, f, indent=1) + + @staticmethod + def _histogram_mean(hist_min, hist_max, bin_values): + """Cross-section-weighted mean of a histogrammed observable over its + range. ``bin_values`` carries the underflow and overflow bins first and + last, as both madspace histogram kinds do; they have no position and + are left out.""" + values = list(bin_values)[1:-1] + total = sum(values) + if not values or total == 0: + return None + width = (hist_max - hist_min) / len(values) + return sum(v * (hist_min + (i + 0.5) * width) + for i, v in enumerate(values)) / total + + def _histogram_means(self) -> dict: + """name -> mean of every [histograms] entry, from the post-processing + histograms (the final event sample) when they were filled, otherwise + from the weighted integration histograms.""" + if self.event_histograms is not None: + data = json.loads(self.event_histograms.to_json(self.systematics)) + hists = [(h["name"], h["min"], h["max"], h["bin_values"]) + for h in data] + else: + hists = [(h.name, h.min, h.max, h.bin_values) + for h in self.event_generator.histograms()] + means = {} + for name, hist_min, hist_max, values in hists: + mean = self._histogram_mean(hist_min, hist_max, values) + if mean is not None: + means[name] = mean + return means + + def get_result(self) -> dict: + """Return the run result: cross-section (pb) with MC error, the number + of (unweighted) events, and the mean of every observable declared in the + [histograms] section. Used to build the scan summary.""" + status = self.event_generator.status() + result = {'cross(pb)': status.mean, 'error(pb)': status.error, + 'nb_event': status.count_unweighted} + try: + if self.systematics is not None and self.systematics.weight_count: + summary = json.loads(self.systematics.summary()) + scale = systematics_summary.scale_percentages(summary) + if scale is not None: + result['scale_up(%)'], result['scale_down(%)'] = scale + pdf = systematics_summary.pdf_percentages(summary) + if pdf: + # one column pair; the nominal set's entry comes first + _entry, up, down = pdf[0] + result['pdf_up(%)'], result['pdf_down(%)'] = up, down + except Exception as err: + logger.warning("could not extract the systematics summary: %s", err) + try: + for name, mean in self._histogram_means().items(): + result['<%s>' % name] = mean + except Exception as err: + logger.warning("could not extract observable means: %s", err) + return result + + def _beam_info(self): + """Return (beam_pdg_ids, beam_energies) for the LHE block. + + `incoming` holds the *partonic* initial state (e.g. gluons), not the + beam particle, so hadronic beams are protons (2212); leptonic beams are + the incoming leptons themselves.""" + half_e = float(self.e_cm) / 2. + if self.is_decay: + # No beams. LHE has no way to say that, so report the decaying + # particle at rest as a single "beam"; the second slot is empty. + data = self.subprocess_data[0] + pdg = clean_pids(data["incoming"])[0] + return [pdg, 0], [float(self.e_cm), 0.0] + if not self.leptonic: + # hadronic collider: proton beams (p-pbar is not distinguished) + return [2212, 2212], [half_e, half_e] + data = self.subprocess_data[0] + incoming = data["incoming"] + flavor0 = data["flavors"][0]["options"][0] + init_pdgs = flavor0[:len(incoming)] + beam_pdgs = [pdg if abs(code) in (81, 82) else code + for code, pdg in zip(incoming, init_pdgs)] + return beam_pdgs, [half_e, half_e] + + def _lhapdf_id(self): + """Central LHAPDF id of the beam PDF set (read from its .info SetIndex), + or -1 for a leptonic beam (no PDF).""" + if self.leptonic or not self.pdf_dir: + return -1 + pdf_set = self.run_card["beam"]["pdf"] + info = os.path.join(self.pdf_dir, pdf_set, "%s.info" % pdf_set) + try: + for line in open(info): + if line.strip().startswith("SetIndex:"): + return int(line.split(":", 1)[1].strip()) + except Exception as err: + logger.warning("could not read LHAPDF id from %s: %s", info, err) + return -1 + + def build_lhe_meta(self, systematics=None): + """Build the LHE header/ metadata: the param_card () and the + run_card.toml () headers plus the beam/PDF/cross-section info + needed by downstream tools (systematics, MadSpin, ...). + + `combine_to_lhe` injects the header itself, so callers that + go through it should leave `systematics` unset; pass it only when + building meta for a writer that bypasses that injection (e.g. + write_lhe_header).""" + beam_pdgs, energies = self._beam_info() + lhaid = self._lhapdf_id() + pdf_group = -1 if self.leptonic else 0 + status = self.event_generator.status() + xsec, err = status.mean, status.error + weight_mode, max_weight = lhe_weight_info(status) + with open(self.param_card_path) as f: + param_text = f.read() + with open(os.path.join("Cards", "run_card.toml")) as f: + run_text = f.read() + headers = [] + # generation commands (model + process): read by MadSpin/reweight/... + proc_card = os.path.join("Cards", "proc_card_mg5.dat") + if os.path.exists(proc_card): + with open(proc_card) as f: + headers.append(ms.LHEHeader(name="MG5ProcCard", content=f.read())) + headers.append(ms.LHEHeader(name="slha", content=param_text)) + headers.append(ms.LHEHeader(name="MG7RunCard", content=run_text)) + # The resolved seed (even when the run_card requested a random one via + # seed = -1), so the run can be reproduced from the LHE file alone. + headers.append(ms.LHEHeader(name="MG7Seed", content=str(self.run_seed))) + if systematics is not None and systematics.weight_ids: + headers.append(ms.LHEHeader( + name="initrwgt", content=systematics.initrwgt(), escape_content=False + )) + return ms.LHEMeta( + beam1_pdg_id=beam_pdgs[0], beam2_pdg_id=beam_pdgs[1], + beam1_energy=energies[0], beam2_energy=energies[1], + beam1_pdf_authors=pdf_group, beam2_pdf_authors=pdf_group, + beam1_pdf_id=lhaid, beam2_pdf_id=lhaid, + weight_mode=weight_mode, + # positional: the pybind arg name for max_weight is non-kwarg-safe + processes=[ms.LHEProcess(xsec, err, max_weight, 1)], + headers=headers, + ) + + def write_lhe_header(self, systematics) -> None: + """Write header.lhe next to events.npy: the
/ blocks + (run card, param card, beam/PDF, cross section) that the npy formats + otherwise drop, with no events.""" + header_path = os.path.join(self.run_path, "header.lhe") + writer = ms.LHEFileWriter(header_path, self.build_lhe_meta(systematics)) + del writer # closes the file (writes the closing tag) + + def build_lhe_completer(self): + all_mcdata = ( + [subproc.build_multi_channel_data() for subproc in self.subprocesses] + if self.merged_subprocess_data is None else + [build_multi_channel_data(meta, self) for meta in self.subprocess_data] + ) + subproc_args = [ + ms.SubprocArgs( + topologies = [topo[0] for topo in mcdata.topologies], + permutations = mcdata.permutations, + diagram_indices = mcdata.diagram_indices, + diagram_color_indices = mcdata.diagram_color_indices, + diagram_propagator_pdgs = mcdata.diagram_propagator_pdgs, + color_flows = meta["color_flows"], + pdg_color_types = { + int(key): value + for key, value in meta["pdg_color_types"].items() + }, + helicities = lhe_helicities(meta, self.run_card), + pdg_ids = [flavor["options"] for flavor in meta["flavors"]], + ) + for mcdata, meta in zip(all_mcdata, self.subprocess_data) + ] + return ms.LHECompleter( + subproc_args=subproc_args, + bw_cutoff=self.run_card["phasespace"]["bw_cutoff"] + ) + + def save_gridpack(self) -> None: + if not self.run_card["gridpack"]["save_gridpack"]: + return + + gridpack_path = os.path.join(self.run_path, "gridpack") + data_path = os.path.join(gridpack_path, "data") + events_path = os.path.join(gridpack_path, "Events") + os.mkdir(gridpack_path) + os.mkdir(data_path) + os.mkdir(events_path) + self.contexts[0].save_globals(os.path.join(data_path, "globals.tar")) + + # one json object per channel, joined by hand so that the channel + # json is not parsed and dumped again + channel_entries = [ + f"{json.dumps(channel.status().name)}:{channel.to_json(self.gridpack_run_mode == "fix_max_weight")}" + for channel in self.event_generator.channels() + ] + with open(os.path.join(data_path, "channels.json"), "w") as f: + f.write("{" + ",".join(channel_entries) + "}") + + lib_path = os.path.join(gridpack_path, "lib") + if self.run_card["gridpack"]["include_source"]: + os.mkdir(lib_path) + shutil.copytree("src", os.path.join(gridpack_path, "src")) + shutil.copytree("SubProcesses", os.path.join(gridpack_path, "SubProcesses")) + shutil.copytree("backend", os.path.join(gridpack_path, "backend")) + else: + shutil.copytree("lib", lib_path) + for _, path in self.custom_function_paths: + target = path if os.path.isabs(path) else os.path.join(gridpack_path, path) + if not os.path.exists(target): + os.makedirs(os.path.dirname(target), exist_ok=True) + shutil.copy(path, target) + + if self.run_card["gridpack"]["include_madspace_source"]: + shutil.copytree( + _MADSPACE_DIR, + os.path.join(gridpack_path, "madspace"), + ignore=shutil.ignore_patterns("build", "install"), + ) + + if self.run_card["gridpack"]["include_madspace"]: + shutil.copytree( + _INSTALL_DIR / "madspace", + os.path.join(gridpack_path, "madspace", "install", "madspace"), + ignore=shutil.ignore_patterns("include"), + ) + + matrix_elements = [] + for subproc in self.subprocess_data: + me_path = subproc["me_path"] + matrix_elements.append(me_path) + + cards_path = os.path.join(gridpack_path, "Cards") + os.mkdir(cards_path) + shutil.copy(os.path.join("Cards", "param_card.dat"), cards_path) + # Full run card with a header noting it is read-only in gridpack context. + import io as _io + _buf = _io.StringIO() + self.run_card.write(_buf) + _header = ( + "# This is the run card used to generate this gridpack.\n" + "# Modifying this file will have no effect on gridpack execution.\n" + "# To change event-generation settings, edit grid_run_card.toml.\n\n" + ) + with open(os.path.join(cards_path, "run_card.toml"), 'w') as _f: + _f.write(_header + _buf.getvalue()) + # Minimal card containing only the settings used by generate_events. + # The gridpack ships the libraries that were actually built, and their + # names carry the resolved backend, so the card it runs from has to name + # that backend as well: a cpu_mode of 'auto' would send the gridpack + # looking for a ..._auto.so that was never built. The full run_card.toml + # written just above keeps 'auto', since that is what was asked for. + device_names = self.run_card["run"]["device"] + if not isinstance(device_names, list): + device_names = [device_names] + resolved_cpu_mode = None + for name, backend in zip(device_names, getattr(self, "backends", [])): + if name.split(":")[0] == "cpu": + resolved_cpu_mode = backend + break + previous_cpu_mode = self.run_card["run"]["cpu_mode"] + if resolved_cpu_mode is not None: + self.run_card["run"]["cpu_mode"] = resolved_cpu_mode + # A gridpack is run to feed a shower or detector chain, which reads + # LHE, so it writes LHE by default whatever format this run used + # (bin/generate_events --output_format still selects an npy output). + previous_output_format = self.run_card["run"]["output_format"] + self.run_card["run"]["output_format"] = "lhe" + # the directories of this run mean nothing on the machine running the gridpack + previous_dirs = (self.run_card["gridpack"]["temp_output_dir"], + self.run_card["gridpack"]["output_dir"]) + self.run_card["gridpack"]["temp_output_dir"] = "" + self.run_card["gridpack"]["output_dir"] = "" + try: + self.run_card.write_gridpack_card( + os.path.join(cards_path, "grid_run_card.toml")) + finally: + self.run_card["run"]["cpu_mode"] = previous_cpu_mode + self.run_card["run"]["output_format"] = previous_output_format + (self.run_card["gridpack"]["temp_output_dir"], + self.run_card["gridpack"]["output_dir"]) = previous_dirs + + bin_path = os.path.join(gridpack_path, "bin") + os.mkdir(bin_path) + gen_events_file = os.path.join(bin_path, "generate_events") + template_dir = os.path.dirname(__file__) + shutil.copy(os.path.join(template_dir, "gridpack.py"), gen_events_file) + shutil.copy(os.path.join(template_dir, "custom_function_loader.py"), bin_path) + os.chmod(gen_events_file, 0o755) + shutil.copy(os.path.join(template_dir, "gridpack_setup.py"), bin_path) + # streaming needs the fixed maximum weights of this mode + if self.gridpack_run_mode == "fix_max_weight": + shutil.copy(os.path.join(template_dir, "event_stream.py"), bin_path) + + data = { + "matrix_elements": matrix_elements, + # loaded before the matrix elements, as in the run + "custom_functions": self.custom_function_paths, + "me_parameters": me_parameters(self.run_card), + "source_hash": ms.SOURCE_HASH, + } + with open(os.path.join(data_path, "data.json"), "w") as f: + json.dump(data, f) + + if self.lhe_completer is None: + self.lhe_completer = self.build_lhe_completer() + self.lhe_completer.save(os.path.join(data_path, "lhe.json")) + self.save_gridpack_lhe_meta(os.path.join(data_path, "lhe_meta.json")) + if self.systematics_data is not None: + systematics_data = dict(self.systematics_data) + # A decay takes alpha_s from a constant grid written into the run + # directory rather than from an LHAPDF set, so the gridpack cannot + # look that file up: ship it and point at the shipped copy. + if self.fixed_alphas_info and \ + systematics_data["nominal_info_file"] == self.fixed_alphas_info: + name = os.path.basename(self.fixed_alphas_info) + shutil.copy(self.fixed_alphas_info, os.path.join(data_path, name)) + systematics_data["nominal_info_file"] = os.path.join("data", name) + with open(os.path.join(data_path, "systematics.json"), "w") as f: + json.dump(systematics_data, f) + if self.run_card["gridpack"]["compress"]: + shutil.make_archive(gridpack_path, "gztar", self.run_path, "gridpack") + shutil.rmtree(gridpack_path) + + def save_gridpack_lhe_meta(self, path) -> None: + """The run-independent part of the LHE header (cards, beams, PDF), for + the gridpack to write complete LHE headers: it adds the cross section + and seed of its own run, which is why those two are left out here.""" + meta = self.build_lhe_meta() + data = {key: getattr(meta, key) for key in ( + "beam1_pdg_id", "beam2_pdg_id", "beam1_energy", "beam2_energy", + "beam1_pdf_authors", "beam2_pdf_authors", "beam1_pdf_id", + "beam2_pdf_id", "weight_mode")} + data["headers"] = [ + {"name": h.name, "content": h.content, + "escape_content": h.escape_content} + for h in meta.headers if h.name != "MG7Seed"] + with open(path, "w") as f: + json.dump(data, f) + + def get_mass(self, pid: int) -> float: + return self.param_card.get_value("mass", pid) + + def get_width(self, pid: int) -> float: + return self.param_card.get_value("width", pid) diff --git a/madgraph/iolibs/template_files/mg7/model_widths.py b/madgraph/iolibs/template_files/mg7/model_widths.py new file mode 100644 index 0000000000..f007473a2a --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/model_widths.py @@ -0,0 +1,155 @@ +"""The model of a process directory and its auto widths. Does not import madspace.""" + +import logging +import os +import re +import subprocess + +from madgraph.various import misc +from madgraph.iolibs.template_files.mg7.bootstrap import MG_ROOT as _MG_ROOT + +logger = logging.getLogger("madgraph7") + + +_model_hash_warned = set() +_model_cache = {} + + +def read_stored_model(me_dir=None) -> "str | None": + """The model this process was generated with, as recorded at output time in + ``SubProcesses/model.txt``: a UFO path carrying its restriction (e.g. + ``.../models/sm-no_b_mass``) or a model name -- in either case something + ``import model`` understands. Returns None when the process was written + without that information. + + Warns (once per directory) when the model on disk no longer matches the + hash stored on the second line, i.e. it changed since the output was made. + """ + base = me_dir or os.getcwd() + model_file = os.path.join(base, "SubProcesses", "model.txt") + if not os.path.exists(model_file): + return None + with open(model_file) as f: + lines = f.read().splitlines() + model = lines[0].strip() if lines else "" + stored_hash = lines[1].strip() if len(lines) > 1 else "" + if not model: + logger.warning("SubProcesses/model.txt is empty; the model of this " + "process is unknown.") + return None + + # verify the model on disk still matches the one used at output time. The + # reference may carry a restriction suffix ('sm-no_b_mass'); the hash was + # taken on the UFO directory itself ('sm'). + model_dir = model if os.path.isdir(model) else model.rsplit('-', 1)[0] + if stored_hash and os.path.isdir(model_dir) \ + and model_file not in _model_hash_warned: + current_hash = misc.hash_model_files(model_dir) + if current_hash and current_hash != stored_hash: + _model_hash_warned.add(model_file) + logger.warning( + "The model at %s has changed since this process was generated " + "(hash mismatch); anything computed from it now (the 'auto' " + "widths, the dependent masses/widths of the param_card) may be " + "inconsistent with the matrix element.", model_dir) + return model + + +def _proc_characteristic(me_dir=None) -> dict: + """SubProcesses/proc_characteristics as a plain dict (empty when the file is + missing or unreadable). Plain, because ProcCharacteristic is a ConfigFile, + whose get() is the parameter accessor and takes no default.""" + from madgraph.various import banner as _banner_mod + path = os.path.join(me_dir or os.getcwd(), 'SubProcesses', + 'proc_characteristics') + if os.path.exists(path): + try: + return dict(_banner_mod.ProcCharacteristic(path)) + except Exception as error: + logger.debug("could not read %s: %s", path, error) + return {} + + +def load_process_model(me_dir=None): + """Import the UFO model this process was generated with, as + :func:`models.import_ufo.import_model` builds it -- the object every + model-level consumer of the param_card expects (see + :meth:`MG7Cmd.get_model`). None when the process recorded no model, or when + that model no longer imports. + + Cached per (reference, complex-mass-scheme): importing a model is slow and + the card question asks for it again at every 'update dependent'.""" + ref = read_stored_model(me_dir) + if not ref: + return None + cms = _proc_characteristic(me_dir).get('complex_mass_scheme', False) + key = (ref, bool(cms)) + if key not in _model_cache: + import models.import_ufo as import_ufo + try: + with misc.MuteLogger(['madgraph.model'], [50]): + _model_cache[key] = import_ufo.import_model( + ref, complex_mass_scheme=bool(cms)) + except Exception as error: + # 'update dependent' calls this under a SIGALRM whose handler + # raises a TimeOutError of its own (a class defined inside the + # caller, so it can only be recognised by name). Let it through: + # the caller then says the model took too long to load and how to + # force it, instead of reporting a model that does not import. + if type(error).__name__ == 'TimeOutError': + raise + logger.debug("could not import the model %s: %s", ref, error) + _model_cache[key] = None + return _model_cache[key] + + +def compute_auto_widths(param_card_path=os.path.join("Cards", "param_card.dat")) -> None: + """Fill any width set to ``auto`` in the param_card, using madgraph and the + model stored at output time (``SubProcesses/model.txt``), and write the + result back into the card. A no-op when the card has no ``auto`` width. + + Called before every generation, so a single run and each scan point (whose + param_card is rewritten just before ``run_single``) both get model-computed + widths.""" + try: + with open(param_card_path) as f: + text = f.read() + except OSError: + return + # matches "DECAY auto" (optionally "auto@NLO"), as in + # common_run_interface.static_check_param_card + pdgs = re.findall(r"(?im)^\s*decay\s+([+-]?\d+)\s+auto", text) + if not pdgs: + return + pdgs = list(dict.fromkeys(pdgs)) # de-duplicate, keep order + + model = read_stored_model() + if not model: + logger.warning( + "The param_card requests 'auto' width(s) for %s but the model of " + "this process is unknown; leaving them as-is.", " ".join(pdgs)) + return + + mg5 = str(_MG_ROOT / "bin" / "madgraph") + if not os.path.exists(mg5): + logger.warning("Cannot find madgraph at %s; 'auto' widths not computed.", mg5) + return + + import tempfile + cmds = "import model %s\ncompute_widths %s --path=%s\n" % ( + model, " ".join(pdgs), os.path.abspath(param_card_path)) + with tempfile.NamedTemporaryFile("w", suffix=".mg5", delete=False) as fh: + fh.write(cmds) + cmdfile = fh.name + try: + logger.info("Computing 'auto' width(s) for %s ...", " ".join(pdgs)) + proc = subprocess.run([mg5, cmdfile]) + if proc.returncode != 0: + logger.warning( + "compute_widths returned a non-zero exit code; the param_card " + "may still contain 'auto' entries.") + finally: + try: + os.remove(cmdfile) + except OSError: + pass diff --git a/madgraph/iolibs/template_files/mg7/observables.py b/madgraph/iolibs/template_files/mg7/observables.py new file mode 100644 index 0000000000..27523e7d73 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/observables.py @@ -0,0 +1,178 @@ +"""Observables, cuts and histograms of the run card, and PDG id helpers.""" + +from dataclasses import dataclass + +from madgraph.iolibs.template_files.mg7.common import ms + +@dataclass +class CutItem: + observable_kwargs: dict + min: float + max: float + mode: str + + +@dataclass +class HistItem: + observable_kwargs: dict + min: float + max: float + bin_count: int + # the distribution of the event weight itself (the reserved "weight" key + # of [histograms]); it is read off the event record, not computed from + # the momenta, so it has no observable + name: str = "" + from_weight: bool = False + + +# Flavor-merged legs carry a group id instead of a pdg. Every member of a group +# shares the same mass -- that is what makes them mergeable -- so any member is +# a valid representative for the kinematics. +_MERGED_PID_REPRESENTATIVE = { + 81: 1, # light quarks d u s c + 82: 11, # charged leptons e mu + 83: 12, # neutrinos ve vm vt +} + + +def clean_pids(pids: list[int]) -> list[int]: + pids_out = [] + for pid in pids: + pid = abs(pid) + if pid in _MERGED_PID_REPRESENTATIVE: + pid = _MERGED_PID_REPRESENTATIVE[pid] + elif 81 <= pid <= 99: + # Reserved for flavor merging. Anything outside this window is a + # real pdg (BSM models use codes in the millions), so let it pass. + raise ValueError( + f"unknown flavor-merged particle id {pid}; add its " + "representative to _MERGED_PID_REPRESENTATIVE" + ) + pids_out.append(pid) + return pids_out + + +# observables that tell a particle from its antiparticle, and so need the signed +# pdg ids of the actual flavours rather than clean_pids' unsigned representatives +_SIGNED_OBSERVABLES = {"sfos_pair_mass"} + + +def _flavor_options(meta: dict, unmerged_meta: dict | None = None) -> list[list[int]]: + """Every flavour assignment (signed pdg ids, incoming then outgoing) a + subprocess evaluates with the same momenta.""" + if "flavors" in meta: + return [ + option for flavor in meta["flavors"] for option in flavor["options"] + ] + return [ + option + for subproc in meta.get("subprocesses", []) + for option in _flavor_options(unmerged_meta[subproc]) + ] + + +def observable_pids( + meta: dict, observable_kwargs: dict, unmerged_meta: dict | None = None +) -> list[int]: + """The pdg ids an observable of this subprocess is built with. + + Most observables only select particles by group, which the unsigned + representatives of clean_pids do. One that pairs a particle with its + antiparticle needs the real signed flavours: those of the first flavour + assignment, provided every assignment of the subprocess forms the same + pairs -- the cut is applied once to momenta that all of them share. + """ + all_pids = clean_pids(meta["incoming"]) + clean_pids(meta["outgoing"]) + if observable_kwargs.get("observable") not in _SIGNED_OBSERVABLES: + return all_pids + options = _flavor_options(meta, unmerged_meta) + if not options: + raise ValueError( + f"{observable_kwargs.get('name')}: no flavour information for " + f"subprocess {meta['incoming']} > {meta['outgoing']}" + ) + n_in = len(meta["incoming"]) + groups = [set(group) for group in observable_kwargs["select_pids"]] + if len(groups) == 1: + groups = groups * 2 + + def sfos_pairs(pids): + out = range(n_in, len(pids)) + return frozenset( + frozenset((i, j)) + for i in out if pids[i] in groups[0] + for j in out if pids[j] in groups[1] + if i != j and pids[i] != 0 and pids[i] == -pids[j] + ) + + reference = sfos_pairs(options[0]) + if any(sfos_pairs(option) != reference for option in options[1:]): + raise ValueError( + f"{observable_kwargs.get('name')}: the flavours grouped into subprocess " + f"{meta['incoming']} > {meta['outgoing']} pair the particles " + "differently, so a same-flavour opposite-sign cut cannot be applied " + "to their common momenta; regenerate the process after " + "'set apply_flavor_grouping False'" + ) + return list(options[0]) + + +def build_observable(pids: list[int], n_in: int, observable_kwargs: dict): + """ms.Observable of a subprocess with `n_in` incoming particles. + + madspace takes the first two momenta to be the beams (ignore_incoming), so + for a 1 -> N decay it would drop the first decay product from every + selection. There the selection runs over all the momenta instead, with the + incoming particle set to pdg id 0, which is in no group. + """ + if n_in == 2: + return ms.Observable(pids, **observable_kwargs) + pids = [0] * n_in + list(pids[n_in:]) + return ms.Observable(pids, **dict(observable_kwargs, ignore_incoming=False)) + + +def decay_products(meta: dict, unmerged_meta: dict | None = None) -> set[int]: + """Outgoing positions (0-based, beams excluded) of the particles that come + from an on-shell (decay-chain) propagator, as MadEvent's check_decay marks + them. Every channel of a subprocess shares its decay chains, so the first + channel tells, read through its first diagram's permutation.""" + if not meta["channels"]: + return set() + channel = meta["channels"][0] + topo_channel = channel + if "vertices" not in channel: + topo_channel = unmerged_meta[channel["subprocess"]]["channels"][channel["channel"]] + on_shell = topo_channel.get("on_shell_propagators", []) + if not on_shell: + return set() + vertices = topo_channel["vertices"] + # the closing vertex creates no propagator + created = {vertex[-1]: vertex[:-1] for vertex in vertices[:-1]} + + def legs(name): + if name in created: + return set().union(*(legs(child) for child in created[name])) + return {name} + + n_in = len(meta["incoming"]) + n_out = len(meta["outgoing"]) + products = set() + for prop in on_shell: + external = legs(f"p{prop}") + incoming = {leg for leg in external if leg.startswith("i")} + outgoing = {int(leg[1:]) for leg in external if leg.startswith("o")} + if not incoming: + products |= outgoing + elif len(incoming) == n_in: + # built from the beams: the decay is everything else + products |= set(range(n_out)) - outgoing + # event[i] = topology[permutation[i]] + permutation = channel["diagrams"][0]["permutation"] + return { + i - n_in for i in range(n_in, len(permutation)) + if permutation[i] - n_in in products + } + + +def pid_is_qcd(pid: int): + return abs(pid) in [21, 1, 2, 3, 4, 5, 6, 81] diff --git a/madgraph/iolibs/template_files/mg7/options.py b/madgraph/iolibs/template_files/mg7/options.py new file mode 100644 index 0000000000..9225c37271 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/options.py @@ -0,0 +1,86 @@ +"""The MG5 options of the process directory and the LHAPDF paths they select.""" + +import os + +from madgraph.various import misc + +_ROOTED_OPTIONS = ('lhapdf', 'lhapdf_py3', 'lhapdf_py2', 'heptools_install_dir') + +def load_mg5_options(me_dir=None) -> dict: + """Read the tool paths from the MadGraph configuration, so the launcher + knows where LHAPDF and the optional programs (Pythia8/Delphes/MadSpin/ + reweight/analysis) live. + + Files are read least- to most-specific, each overriding the previous, which + is the same layering CommonRunCmd.set_configuration uses: the card inside + this process directory has the last word, then this installation, then the + per-user file. Relative values are resolved against the root of the file + they came from. + + ``me_dir`` is the process directory whose Cards/me5_configuration.txt has + the last word. It defaults to the working directory, which is right for + bin/generate_events (it chdirs there first) but not for an in-process + caller, which has not moved yet. + """ + + import madgraph + mg5dir = os.path.dirname(os.path.dirname(os.path.abspath(madgraph.__file__))) + me_dir = os.getcwd() if me_dir is None else me_dir + + options = { + 'pythia-pgs_path': None, 'pythia8_path': None, + 'madanalysis5_path': None, 'exrootanalysis_path': None, 'delphes_path': None, + 'rivet_path': None, 'contur_path': None, 'f2py_compiler': None, + 'lhapdf': None, 'lhapdf_py3': None, 'lhapdf_py2': None, 'timeout': 0, + 'heptools_install_dir': None, 'hepmc3_path': None, + } + config_files = [] + base_config = misc.base_config_file() + if base_config: + config_files.append((base_config, mg5dir)) + user_config = misc.user_config_file() + if user_config: + config_files.append((user_config, mg5dir)) + config_files.append((misc.install_config_file(mg5dir), mg5dir)) + config_files.append((os.path.join(me_dir, 'Cards', 'me5_configuration.txt'), + me_dir)) + for cfg, root in config_files: + if not os.path.exists(cfg): + continue + with open(cfg) as fsock: + for line in fsock: + line = line.split('#', 1)[0] + if '=' not in line: + continue + name, value = (x.strip() for x in line.split('=', 1)) + if name not in options or value in ('', 'None'): + continue + if (name.endswith('_path') or name in _ROOTED_OPTIONS) and \ + not os.path.isabs(value) and \ + (os.sep in value or value.startswith('.')): + value = os.path.normpath(os.path.join(root, value)) + options[name] = value + options['mg5_path'] = mg5dir # enables MadSpin/reweight + return options + + +_LHAPDF = None + +def lhapdf_paths(refresh: bool = False) -> misc.LhapdfPaths: + """The LHAPDF installation this run should use, resolved once. + + Uses the same helper as the 'launch' command, so driving a process + directory straight through bin/generate_events behaves identically to + launching it from MadGraph. + """ + + global _LHAPDF + if _LHAPDF is None or refresh: + _LHAPDF = misc.resolve_lhapdf(load_mg5_options()) + # Export what we found: the real LHAPDF library, systematics and the + # MadSpin/reweight subprocesses read these from the environment. + if _LHAPDF.data_paths: + os.environ["LHAPDF_DATA_PATH"] = os.pathsep.join(_LHAPDF.data_paths) + if _LHAPDF.config: + os.environ["MADGRAPH_LHAPDF_CONFIG"] = _LHAPDF.config + return _LHAPDF diff --git a/madgraph/iolibs/template_files/mg7/postprocess.py b/madgraph/iolibs/template_files/mg7/postprocess.py new file mode 100644 index 0000000000..38bb38cf76 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/postprocess.py @@ -0,0 +1,380 @@ +"""Post-processing tools run on the events, and their logging.""" + +import logging +import os +import sys +import time + +from madgraph.iolibs.template_files.mg7 import npy_to_lhe +from madgraph.iolibs.template_files.mg7.options import lhapdf_paths, load_mg5_options + +def _find_event_file(run_path): + """Locate the LHE event file produced by generate_events (if any).""" + for name in ("events.lhe", "events.lhe.gz"): + path = os.path.join(run_path, name) + if os.path.exists(path): + return path + return None + + +def _find_or_convert_event_file(run_path, log): + """The run's LHE event file, converting its npy events when that is all + the run wrote (see npy_to_lhe).""" + lhe_path = _find_event_file(run_path) + if lhe_path is None and npy_to_lhe.can_convert(run_path): + log.info("No LHE event file in %s: converting the npy events.", run_path) + try: + lhe_path = npy_to_lhe.convert(run_path) + except Exception as error: + _report_failure(log, "npy->LHE conversion", error, run_path) + return lhe_path + + +def _report_failure(log, what, error, directory=None): + """Log a post-processing failure and write the full traceback to a file + (whose path is printed) so the problem can be investigated.""" + import traceback + if not directory or not os.path.isdir(directory): + directory = os.getcwd() + path = os.path.join(directory, "%s_crash.log" % what.replace(" ", "_")) + try: + with open(path, "w") as fsock: + fsock.write(traceback.format_exc()) + except Exception: + path = None + log.warning("%s failed: %s", what, error) + if path: + log.warning("full traceback written to: %s", path) + + +def _off(value) -> bool: + """True when a switch value means the tool was not selected.""" + return value in (None, "OFF", "Not Avail.", "Not Avail. (numpy missing)") + + +def _has_effective_handler(lg) -> bool: + """True when ``lg`` (or an ancestor it propagates to) already has a handler, + i.e. some other party is already displaying these records.""" + while lg: + if lg.handlers: + return True + if not lg.propagate: + return False + lg = lg.parent + return False + + +_TOOL_LOGGING_READY = False + + +def _setup_logging(): + """Make the reused madevent tool drivers' progress visible on screen. + + The drivers already narrate what they do through logger.info / + update_status ("Running Pythia8 [arXiv:...]", "Splitting .lhe event + file...", the live Idle/Running/Completed job counters, "Running + MadSpin", ...), but the standalone mg7 launcher never attaches a handler + to the 'madgraph'/'madevent' loggers, so all of it is swallowed. Attach a + colored INFO console handler to them (idempotent).""" + global _TOOL_LOGGING_READY + if _TOOL_LOGGING_READY: + return + if _has_effective_handler(logging.getLogger("madgraph")): + # Someone already configured logging for us -- in practice MG5, which + # now runs the launch in its own process. Its configuration is + # authoritative (it owns stdout_level and the tutorial logger); adding + # a second handler here would print every line twice. + _TOOL_LOGGING_READY = True + return + try: + import madgraph.interface.coloring_logging # registers ColorFormatter + formatter = logging.ColorFormatter("%(message)s") + except Exception: + formatter = logging.Formatter("%(message)s") + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter(formatter) + handler._mg7_tool_handler = True + for name in ("madgraph", "madevent", "cmdprint", "madgraph7"): + lg = logging.getLogger(name) + if not any(getattr(h, "_mg7_tool_handler", False) for h in lg.handlers): + lg.addHandler(handler) + lg.setLevel(logging.INFO) + lg.propagate = False + _TOOL_LOGGING_READY = True + + +def selected_tools(switch) -> list: + """The post-processing drivers run_selected_tools runs for this switch.""" + # MadAnalysis5 hadron level analyses the shower/detector output, so it only + # makes sense when a shower ran (mirrors madevent's card gating: + # analysis == 'MadAnalysis5' and shower != 'OFF'). + ma5 = switch.get("analysis") == "MadAnalysis5" + showered = not _off(switch.get("shower")) + return [t for t, on in ( + ("reweighting", not _off(switch.get("reweight"))), + ("MadSpin", not _off(switch.get("madspin"))), + ("MadAnalysis5 (parton level)", ma5), + ("Pythia8 shower", switch.get("shower") == "Pythia8"), + ("Delphes", switch.get("detector") == "Delphes"), + ("MadAnalysis5 (hadron level)", ma5 and showered), + ("Rivet", switch.get("analysis") == "Rivet"), + ) if on] + + +def run_selected_tools(switch, process) -> None: + """Run the optional post-processing programs selected in the merged + question on the generated events. + + All the tools are driven through :class:`MG7RunCmd`, a thin adapter around + the standard madevent run interface (``MadEventCmd``). This means the mg7 + output reuses the *exact* madevent tool drivers -- in particular Pythia8 is + showered in parallel over the cluster/multicore backend, identically to a + madevent run -- instead of re-implementing each tool here. The command + sequence mirrors ``MadEventCmd.do_launch`` (reweight, MadSpin, MA5 parton, + shower, Delphes, MA5 hadron, Rivet); each tool checks its own card, so a + tool is invoked only when its switch is on. + """ + log = logging.getLogger("madevent") + + tools = selected_tools(switch) + if not tools: + # A switch that is not off does not necessarily select a driver here: + # "Not Avail." is not off, analysis = ExRoot has no mg7 driver, and a + # shower switch set to something other than Pythia8 selects nothing + # either. Without this a plain generate/output/launch announced a + # post-processing step with nothing after the colon, and paid for + # building the run interface (or an npy->LHE conversion) to do nothing. + return + + lhe_path = _find_or_convert_event_file(process.run_path, log) + if lhe_path is None: + log.warning("No LHE event file in %s; cannot run %s.", + process.run_path, ", ".join(tools)) + return + + run_name = os.path.basename(os.path.dirname(os.path.abspath(lhe_path))) + run_dir = os.path.dirname(os.path.abspath(lhe_path)) + + ma5 = switch.get("analysis") == "MadAnalysis5" + showered = not _off(switch.get("shower")) + + log.info("") + log.info("Post-processing the generated events with: %s", ", ".join(tools)) + + try: + from madgraph.iolibs.template_files.mg7.run_interface import MG7RunCmd + cmd = MG7RunCmd(os.getcwd(), load_mg5_options(), run_name, lhe_path) + cmd.set_run_name(run_name, None, "parton") + except Exception as error: + _report_failure(log, "post-processing setup", error, run_dir) + return + + def run(tool, command): + bar = "=" * 60 + log.info("") + log.info(bar) + log.info(" post-processing step: %s", tool) + log.info(bar) + start = time.time() + try: + cmd.exec_cmd(command, postcmd=False, printcmd=False) + except Exception as error: + _report_failure(log, tool, error, run_dir) + else: + log.info(" -> %s done (%.1fs)", tool, time.time() - start) + + # Same order as MadEventCmd.do_launch. The commands take no run name and so + # act on cmd.run_name: reweight runs on the generated events, then + # decay_events (MadSpin) repoints cmd.run_name to the "_decayed_i" run, + # so the shower and the analyses that follow act on the decayed events -- + # exactly as a madevent run does. Each driver still guards on its own card. + if not _off(switch.get("reweight")): + run("reweight", "reweight -from_cards") + if not _off(switch.get("madspin")): + run("MadSpin", "decay_events -from_cards") + if ma5: + run("MadAnalysis5 (parton)", "madanalysis5_parton --no_default") + if switch.get("shower") == "Pythia8": + run("Pythia8 shower", "shower --no_default") + if switch.get("detector") == "Delphes": + run("Delphes", "delphes --no_default") + # hadron-level MA5 needs the shower/detector output: only run it if a shower + # was requested (otherwise there is no hadron-level event file to analyse). + if ma5 and showered: + run("MadAnalysis5 (hadron)", "madanalysis5_hadron --no_default") + if switch.get("analysis") == "Rivet": + run("Rivet", "rivet --no_default") + + # Finalize the run exactly as MadEventCmd.do_launch does after the + # shower/detector/analysis tools: store_result() processes the deferred + # "to_store" actions -- in particular it gzips the Pythia8 HepMC to + # _pythia8_events.hepmc.gz (the pythia8_card default is HEPMCoutput:file + # = hepmc.gz). The deferred Rivet job reads exactly that .gz path, so this + # must run before the Rivet/Contur postprocessor below. + try: + cmd.store_result() + except Exception as error: + _report_failure(log, "store_result", error, run_dir) + + if switch.get("analysis") == "Rivet": + # do_rivet only *prepared* the run (wrote Events//run_rivet.sh) and + # deferred execution to the postprocessor -- the rivet_card default has + # run_rivet_later = True, so with --no_default it logged "Skipping Rivet + # for now, passing it to postprocessor" and appended the run to + # cmd.postprocessing_dirs. MadEventCmd.do_launch runs that postprocessor + # at the end of a run; do the same here so the .yoda (and, when + # run_contur = True, the Contur limits) are actually produced. + bar = "=" * 60 + log.info("") + log.info(bar) + log.info(" post-processing step: Rivet/Contur postprocessing") + log.info(bar) + start = time.time() + try: + cmd.postprocessing() + except Exception as error: + _report_failure(log, "Rivet-Contur postprocessing", error, run_dir) + else: + log.info(" -> Rivet/Contur postprocessing done (%.1fs)", + time.time() - start) + + +def _add_time_of_flight(lhe_path, threshold, param_card_path, log): + """Add invariant-lifetime (vtim) information to the LHE events, drawing each + unstable particle's decay length from its width (in mm). Mirrors + common_run_interface.CommonRunCmd.do_add_time_of_flight but reads the widths + from the local param_card instead of the LHE banner.""" + import random + try: + import madgraph.various.lhe_parser as lhe_parser + except ImportError: + import internal.lhe_parser as lhe_parser + from models import check_param_card as param_card_mod + from madgraph.various import misc as _misc + import madgraph.iolibs.files as files + + need_zip = lhe_path.endswith('.gz') + if need_zip: + _misc.gunzip(lhe_path) + lhe_path = lhe_path[:-3] + + param_card = param_card_mod.ParamCard(param_card_path) + cst = 6.58211915e-25 # hbar in GeV s + c = 299792458000 # speed of light in mm/s + log.info('Adding time of flight information on %s', lhe_path) + lhe = lhe_parser.EventFile(lhe_path) + out = open('%s_2vertex.lhe' % lhe_path, 'w') + out.write(lhe.banner) + for event in lhe: + for particle in event: + # default=0 -> particles without a decay entry are treated as stable + width = param_card['decay'].get((abs(particle.pid),), 0.).value + if width: + vtim = c * random.expovariate(width / cst) + if vtim > threshold: + particle.vtim = vtim + out.write(str(event)) + out.write('\n') + out.close() + lhe.close() + files.mv('%s_2vertex.lhe' % lhe_path, lhe_path) + if need_zip: + _misc.gzip(lhe_path) + + +def _lhapdf_config_path(): + """Path to lhapdf-config so systematics can import the python lhapdf + module (required to compute PDF/scale variations).""" + return lhapdf_paths().config + + +def _run_systematics(lhe_path, cfg, log): + """Run systematics.py (scale/PDF variations) on the LHE file, reusing the + exact worker madevent uses (systematics.call_systematics).""" + try: + import madgraph.various.systematics as systematics + except ImportError: + import internal.systematics as systematics + + def fmt(vals): + return ','.join(str(v) for v in vals) + + opts = [] + if cfg.get('systematics_mur'): + opts.append('--mur=%s' % fmt(cfg['systematics_mur'])) + if cfg.get('systematics_muf'): + opts.append('--muf=%s' % fmt(cfg['systematics_muf'])) + if cfg.get('systematics_pdf'): + opts.append('--pdf=%s' % fmt(cfg['systematics_pdf'])) + extra = cfg.get('systematics_str_options', '') + if extra: + opts.extend(extra.split()) + # The mg7 LHE has no block, so systematics cannot read the per-event + # LO reweighting info. For a process with a single alpha_s power (stored in + # SubProcesses/proc_characteristics at output time) it can be reconstructed + # from the events: pass that power as --lo_nqcd. -1 means the QCD power is + # not uniform, so the reconstruction is not applicable. + if not any(o.startswith('--lo_nqcd') for o in opts): + try: + from madgraph.various import banner as _banner_mod + except ImportError: + import internal.banner as _banner_mod + pc_path = os.path.join('SubProcesses', 'proc_characteristics') + if os.path.exists(pc_path): + try: + nqcd = int(_banner_mod.ProcCharacteristic(pc_path)['single_qcd_order']) + except Exception: + nqcd = -1 + if nqcd >= 0: + opts.append('--lo_nqcd=%d' % nqcd) + # tell systematics where to find lhapdf (so it can link the python module) + if not any(o.startswith('--lhapdf_config') for o in opts): + lhapdf_config = _lhapdf_config_path() + if lhapdf_config: + opts.append('--lhapdf_config=%s' % lhapdf_config) + + log.info('Running systematics on %s %s', lhe_path, ' '.join(opts)) + systematics.call_systematics([lhe_path, lhe_path] + opts, + log=lambda x: log.info(str(x))) + + +def run_lhe_postprocessing(process) -> None: + """Run the LHE-level post-processings configured in the run_card + [postprocessing] section on the generated event file (displaced-vertex + time-of-flight and systematics). Only applies when an LHE file exists.""" + lhe_path = _find_event_file(process.run_path) + if lhe_path is None: + return + # process.run_card is a RunCardMG7; the section view exposes .get(key, def) + try: + cfg = process.run_card["postprocessing"] + except Exception: + return + log = logging.getLogger('madevent') + + if cfg.get('systematics') and getattr(process, 'is_decay', False): + # Scale and PDF variations are a beam quantity. A decay has neither, so + # systematics can only fail here ("not supported for pdlabel=none") -- + # and it is not free: MadSpin reruns the launcher for every pool refill. + log.info("decay mode: skipping systematics (no beams to vary)") + elif cfg.get('systematics'): + native = getattr(process, 'systematics', None) + if native is not None and native.weight_count > 0: + log.warning("[postprocessing] systematics is ignored: the scale/PDF " + "weights were already computed by madspace ([systematics] " + "section); set systematics.enable = false to use the legacy " + "systematics.py path instead.") + else: + try: + _run_systematics(lhe_path, cfg, log) + except Exception as error: + _report_failure(log, "systematics computation", error, + os.path.dirname(lhe_path)) + + tof = cfg.get('time_of_flight', -1.0) + try: + if tof is not None and float(tof) >= 0: + _add_time_of_flight(lhe_path, float(tof), process.param_card_path, log) + except Exception as error: + _report_failure(log, "add_time_of_flight", error, + os.path.dirname(lhe_path)) diff --git a/madgraph/iolibs/template_files/mg7/run_interface.py b/madgraph/iolibs/template_files/mg7/run_interface.py index 4ba4a80ade..bd71bbd364 100644 --- a/madgraph/iolibs/template_files/mg7/run_interface.py +++ b/madgraph/iolibs/template_files/mg7/run_interface.py @@ -198,7 +198,7 @@ def get_model(self): tools this class drives (MadSpin, reweight, the auto widths and the 'update dependent' pass of their card questions) work here too. """ - from madgraph.iolibs.template_files.mg7.launch import load_process_model + from madgraph.iolibs.template_files.mg7.model_widths import load_process_model return load_process_model(self.me_dir) def do_set(self, line, log=True): diff --git a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py new file mode 100644 index 0000000000..8352894842 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py @@ -0,0 +1,1131 @@ +"""Subprocesses with madmatrix matrix elements and MG diagram based phase space.""" + +from __future__ import annotations + +import os +import re +from dataclasses import dataclass, field +from typing import Literal, NamedTuple, TYPE_CHECKING + +from madgraph.iolibs.template_files.mg7.common import logger, me_parameters, ms +from madgraph.iolibs.template_files.mg7.custom_function_loader import load_surrogate +from madgraph.iolibs.template_files.mg7.observables import build_observable, clean_pids, decay_products, observable_pids, pid_is_qcd + +if TYPE_CHECKING: + from madgraph.iolibs.template_files.mg7.madgraph_process import MadgraphProcess + +@dataclass +class Channel: + phasespace_mapping: ms.PhaseSpaceMapping + adaptive_mapping: ms.Flow | ms.VegasMapping + discrete_sym: ms.DiscreteSampler | ms.DiscreteFlow | None + discrete_flavor: ms.DiscreteSampler | ms.DiscreteFlow | None + channel_weight_indices: list[int] | None + name: str + active_flavors: list[int] + event_generator: ms.ChannelEventGenerator | None = None + + +@dataclass +class PhaseSpace: + mode: Literal["multichannel", "flat", "both"] + channels: list[Channel] + symfact: list[int | None] + first_chan_weight_remap: list[list[int]] = field(default_factory=list) + first_remapped_chan_count: int = 0 + second_chan_weight_remap: list[int] = field(default_factory=list) + second_remapped_chan_count: int = 0 + prop_chan_weights: ms.PropagatorChannelWeights | None = None + subchan_weights: ms.SubchannelWeights | None = None + cwnet: ms.ChannelWeightNetwork | None = None + + +class MultiChannelData(NamedTuple): + amp2_remaps: list[list[int]] + symfact: list[int | None] + topologies: list[list[ms.Topology]] + permutations: list[list[list[int]]] + channel_indices: list[list[int]] + channel_weight_indices: list[list[list[int]]] + diagram_indices: list[list[int]] + diagram_color_indices: list[list[list[int]]] + diagram_propagator_pdgs: list[list[list[int]]] + active_flavors: list[list[list[int]]] + qcd_s_channel_count: list[int] + + +def build_topologies( + incoming_masses: list[float], + outgoing_masses: list[float], + channel: dict, + process: MadgraphProcess +) -> list[ms.Topology]: + propagators = [] + for i, (pid, signed_pid) in enumerate(zip( + clean_pids(channel["propagators"]), channel["propagators"] + )): + mass = process.get_mass(pid) + width = process.get_width(pid) + bw_cutoff = process.run_card["phasespace"]["bw_cutoff"] + if i in channel["on_shell_propagators"]: + e_min = mass - bw_cutoff * width + e_max = mass + bw_cutoff * width + else: + e_min = 0 + e_max = 0 + # $-excluded propagator: the matrix element vanishes within bw_cutoff + # widths of the pole, so do not pile the channel's points up there + # (a flat density in that window instead of the Breit-Wigner peak). + # Channels written before the exporter recorded them have none. + flat_window = 0. + if i in channel.get("dollar_propagators", ()) and width > 0: + flat_window = float(bw_cutoff) + propagators.append(ms.Propagator( + mass=mass, + width=width, + integration_order=0, + e_min=e_min, + e_max=e_max, + pdg_id=signed_pid, + flat_window=flat_window, + )) + vertices = channel["vertices"] + diag = ms.Diagram( + incoming_masses, outgoing_masses, propagators, vertices + ) + return ms.Topology.topologies(diag) + + +def build_multi_channel_data( + meta: dict, process: MadgraphProcess, unmerged_meta: dict | None = None +) -> MultiChannelData: + incoming_masses = [ + process.get_mass(pid) for pid in clean_pids(meta["incoming"]) + ] + outgoing_masses = [ + process.get_mass(pid) for pid in clean_pids(meta["outgoing"]) + ] + + if unmerged_meta is None: + diagram_count = meta["diagram_count"] + amp2_remaps = [[-1] * diagram_count] + else: + amp2_remaps = [ + [-1] * unmerged_meta[subproc]["diagram_count"] + for subproc in meta["subprocesses"] + ] + symfact = [] + topologies = [] + permutations = [] + channel_indices = [] + channel_weight_indices = [] + diagram_indices = [] + diagram_color_indices = [] + diagram_propagator_pdgs = [] + active_flavors = [] + channel_index = 0 + qcd_s_channel_count = [] + + for channel in meta["channels"]: + if unmerged_meta is None: + topo_channel = channel + else: + topo_subproc = channel["subprocess"] + topo_channel_index = channel["channel"] + topo_channel = unmerged_meta[topo_subproc]["channels"][topo_channel_index] + chan_topologies = build_topologies( + incoming_masses, outgoing_masses, topo_channel, process + ) + topo_count = len(chan_topologies) + if topo_count == 0: + continue + + topo = chan_topologies[0] + s_chan_count = [0] * len(topo.decays) + non_qcd = [False] * len(topo.decays) + pdg_ids = [decay.pdg_id for decay in topo.decays] + for i, pid in zip(topo.outgoing_indices, meta["outgoing"]): + pdg_ids[i] = pid + for decay in reversed(topo.decays): + if len(decay.child_indices) == 0: + continue + if decay.index == 0 and topo.t_propagator_count > 0: + s_chan_count[0] = sum(s_chan_count[i] for i in decay.child_indices) + break + + is_qcd = pid_is_qcd(pdg_ids[decay.index]) and all( + pid_is_qcd(pdg_ids[i]) for i in decay.child_indices + ) + is_non_qcd = decay.on_shell or not is_qcd or any( + non_qcd[i] for i in decay.child_indices + ) + non_qcd[decay.index] = is_non_qcd + s_chan_count[decay.index] = ( + 0 if is_non_qcd else sum(s_chan_count[i] for i in decay.child_indices) + 1 + ) + qcd_s_channel_count.append(s_chan_count[0]) + + diagrams = channel["diagrams"] + chan_permutations = [d["permutation"] for d in diagrams] + if unmerged_meta is None: + amp2_remaps[0][diagrams[0]["diagram"]] = channel_index + else: + for amp2_remap, diag in zip(amp2_remaps, diagrams[0]["diagram"]): + if diag != -1: + amp2_remap[diag] = channel_index + + channel_index_first = channel_index + symfact_index_first = len(symfact) + channel_index += 1 + symfact.extend([None] * topo_count) + for d in diagrams[1:]: + if unmerged_meta is None: + amp2_remaps[0][d["diagram"]] = channel_index + else: + for amp2_remap, diag in zip(amp2_remaps, d["diagram"]): + if diag != -1: + amp2_remap[diag] = channel_index + channel_index += 1 + symfact.extend(range(symfact_index_first, symfact_index_first + topo_count)) + + topologies.append(chan_topologies) + permutations.append(chan_permutations) + channel_indices.append(list(range(channel_index_first, channel_index))) + channel_weight_indices.append([ + [ + symfact_index_first + topo_index + i * topo_count + for i in range(len(chan_permutations)) + ] + for topo_index in range(topo_count) + ]) + diagram_indices.append([d["diagram"] for d in diagrams]) + if unmerged_meta is None: + diagram_color_indices.append([d["active_colors"] for d in diagrams]) + diagram_propagator_pdgs.append( + [d["propagator_pdgs"] for d in diagrams] + ) + active_flavors.append([d["active_flavors"] for d in diagrams]) + + return MultiChannelData( + amp2_remaps, + symfact, + topologies, + permutations, + channel_indices, + channel_weight_indices, + diagram_indices, + diagram_color_indices, + diagram_propagator_pdgs, + active_flavors, + qcd_s_channel_count, + ) + + +class MadgraphSubprocess: + def __init__( + self, + process: MadgraphProcess, + meta: dict, + subproc_id: int, + unmerged_meta: dict | None = None + ): + self.process = process + self.meta = meta + self.subproc_id = subproc_id + self.multi_channel_data = None + + self.unmerged_meta = None + if unmerged_meta is None: + api_path_formats = [self.meta["me_path"]] + member_metas = [self.meta] + else: + api_path_formats = [] + member_metas = [] + for subproc in self.meta["subprocesses"]: + submeta = unmerged_meta[subproc] + api_path_formats.append(submeta["me_path"]) + member_metas.append(submeta) + if len(api_path_formats) == 1: + self.meta = submeta + else: + self.unmerged_meta = unmerged_meta + + # The libraries were all built up front by MadgraphProcess.compile_matrix_elements + all_api_paths = [ + [api_path_format.format(device=backend) for backend in self.process.backends] + for api_path_format in api_path_formats + ] + + self.incoming_masses = [ + self.process.get_mass(pid) for pid in clean_pids(self.meta["incoming"]) + ] + self.outgoing_masses = [ + self.process.get_mass(pid) for pid in clean_pids(self.meta["outgoing"]) + ] + self.particle_count = len(self.incoming_masses) + len(self.outgoing_masses) + # MadEvent's cut_decays = F: the decay products of on-shell propagators + # are left out of every cut selection (a pair cut then needs both + # members to be cut), while the histograms still see them + if self.process.run_card["phasespace"]["cut_decays"]: + uncut = set() + else: + uncut = decay_products(self.meta, self.unmerged_meta) + n_in = len(self.incoming_masses) + + def cut_pids(observable_kwargs): + pids = observable_pids(self.meta, observable_kwargs, self.unmerged_meta) + # pdg id 0 is in no group, so the particle is never selected + return [0 if i - n_in in uncut else pid for i, pid in enumerate(pids)] + + # custom cuts and scales see the PDG ids of the first flavor, like the + # built-in cuts, which select by flavor-independent groups + representative_pids = [self.flavor_pids()[0]] + custom_cut = self.process.custom_function( + "cut", self.particle_count, representative_pids + ) + if len(self.process.cut_data) > 0: + self.cuts = ms.Cuts([ + ms.CutItem( + observable=build_observable( + cut_pids(cut_item.observable_kwargs), + n_in, cut_item.observable_kwargs + ), + min=cut_item.min, + max=cut_item.max, + mode=cut_item.mode, + ) + for cut_item in self.process.cut_data + ], custom_cut) + elif custom_cut is not None: + self.cuts = ms.Cuts(self.particle_count, custom_cut) + else: + self.cuts = None + # the integration histograms are functions of the momenta; the weight + # distribution is a property of the final event sample, so it is only + # filled by MadgraphProcess.build_event_histograms at combine time + # They cost an observable evaluation per integration point, and the + # plots/HwU are drawn from the post-processing ones instead, so they + # are only filled on request ([run] weighted_histograms). + momentum_hists = [item for item in self.process.hist_data + if not item.from_weight] \ + if self.process.run_card["run"]["weighted_histograms"] else [] + self.histograms = ( + ms.ObservableHistograms([ + ms.HistItem( + observable=build_observable( + observable_pids( + self.meta, hist_item.observable_kwargs, self.unmerged_meta + ), + n_in, hist_item.observable_kwargs + ), + min=hist_item.min, + max=hist_item.max, + bin_count=hist_item.bin_count, + ) + for hist_item in momentum_hists + ]) + if momentum_hists + else None + ) + + self.scale = ms.EnergyScale( + particle_count=self.particle_count, + custom_scale=self.process.custom_function( + "scale", self.particle_count, representative_pids + ), + **self.process.scale_kwargs + ) + + if self.process.run_card["run"]["dummy_matrix_element"]: + self.matrix_elements = [None] * len(all_api_paths) + else: + self.matrix_elements = [] + for api_paths in all_api_paths: + for context, api_path in zip(self.process.contexts, api_paths): + mat = context.load_matrix_element( + api_path, self.process.param_card_path, + me_parameters(self.process.run_card) + ) + self.matrix_elements.append(mat) + self.surrogates = [] + if self.process.surrogate is not None: + for path_format, member_meta in zip(api_path_formats, member_metas): + self.surrogates.append(self.load_surrogate(path_format, member_meta)) + + def load_surrogate(self, me_path_format: str, member_meta: dict): + """The [surrogate] of one matrix element of the subprocess.""" + name = re.sub( + r"^libmadmatrix_(.*)_\{device\}\.so$", r"\1", os.path.basename(me_path_format) + ) + # PDG ids for each flavor index of the matrix element + flavors = member_meta["flavors"] + pdg_ids = [flavors[0]["options"][0]] * (max(f["index"] for f in flavors) + 1) + for flavor in flavors: + pdg_ids[flavor["index"]] = flavor["options"][0] + settings = self.process.surrogate + apis = load_surrogate( + ms, + self.process.contexts, + self.process.backends, + self.process.device_types, + settings["function"], + name, + pdg_ids, + settings["param_card"], + ) + if len({api.index() for api in apis}) != 1: + raise RuntimeError("surrogate loaded at different indices") + return apis[0] + + def flavor_pids(self) -> list[list[int]]: + """PDG ids of the external particles for every flavor of the + integrand, in the order of its flavor index.""" + pids = [] + for flav in self.meta["flavors"]: + if self.unmerged_meta is not None: + subproc_index = self.meta["subprocesses"][flav["subprocess"]] + flav = self.unmerged_meta[subproc_index]["flavors"][flav["flavor"]] + pids.append(flav["options"][0]) + return pids + + def build_multi_channel_data(self) -> MultiChannelData: + if self.multi_channel_data is not None: + return self.multi_channel_data + self.multi_channel_data = build_multi_channel_data( + self.meta, self.process, self.unmerged_meta + ) + return self.multi_channel_data + + def build_multichannel_phasespace(self) -> PhaseSpace: + mcdata = self.build_multi_channel_data() + channel_count = sum(len(topos) for topos in mcdata.topologies) + drop_threshold = self.process.run_card["phasespace"]["drop_qcd_s_channel"] + if drop_threshold >= 0 and channel_count > drop_threshold: + mcdata = self.drop_qcd_s_channels(mcdata) + mcdata = self.drop_empty_channels(mcdata) + + channels = [] + t_channel_mode = self.t_channel_mode( + self.process.run_card["phasespace"]["t_channel"] + ) + for channel_id, (chan_topologies, chan_permutations, chan_indices, active_flavors) in enumerate(zip( + mcdata.topologies, mcdata.permutations, mcdata.channel_weight_indices, + mcdata.active_flavors + )): + topo_count = len(chan_topologies) + for topo_index, (topo, indices) in enumerate(zip(chan_topologies, chan_indices)): + mapping = ms.PhaseSpaceMapping( + chan_topologies[0], + self.process.e_cm, + t_channel_mode=t_channel_mode, + cuts=self.cuts, + invariant_power=self.process.run_card["phasespace"]["invariant_power"], + permutations=chan_permutations, + leptonic=self.process.leptonic, + ) + prefix = f"subproc{self.subproc_id}.channel{channel_id}" + if topo_count > 1: + prefix += f".subchan{topo_index}" + discrete_sym, discrete_flavor = self.build_discrete( + len(chan_permutations), len(self.meta["flavors"]), prefix + ) + channels.append(Channel( + phasespace_mapping = mapping, + adaptive_mapping = self.build_vegas(mapping, prefix), + discrete_sym = discrete_sym, + discrete_flavor = discrete_flavor, + channel_weight_indices = indices, + name = f"{channel_id}", + active_flavors = active_flavors, + )) + + remapped_chan_count = sum( + len(indices) for indices in mcdata.channel_indices + ) + if self.process.run_card["phasespace"]["sde_strategy"] == "denominators": + prop_chan_weights = ms.PropagatorChannelWeights( + [topo[0] for topo in mcdata.topologies], mcdata.permutations, + mcdata.channel_indices + ) + chan_weight_remap = [] + else: + prop_chan_weights = None + chan_weight_remap = [ + [ + len(mcdata.symfact) if remap == -1 else remap + for remap in amp2_remap + ] + for amp2_remap in mcdata.amp2_remaps + ] + + if any(len(topos) > 1 for topos in mcdata.topologies): + subchan_weights = ms.SubchannelWeights( + mcdata.topologies, mcdata.permutations, mcdata.channel_indices + ) + else: + subchan_weights = None + + return PhaseSpace( + mode="multichannel", + channels=channels, + first_chan_weight_remap=chan_weight_remap, + first_remapped_chan_count=remapped_chan_count, + symfact=mcdata.symfact, + prop_chan_weights=prop_chan_weights, + subchan_weights=subchan_weights, + ) + + def build_flat_phasespace(self) -> PhaseSpace: + mapping = ms.PhaseSpaceMapping( + self.incoming_masses + self.outgoing_masses, + self.process.e_cm, + mode=self.t_channel_mode(self.process.run_card["phasespace"]["flat_mode"]), + cuts=self.cuts, + leptonic=self.process.leptonic, + ) + prefix = f"subproc{self.subproc_id}.flat" + discrete_sym, discrete_flavor = self.build_discrete( + 1, len(self.meta["flavors"]), prefix + ) + channel = Channel( + phasespace_mapping = mapping, + adaptive_mapping = self.build_vegas(mapping, prefix), + discrete_sym = discrete_sym, + discrete_flavor = discrete_flavor, + channel_weight_indices = [0], + name = "F", + active_flavors = [], + ) + if self.unmerged_meta is None: + remap = [list(range(self.meta["diagram_count"]))] + else: + remap = [ + list(range(self.unmerged_meta[subproc]["diagram_count"])) + for subproc in self.meta["subprocesses"] + ] + return PhaseSpace( + mode="flat", + channels=[channel], + first_chan_weight_remap=remap, + first_remapped_chan_count=1, + symfact=[None], + ) + + def simplify_phasespace( + self, + multi_phasespace: PhaseSpace, + cross_sections: list[float] + ) -> PhaseSpace | None: + assert multi_phasespace.mode == "multichannel" + + threshold = 1 - self.process.run_card["phasespace"]["combine_channel_threshold"] + kept_channels = [] + tot_cs = sum(cross_sections) + cum_cs = 0. + seen_active_flavors = set() + #seen_resonances = set() + for index, (cs, chan) in sorted( + enumerate(zip(cross_sections, multi_phasespace.channels)), + key=lambda pair: pair[1][0], + reverse=True + ): + cum_cs += cs + has_unseen_flavors = False + has_unseen_resonances = False + for flavs in chan.active_flavors: + for flav in flavs: + if flav not in seen_active_flavors: + has_unseen_flavors = True + seen_active_flavors.add(flav) + #for resonance in chan.resonances: + # if resonance not in seen_resonances: + # has_unseen_resonances = True + # seen_resonances.add(flav) + if has_unseen_flavors or has_unseen_resonances or not tot_cs \ + or cum_cs / tot_cs < threshold: + kept_channels.append(index) + if len(kept_channels) >= len(cross_sections) - 1: + return multi_phasespace + + channels = [] + channel_map = {} + symfact = [] + for old_chan_index in kept_channels: + channel = multi_phasespace.channels[old_chan_index] + perm_count = max(1, channel.phasespace_mapping.channel_count()) + channel_index = len(symfact) + symfact.append(None) + symfact.extend([channel_index] * (perm_count - 1)) + channel_map.update({ + old_index: new_index + for new_index, old_index in enumerate( + channel.channel_weight_indices, start=channel_index + ) + }) + channels.append(Channel( + phasespace_mapping = channel.phasespace_mapping, + adaptive_mapping = channel.adaptive_mapping, + discrete_sym = channel.discrete_sym, + discrete_flavor = channel.discrete_flavor, + channel_weight_indices = list(range( + channel_index, channel_index + perm_count + )), + name = channel.name, + active_flavors = channel.active_flavors, + event_generator = channel.event_generator, + )) + + flat_phasespace = self.build_flat_phasespace() + flat_channel = flat_phasespace.channels[0] + channels.append(Channel( + phasespace_mapping = flat_channel.phasespace_mapping, + adaptive_mapping = flat_channel.adaptive_mapping, + discrete_sym = flat_channel.discrete_sym, + discrete_flavor = flat_channel.discrete_flavor, + channel_weight_indices = [len(symfact)], + name = flat_channel.name, + active_flavors = flat_channel.active_flavors, + )) + flat_index = len(symfact) + symfact.append(None) + channel_map[len(multi_phasespace.symfact)] = len(symfact) + if multi_phasespace.subchan_weights is None and len(multi_phasespace.first_chan_weight_remap) > 0: + first_chan_weight_remap = [ + [ + channel_map.get(remap, flat_index) + for remap in cw_remap + ] + for cw_remap in multi_phasespace.first_chan_weight_remap + ] + first_remapped_chan_count = len(symfact) + second_chan_weight_remap = [] + second_remapped_chan_count = 0 + else: + first_chan_weight_remap = multi_phasespace.first_chan_weight_remap + first_remapped_chan_count = multi_phasespace.first_remapped_chan_count + chan_count = multi_phasespace.first_remapped_chan_count if multi_phasespace.subchan_weights is None else multi_phasespace.subchan_weights.channel_count() + second_chan_weight_remap = [ + channel_map.get(i, flat_index) + for i in range(chan_count) + ] + second_remapped_chan_count = len(symfact) + + return PhaseSpace( + mode="both", + channels=channels, + first_chan_weight_remap=first_chan_weight_remap, + first_remapped_chan_count=first_remapped_chan_count, + second_chan_weight_remap=second_chan_weight_remap, + second_remapped_chan_count=second_remapped_chan_count, + symfact=symfact, + prop_chan_weights=multi_phasespace.prop_chan_weights, + subchan_weights=multi_phasespace.subchan_weights, + ) + + def drop_qcd_s_channels(self, mcdata: MultiChannelData) -> MultiChannelData: + """Drop channels with non-resonant QCD s-channel propagators, to reduce the + channel count for processes with many diagrams. Channels are grouped by the + number of QCD s-channels, allowing for more if necessary to map out all flavor + indices. Channel weights belonging to a dropped channel are left unmapped.""" + groups_by_s_count = {} + for index, count in enumerate(mcdata.qcd_s_channel_count): + groups_by_s_count.setdefault(count, []).append(index) + + covered_flavors = set() + kept_groups = set() + for s_count in sorted(groups_by_s_count): + s_count_groups = groups_by_s_count[s_count] + if s_count == 0: + selected = s_count_groups + else: + selected = [ + index + for index in s_count_groups + if any( + flav not in covered_flavors + for flavs in mcdata.active_flavors[index] + for flav in flavs + ) + ] + kept_groups.update(selected) + for index in selected: + for flavs in mcdata.active_flavors[index]: + covered_flavors.update(flavs) + + return self.keep_channel_groups(mcdata, sorted(kept_groups)) + + def drop_empty_channels(self, mcdata: MultiChannelData) -> MultiChannelData: + """Drop the channels whose phase space the cuts exclude entirely: a pair + mass cut whose floor lies at or above the top of an on-shell window of + the channel (ms.PhaseSpaceMapping.empty). Every point such a channel + generates fails the cuts, so it could only ever fail to find one. + + If that leaves a flavour with no channel at all, its cross section is + zero -- typically a cut that excludes the window of a decay chain every + diagram shares -- and the run cannot sample it: that is reported as an + error naming the cause, rather than as the survey finding no point + passing the cuts.""" + if self.cuts is None: + return mcdata + t_channel_mode = self.t_channel_mode( + self.process.run_card["phasespace"]["t_channel"] + ) + empty = [ + ms.PhaseSpaceMapping( + chan_topologies[0], + self.process.e_cm, + t_channel_mode=t_channel_mode, + cuts=self.cuts, + invariant_power=self.process.run_card["phasespace"]["invariant_power"], + permutations=chan_permutations, + leptonic=self.process.leptonic, + ).empty() + for chan_topologies, chan_permutations in zip( + mcdata.topologies, mcdata.permutations + ) + ] + if not any(empty): + return mcdata + kept_groups = [index for index, is_empty in enumerate(empty) if not is_empty] + covered_flavors = { + flav + for index in kept_groups + for flavs in mcdata.active_flavors[index] + for flav in flavs + } + uncovered = sorted({ + flav + for index, is_empty in enumerate(empty) if is_empty + for flavs in mcdata.active_flavors[index] + for flav in flavs + } - covered_flavors) + if uncovered: + raise ValueError( + f"subprocess {self.meta['incoming']} > {self.meta['outgoing']}: " + "the pair mass cuts exclude the whole on-shell window " + f"(bw_cutoff = {self.process.run_card['phasespace']['bw_cutoff']}) " + f"of every channel of flavour(s) {uncovered}, so their cross " + "section is zero; loosen the cut or remove that decay" + ) + logger.info( + "subprocess %d: %d channel(s) dropped, the cuts exclude the " + "on-shell window they sample", self.subproc_id, len(empty) - len(kept_groups) + ) + return self.keep_channel_groups(mcdata, kept_groups) + + def keep_channel_groups( + self, mcdata: MultiChannelData, kept_groups: list[int] + ) -> MultiChannelData: + """Rebuild the multi-channel data with only the channel groups (indices + into mcdata.topologies) in kept_groups. Channel weights belonging to a + dropped channel are left unmapped.""" + if len(kept_groups) == len(mcdata.topologies): + return mcdata + + amp2_remaps = [[-1] * len(remap) for remap in mcdata.amp2_remaps] + symfact = [] + topologies = [] + permutations = [] + channel_indices = [] + channel_weight_indices = [] + diagram_indices = [] + diagram_color_indices = [] + diagram_propagator_pdgs = [] + active_flavors = [] + qcd_s_channel_count = [] + channel_index = 0 + + for group in kept_groups: + chan_topologies = mcdata.topologies[group] + chan_permutations = mcdata.permutations[group] + chan_diagram_indices = mcdata.diagram_indices[group] + topo_count = len(chan_topologies) + + channel_index_first = channel_index + symfact_index_first = len(symfact) + for i, diag in enumerate(chan_diagram_indices): + if self.unmerged_meta is None: + amp2_remaps[0][diag] = channel_index + else: + for amp2_remap, d in zip(amp2_remaps, diag): + if d != -1: + amp2_remap[d] = channel_index + if i == 0: + symfact.extend([None] * topo_count) + else: + symfact.extend(range(symfact_index_first, symfact_index_first + topo_count)) + channel_index += 1 + + topologies.append(chan_topologies) + permutations.append(chan_permutations) + channel_indices.append(list(range(channel_index_first, channel_index))) + channel_weight_indices.append([ + [ + symfact_index_first + topo_index + i * topo_count + for i in range(len(chan_permutations)) + ] + for topo_index in range(topo_count) + ]) + diagram_indices.append(chan_diagram_indices) + if mcdata.diagram_color_indices: + diagram_color_indices.append(mcdata.diagram_color_indices[group]) + if mcdata.diagram_propagator_pdgs: + diagram_propagator_pdgs.append( + mcdata.diagram_propagator_pdgs[group] + ) + active_flavors.append(mcdata.active_flavors[group]) + qcd_s_channel_count.append(mcdata.qcd_s_channel_count[group]) + + return MultiChannelData( + amp2_remaps, + symfact, + topologies, + permutations, + channel_indices, + channel_weight_indices, + diagram_indices, + diagram_color_indices, + diagram_propagator_pdgs, + active_flavors, + qcd_s_channel_count, + ) + + def set_madnis_auto_settings(self, rsd: float): + madnis_args = self.process.run_card["madnis"] + n_out = len(self.meta["outgoing"]) + n_events = self.process.run_card["generation"]["events"] + is_gridpack = self.process.run_card["gridpack"]["save_gridpack"] + train_batches = madnis_args["train_batches"] + hidden_dim = min(max(int((7 * rsd) / 32) * 32 + 64, 64), 256) + flow_layers = 3 if rsd < 32 else 4 + lr = min(max((10000 - train_batches) / 8000 * 7e-4 + 3e-4, 3e-4), 1e-3) + if n_out <= 2: + enable = False + elif n_out == 3: + enable = rsd > 10. or n_events > 1000000 or is_gridpack + else: + enable = True + fixed_cwnet_fraction = max(0.33, 1.0 - 10000. / train_batches) + batch_size_per_channel = min(max(int((7 * rsd) / 32) * 32 + 64, 128), 512) + + self.madnis_settings = { + "enable": enable, + "flow_layers": flow_layers, + "flow_hidden_dim": hidden_dim, + "discrete_hidden_dim": hidden_dim, + "cwnet_hidden_dim": hidden_dim, + "train_batches": train_batches, + "lr": lr, + "fixed_cwnet_fraction": fixed_cwnet_fraction, + "batch_size_per_channel": batch_size_per_channel, + } + for key, value in self.madnis_settings.items(): + run_card_value = madnis_args[key] + if run_card_value != "auto": + self.madnis_settings[key] = run_card_value + + def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace: + madnis_args = self.process.run_card["madnis"] + # Shared across all networks below: initialize_globals() derives an + # independent, non-colliding stream per tensor from this one base seed. + seed = self.process.run_seed + channels = [] + for channel_id, channel in enumerate(phasespace.channels): + prefix = f"subproc{self.subproc_id}.channel{channel_id}" + cond_dim = 0 + + flow_dim = channel.phasespace_mapping.random_dim() + flow = ms.Flow( + input_dim=flow_dim, + condition_dim=cond_dim, + prefix=prefix, + bin_count=madnis_args["flow_spline_bins"], + subnet_hidden_dim=self.madnis_settings["flow_hidden_dim"], + subnet_layers=self.madnis_settings["flow_layers"], + subnet_activation=self.activation(madnis_args["flow_activation"]), + invert_spline=madnis_args["flow_invert_spline"], + ) + if channel.adaptive_mapping is None: + flow.initialize_globals(self.process.contexts[0], seed) + else: + flow.initialize_from_vegas( + self.process.contexts[0], channel.adaptive_mapping.grid_name(), + seed + ) + cond_dim += flow_dim + + # discrete_sym runs after the adaptive map, so it can condition on its latent. + discrete_sym = channel.discrete_sym + if discrete_sym is not None: + perm_count = channel.phasespace_mapping.channel_count() + discrete_sym = ms.DiscreteFlow( + option_counts=[perm_count], + prefix=f"{prefix}.discrete_flow_sym", + dims_with_prior=[], + condition_dim=cond_dim, + subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], + subnet_layers=madnis_args["discrete_layers"], + subnet_activation=self.activation(madnis_args["discrete_activation"]), + ) + discrete_sym.initialize_globals(self.process.contexts[0], seed) + cond_dim += perm_count + + discrete_flavor = channel.discrete_flavor + if discrete_flavor is not None: + discrete_flavor = ms.DiscreteFlow( + option_counts=[len(self.meta["flavors"])], + prefix=f"{prefix}.discrete_flow_flavor", + dims_with_prior=[0], + condition_dim=cond_dim, + subnet_hidden_dim=self.madnis_settings["discrete_hidden_dim"], + subnet_layers=madnis_args["discrete_layers"], + subnet_activation=self.activation(madnis_args["discrete_activation"]), + ) + discrete_flavor.initialize_globals(self.process.contexts[0], seed) + + channels.append(Channel( + phasespace_mapping = channel.phasespace_mapping, + adaptive_mapping = flow, + discrete_sym = discrete_sym, + discrete_flavor = discrete_flavor, + channel_weight_indices = channel.channel_weight_indices, + name = channel.name, + active_flavors = channel.active_flavors, + )) + + return PhaseSpace( + mode="both", + channels=channels, + first_chan_weight_remap=phasespace.first_chan_weight_remap, + first_remapped_chan_count=phasespace.first_remapped_chan_count, + second_chan_weight_remap=phasespace.second_chan_weight_remap, + second_remapped_chan_count=phasespace.second_remapped_chan_count, + symfact=phasespace.symfact, + cwnet=self.build_cwnet(len(phasespace.symfact)), + prop_chan_weights=phasespace.prop_chan_weights, + subchan_weights=phasespace.subchan_weights, + ) + + def build_vegas(self, mapping: ms.PhaseSpaceMapping, prefix: str) -> ms.VegasMapping: + if not self.process.run_card["vegas"]["enable"]: + return None + + vegas = ms.VegasMapping( + mapping.random_dim(), + self.process.run_card["vegas"]["bins"], + prefix, + ) + for context in self.process.contexts: + vegas.initialize_globals(context) + return vegas + + def build_discrete( + self, permutation_count: int, flavor_count: int, prefix: str + ) -> tuple[ms.DiscreteSampler | None, ms.DiscreteSampler | None]: + is_adaptive = self.process.run_card["phasespace"]["adaptive_symmetry_sampling"] + if is_adaptive and permutation_count > 1: + discrete_sym = ms.DiscreteSampler( + [permutation_count], f"{prefix}.discrete_sym" + ) + for context in self.process.contexts: + discrete_sym.initialize_globals(context) + else: + discrete_sym = None + + if flavor_count > 1: + # the flavor probabilities are conditioned on the PDF prior; a + # leptonic process (or a decay) has none, and the Integrand then + # passes no condition: declaring one anyway fails to build + # ("keys and values must have the same size"), e.g. z > q q~ with + # both u- and d-type rows + discrete_flavor = ms.DiscreteSampler( + [flavor_count], f"{prefix}.discrete_flavor", + [] if self.process.leptonic else [0] + ) + for context in self.process.contexts: + discrete_flavor.initialize_globals(context) + else: + discrete_flavor = None + + return discrete_sym, discrete_flavor + + def build_cwnet(self, channel_count: int) -> ms.ChannelWeightNetwork: + #if channel_count == 1: + # return None + madnis_args = self.process.run_card["madnis"] + cwnet = ms.ChannelWeightNetwork( + channel_count=channel_count, + particle_count=self.particle_count, + hidden_dim=self.madnis_settings["cwnet_hidden_dim"], + layers=madnis_args["cwnet_layers"], + activation=self.activation(madnis_args["cwnet_activation"]), + prefix=f"subproc{self.subproc_id}.cwnet", + ) + cwnet.initialize_globals( + self.process.contexts[0], self.process.run_seed + ) + return cwnet + + def t_channel_mode(self, name: str) -> ms.PhaseSpaceMapping.TChannelMode: + modes = { + "propagator": ms.PhaseSpaceMapping.propagator, + "rambo": ms.PhaseSpaceMapping.rambo, + "chili": ms.PhaseSpaceMapping.chili, + } + if name in modes: + return modes[name] + else: + raise ValueError(f"Invalid t-channel mode '{name}'") + + def activation(self, name: str) -> ms.MLP.Activation: + activations = { + "relu": ms.MLP.relu, + "leaky_relu": ms.MLP.leaky_relu, + "elu": ms.MLP.elu, + "gelu": ms.MLP.gelu, + "sigmoid": ms.MLP.sigmoid, + "softplus": ms.MLP.softplus, + } + if name in activations: + return activations[name] + else: + raise ValueError(f"Invalid activation function '{name}'") + + def build_integrands( + self, + phasespace: PhaseSpace, + madnis_training: bool = False, + drop_cuts_and_rescale: bool = False + ) -> list[ms.Integrand]: + flavors = [] + flavor_remap = [] + flavor_factors = [] + flavor_mirror = [] + flavor_diff_xs_indices = [] + flavor_subproc_indices = [] + flavor_per_subproc_remap = [] + + for flav in self.meta["flavors"]: + if self.unmerged_meta is not None: + diff_xs_index = flav["subprocess"] + subproc_index = self.meta["subprocesses"][diff_xs_index] + ps_flavor = flav["flavor"] + flavor_diff_xs_indices.append(diff_xs_index) + flavor_subproc_indices.append(subproc_index) + flavor_per_subproc_remap.append(ps_flavor) + flav = self.unmerged_meta[subproc_index]["flavors"][ps_flavor] + flavors.append(flav["options"][0]) + flavor_remap.append(flav["index"]) + flavor_factors.append(len(flav["options"])) + flavor_mirror.append(flav["mirror"]) + + cross_sections = [] + self.step_matrix_elements = [] + for me_index, matrix_element in enumerate(self.matrix_elements): + if self.surrogates: + # the surrogate in the integrand, the full matrix element in + # the second pass, with the random numbers passed in + mat = ms.MatrixElement( + self.surrogates[me_index].index(), + self.particle_count, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in], + [ms.MatrixElement.matrix_element_out], + 1, + False, + self.process.me_frame, + self.process.incoming_count, + ) + self.step_matrix_elements.append(ms.MatrixElement( + matrix_element, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in, ms.MatrixElement.random_color_in, + ms.MatrixElement.random_helicity_in, + ms.MatrixElement.random_diagram_in], + [ms.MatrixElement.matrix_element_out, + ms.MatrixElement.color_index_out, + ms.MatrixElement.helicity_index_out, + ms.MatrixElement.diagram_index_out], + False, + self.process.me_frame, + self.process.incoming_count, + )) + elif matrix_element: + mat = ms.MatrixElement( + matrix_element, + ms.Integrand.matrix_element_inputs, + ms.Integrand.matrix_element_outputs, + True, + self.process.me_frame, + self.process.incoming_count, + ) + else: + #TODO: not working in merged mode + mat = ms.MatrixElement( + 0xBADCAFE, + self.particle_count, + ms.Integrand.matrix_element_inputs, + ms.Integrand.matrix_element_outputs, + self.meta["diagram_count"], + True, + self.process.me_frame, + self.process.incoming_count, + ) + pdf_grid = None if self.process.leptonic else self.process.pdf_grid + pdf_arg = None if self.process.leptonic else ms.CachedPdf() + cross_sections.append( + ms.DifferentialCrossSection( + matrix_element=mat, + # For a decay this is the decaying particle's mass, and the + # flux becomes 1/(2M): the result is a partial width in + # GeV, not a cross section in pb. + cm_energy=self.process.e_cm, + running_coupling=None, + energy_scale=ms.CachedScale(), + pid_options=[], + pdf1=pdf_arg, + pdf2=pdf_arg, + input_momentum_fraction=True, + decay=self.process.is_decay, + return_matrix_element=bool(self.surrogates), + ) + ) + partial_weights = self.process.systematics_enabled() + madnis_args = self.process.run_card["madnis"] + custom_bias = self.process.custom_function("bias", self.particle_count, flavors) + integrands = [] + for channel in phasespace.channels: + integrands.append(ms.Integrand( + channel.phasespace_mapping, + cross_sections, + channel.adaptive_mapping, + channel.discrete_sym, + channel.discrete_flavor, + flavors, + pdf_grid, + self.process.running_coupling, + self.scale, + phasespace.prop_chan_weights, + phasespace.subchan_weights, + phasespace.cwnet, + phasespace.first_chan_weight_remap, + phasespace.first_remapped_chan_count, + phasespace.second_chan_weight_remap, + phasespace.second_remapped_chan_count, + madnis_training, + drop_cuts_and_rescale, + partial_weights, + channel.channel_weight_indices, + channel.active_flavors, + flavor_remap, + flavor_factors, + flavor_mirror, + flavor_diff_xs_indices, + flavor_subproc_indices, + flavor_per_subproc_remap, + madnis_args["compressed_channel_weight_count"], + custom_bias, + )) + #print(integrands[1].function()) + #for i in integrands: print(i.function()) + return integrands diff --git a/tests/unit_tests/interface/test_mg7_launch.py b/tests/unit_tests/interface/test_mg7_launch.py index 8f7fd39788..12f851e099 100644 --- a/tests/unit_tests/interface/test_mg7_launch.py +++ b/tests/unit_tests/interface/test_mg7_launch.py @@ -846,14 +846,15 @@ def test_setup_logging_defers_to_an_existing_handler(self): lg = logging.getLogger('madgraph') handler = logging.NullHandler() lg.addHandler(handler) - saved = self.launch._TOOL_LOGGING_READY - self.launch._TOOL_LOGGING_READY = False + from madgraph.iolibs.template_files.mg7 import postprocess + saved = postprocess._TOOL_LOGGING_READY + postprocess._TOOL_LOGGING_READY = False try: before = list(lg.handlers) self.launch._setup_logging() self.assertEqual(list(lg.handlers), before) finally: - self.launch._TOOL_LOGGING_READY = saved + postprocess._TOOL_LOGGING_READY = saved lg.removeHandler(handler) From 636deadf4d4a5dcf6c0bd433cbe1f610d9171641 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:08:46 +0200 Subject: [PATCH 05/22] mg7: weighting steps of each subprocess The subprocess returns its weighting steps, so the weighting groups do not depend on how a subprocess builds them. A [surrogate] subprocess returns its second pass with the full matrix element, as before. --- .../template_files/mg7/madgraph_process.py | 27 ++++++------------- .../mg7/subprocess_madmatrix.py | 13 +++++++++ 2 files changed, 21 insertions(+), 19 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/madgraph_process.py b/madgraph/iolibs/template_files/mg7/madgraph_process.py index 7321b86638..f1f70ae409 100644 --- a/madgraph/iolibs/template_files/mg7/madgraph_process.py +++ b/madgraph/iolibs/template_files/mg7/madgraph_process.py @@ -330,25 +330,14 @@ def init_surrogate(self) -> None: ) def weighting_groups(self, subprocess_channels: list[list[int]]) -> list: - """One ms.WeightingGroup per subprocess with a [surrogate]: the second - pass with the full matrix element over its channels, given by their - indices in the event generator.""" - if self.surrogate is None: - return [] - gen_args = self.run_card["generation"] - return [ - ms.WeightingGroup( - channel_indices, - [ms.WeightingStep( - subproc.step_matrix_elements, - mode=self.surrogate["mode"], - max_overweight_truncation=gen_args["max_overweight_truncation"], - freeze_max_weight_after=gen_args["freeze_max_weight_after"], - )], - ) - for subproc, channel_indices in zip(self.subprocesses, subprocess_channels) - if channel_indices - ] + """One ms.WeightingGroup per subprocess with weighting steps, over its + channels, given by their indices in the event generator.""" + groups = [] + for subproc, channel_indices in zip(self.subprocesses, subprocess_channels): + steps = subproc.weighting_steps() + if channel_indices and steps: + groups.append(ms.WeightingGroup(channel_indices, steps)) + return groups @staticmethod def build_custom_function(build_dir: str) -> None: diff --git a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py index 8352894842..98eb957a93 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py +++ b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py @@ -377,6 +377,19 @@ def load_surrogate(self, me_path_format: str, member_meta: dict): raise RuntimeError("surrogate loaded at different indices") return apis[0] + def weighting_steps(self) -> list: + """The steps after the first unweighting: with a [surrogate], the + second pass with the full matrix element.""" + if not self.surrogates: + return [] + gen_args = self.process.run_card["generation"] + return [ms.WeightingStep( + self.step_matrix_elements, + mode=self.process.surrogate["mode"], + max_overweight_truncation=gen_args["max_overweight_truncation"], + freeze_max_weight_after=gen_args["freeze_max_weight_after"], + )] + def flavor_pids(self) -> list[list[int]]: """PDG ids of the external particles for every flavor of the integrand, in the order of its flavor index.""" From 68714e6f104596723a052f723bceff14aa95a4dc Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 22:52:43 +0200 Subject: [PATCH 06/22] interface: '[amplicol]' processes 'generate p p > j j j [amplicol]' keeps the process definition for the mg7 output, which computes its matrix elements with AmpliCol, instead of generating MG diagrams. Such processes can be mixed with ordinary ones, and are refused by the other outputs and for decay chains, coupling orders, s-channel or particle restrictions and polarizations. --- madgraph/interface/madgraph_interface.py | 62 +++++++++++++- madgraph/interface/master_interface.py | 4 + .../interface/test_amplicol_syntax.py | 85 +++++++++++++++++++ 3 files changed, 150 insertions(+), 1 deletion(-) create mode 100644 tests/unit_tests/interface/test_amplicol_syntax.py diff --git a/madgraph/interface/madgraph_interface.py b/madgraph/interface/madgraph_interface.py index 836112fc28..32f008229c 100755 --- a/madgraph/interface/madgraph_interface.py +++ b/madgraph/interface/madgraph_interface.py @@ -178,6 +178,26 @@ def interference_mode_of(option, orders): return '' +# '[amplicol]' takes the matrix elements of a process from AmpliCol, a +# colour-ordered generator, instead of MG diagrams (mg7 output only) +ME_BACKENDS = ('amplicol',) +_BACKEND_BRACKET = re.compile(r"\[\s*(?P\w+)\s*\]") + + +def split_backend_line(line): + """(line without the bracket, backend) for a process with a matrix-element + backend bracket like '[amplicol]', (line, '') for any other line.""" + + match = _BACKEND_BRACKET.search(line) + if not match or match.group('backend').lower() not in ME_BACKENDS: + return line, '' + rest = line[:match.start()] + ' ' + line[match.end():] + if '[' in rest or ']' in rest: + raise madgraph.InvalidCmd( + "[%s] cannot be combined with another '[...]'." % match.group('backend')) + return ' '.join(rest.split()), match.group('backend').lower() + + def split_interference_line(line): """Split a process definition of the form @@ -2345,7 +2365,11 @@ def check_output(self, args, default='mg7'): return - if not self._curr_amps: + if self._curr_amplicol_defs and self._export_format != 'mg7': + raise self.InvalidCmd( + 'The [amplicol] processes are only supported by the mg7 output.') + + if not self._curr_amps and not self._curr_amplicol_defs: text = 'No processes generated. Please generate a process first.' raise self.InvalidCmd(text) @@ -3866,6 +3890,8 @@ def _tutorial_opts(self): _curr_model = None #base_objects.Model() _curr_amps = diagram_generation.AmplitudeList() _curr_proc_defs = base_objects.ProcessDefinitionList() + # processes whose matrix elements come from AmpliCol, see split_backend_line + _curr_amplicol_defs = base_objects.ProcessDefinitionList() _curr_matrix_elements = helas_objects.HelasMultiProcess() _curr_helas_model = None _curr_exporter = None @@ -3954,6 +3980,7 @@ def setup(self): self._done_export=False self._curr_amps = diagram_generation.AmplitudeList() self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() self._curr_matrix_elements = helas_objects.HelasMultiProcess() self._v4_export_formats = ['madevent', 'standalone_fortran','standalone_msP','standalone_msF', @@ -3978,6 +4005,30 @@ def do_quit(self, line): # Add a process to the existing multiprocess definition # Generate a new amplitude + def add_backend_process(self, procdef, backend): + """Keep a process whose matrix elements come from another backend than + MG diagrams. It is exported by the mg7 output, without diagrams.""" + + if procdef.get('decay_chains') or procdef.get('required_s_channels') or \ + procdef.get('forbidden_onsh_s_channels') or \ + procdef.get('forbidden_s_channels') or \ + procdef.get('forbidden_particles') or procdef.get('orders') or \ + procdef.get('squared_orders') or procdef.get('is_decay_chain') or \ + any(leg.get('polarization') for leg in procdef.get('legs')) or \ + procdef.get_ninitial() != 2: + raise self.InvalidCmd( + "[%s] supports 2 -> n processes without decay chains, coupling " + "orders, s-channel or particle restrictions and polarizations." + % backend) + initial = [amp.get_ninitial() for amp in self._curr_amps] + if initial and initial[0] != procdef.get_ninitial(): + raise self.InvalidCmd( + "Can not mix processes with different number of initial states.") + self._curr_amplicol_defs.append(procdef) + self._done_export = False + logger.info("Process %s is computed by %s at output time" % + (procdef.nice_string(prefix=False), backend)) + def do_add(self, line, counter=0): """Generate an amplitude for a given process and add to existing amplitudes @@ -4079,7 +4130,10 @@ def do_add(self, line, counter=0): else: nb_proc = len([l for l in self.history if l.startswith(('generate','add process'))]) nb_proc += counter + line, backend = split_backend_line(line) myprocdef = self.extract_process(line, proc_number=nb_proc) + if backend: + return self.add_backend_process(myprocdef, backend) # 'left [treextree] right': hidden order selecting the # interference, minimal orders of the right-hand process diagram_generation.prepare_interference_process(myprocdef) @@ -6045,6 +6099,7 @@ def clean_process(self): self._curr_amps = diagram_generation.AmplitudeList() # Reset Process definition self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() # Reset Helas matrix elements self._curr_matrix_elements = helas_objects.HelasMultiProcess() self._generate_info = "" @@ -7555,6 +7610,7 @@ def do_import(self, line, force=False, options={}): self._curr_amps = diagram_generation.AmplitudeList() # Reset proc defs self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() self._curr_matrix_elements = helas_objects.HelasMultiProcess() process_checks.store_aloha = [] self.advise_neglected_masses() @@ -10938,6 +10994,7 @@ def set2_group_subprocesses(self, args, log=True): logger.info('Note that you need to regenerate all processes') self._curr_amps = diagram_generation.AmplitudeList() self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() self._curr_matrix_elements = helas_objects.HelasMultiProcess() if self.options[args[0]] == 'gpu': export_v4.ProcessExporterFortranMEGroup.grouped_mode = 'gpu' @@ -11113,6 +11170,7 @@ def set2_gauge(self, args, log=True): self._curr_model = None self._curr_amps = diagram_generation.AmplitudeList() self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() self._curr_matrix_elements = helas_objects.HelasMultiProcess() self._curr_helas_model = None self._curr_exporter = None @@ -11592,6 +11650,7 @@ def set2_OLP(self, args, log=True): # depend on this option self._curr_amps = diagram_generation.AmplitudeList() self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() self._curr_matrix_elements = helas_objects.HelasMultiProcess() self._curr_exporter = None self._second_exporter = None @@ -13544,6 +13603,7 @@ def do_decay_diagram(self, line, skip_2body=False, model=None): # Reset amplitudes self._curr_amps = diagram_generation.AmplitudeList() self._curr_proc_defs = base_objects.ProcessDefinitionList() + self._curr_amplicol_defs = base_objects.ProcessDefinitionList() # Reset Helas matrix elements self._curr_matrix_elements = helas_objects.HelasMultiProcess() # Reset _done_export, since we have new process diff --git a/madgraph/interface/master_interface.py b/madgraph/interface/master_interface.py index 7db93f4e83..dd440419a9 100755 --- a/madgraph/interface/master_interface.py +++ b/madgraph/interface/master_interface.py @@ -173,6 +173,10 @@ def extract_process_type(line): space_before = re.compile(r"(?P\S)(?P[\\[\\]/\,\\$\\>|])(?P\S)") line2 = space_before.sub(r'\g \g \g', line) + # a matrix-element backend bracket like [amplicol] is a tree process + if MGcmd.split_backend_line(re.split(r'\s\-\-', line, 1)[0])[1]: + return ('tree', None, []) + # Use regular expressions to extract the loop mode (if present) and its # option, specified in the line with format [ option = loop_orders ] or # [ loop_orders ] which implicitly select the 'all' option. diff --git a/tests/unit_tests/interface/test_amplicol_syntax.py b/tests/unit_tests/interface/test_amplicol_syntax.py new file mode 100644 index 0000000000..1e20f0d35e --- /dev/null +++ b/tests/unit_tests/interface/test_amplicol_syntax.py @@ -0,0 +1,85 @@ +################################################################################ +# +# Copyright (c) 2026 The MadGraph5_aMC@NLO Development team and Contributors +# +# This file is a part of the MadGraph5_aMC@NLO project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph5_aMC@NLO license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""Syntax of the processes computed by AmpliCol + + generate p p > j j j [amplicol] + +The process is kept as a definition for the mg7 output, which generates its +matrix elements with pyamplicol; no MG diagrams are generated for it. +""" + +import os +import sys + +root_path = os.path.split(os.path.dirname(os.path.realpath(__file__)))[0] +sys.path.append(os.path.join(root_path, os.path.pardir, os.path.pardir)) + +import tests.unit_tests as unittest + +import madgraph.interface.madgraph_interface as mg_interface +import madgraph.interface.master_interface as MGCmd + +from madgraph import InvalidCmd + +split = mg_interface.split_backend_line + + +class BackendLineSplitTest(unittest.TestCase): + """split_backend_line, without any model.""" + + def test_split(self): + self.assertEqual(split('p p > j j j [amplicol]'), ('p p > j j j', 'amplicol')) + self.assertEqual(split('p p > j j j [ AmpliCol ] @2'), + ('p p > j j j @2', 'amplicol')) + self.assertEqual(split('p p > j j j'), ('p p > j j j', '')) + self.assertEqual(split('p p > t t~ [QCD]'), ('p p > t t~ [QCD]', '')) + + def test_other_brackets_are_refused(self): + self.assertRaises(InvalidCmd, split, 'p p > j j [amplicol] [QCD]') + + def test_process_type(self): + process_type = MGCmd.MasterCmd.extract_process_type + self.assertEqual(process_type('p p > j j j [amplicol]'), ('tree', None, [])) + self.assertEqual(process_type('p p > t t~ [QCD]'), ('NLO', 'all', ['QCD'])) + + +class AmplicolProcessTest(unittest.TestCase): + """generate/add process with [amplicol], with the sm model.""" + + def setUp(self): + self.cmd = MGCmd.MasterCmd() + self.cmd.exec_cmd('import model sm', printcmd=False, precmd=True) + + def test_no_diagrams_are_generated(self): + self.cmd.exec_cmd('generate g g > g g g [amplicol]', printcmd=False, + precmd=True) + self.cmd.exec_cmd('add process p p > e+ e-', printcmd=False, precmd=True) + self.assertEqual(len(self.cmd._curr_amplicol_defs), 1) + self.assertEqual( + self.cmd._curr_amplicol_defs[0].nice_string(prefix=False).split('@')[0].strip(), + 'g g > g g g') + # the madmatrix process keeps its own process number + self.assertTrue(self.cmd._curr_amps) + self.assertEqual( + {amp.get('process').get('id') for amp in self.cmd._curr_amps}, {2}) + self.cmd.exec_cmd('generate p p > e+ e-', printcmd=False, precmd=True) + self.assertEqual(len(self.cmd._curr_amplicol_defs), 0) + + def test_unsupported_options_are_refused(self): + for line in ['p p > j j QED=0 [amplicol]', + 'p p > z > e+ e- [amplicol]', + 'p p > j j / b [amplicol]', + 'z > e+ e- [amplicol]']: + self.assertRaises(InvalidCmd, self.cmd.do_generate, line) From 0b2f9c3556ac8dba1b4eab1635fe6d47dd669a5a Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 22:52:43 +0200 Subject: [PATCH 07/22] mg7: output of the [amplicol] processes Each [amplicol] process gets a SubProcesses/A_ directory with the leading- and full-colour pyamplicol artifacts, the sources of their ratio library and a makefile building both libraries into lib/. Every colour ordering becomes its own entry of subprocesses.json, with backend 'amplicol', and is never merged with other subprocesses. --- madgraph/iolibs/export_amplicol.py | 247 +++++++++++++++++++++++++++++ madgraph/iolibs/export_cpp.py | 20 ++- 2 files changed, 266 insertions(+), 1 deletion(-) create mode 100644 madgraph/iolibs/export_amplicol.py diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py new file mode 100644 index 0000000000..585d4c9b7a --- /dev/null +++ b/madgraph/iolibs/export_amplicol.py @@ -0,0 +1,247 @@ +################################################################################ +# +# Copyright (c) 2026 The MadGraph5_aMC@NLO Development team and Contributors +# +# This file is a part of the MadGraph5_aMC@NLO project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph5_aMC@NLO license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""mg7 output of the processes computed by AmpliCol ('[amplicol]'). + +pyamplicol generates the colour-ordered matrix elements of a process definition +in SubProcesses/A_: a leading-colour artifact (lc), a full-colour one +(full), and the sources of the UMAMI library returning their ratio (reweight). +Its makefile builds the libraries into lib/. Every colour ordering (UMAMI +channel) of every provider becomes one entry of subprocesses.json, integrated +on its own with a colour-ordered phase-space mapping. +""" + +import json +import logging +import os +import re +import shutil + +import madgraph + +logger = logging.getLogger('madgraph.export_amplicol') +pjoin = os.path.join + +MODEL_FILE = 'amplicol_model.pyamplicol-model' + +MAKEFILE = """\ +# Builds the AmpliCol matrix elements of this process into ../../lib. +RUSTICOL_CONFIG ?= %(rusticol_config)s +BUILD := $(CURDIR)/build + +all: +\t$(MAKE) -C lc/API/umami shared BUILD_DIR=$(BUILD)/lc RUSTICOL_CONFIG=$(RUSTICOL_CONFIG) +\t$(MAKE) -C reweight shared BUILD_DIR=$(BUILD)/reweight RUSTICOL_CONFIG=$(RUSTICOL_CONFIG) +%(copies)s +clean: +\trm -rf $(BUILD) +""" + + +def import_pyamplicol(): + try: + import pyamplicol + from pyamplicol.artifacts import umami_reweight + except ImportError as error: + raise madgraph.InvalidCmd( + 'The [amplicol] processes need pyamplicol (%s). Install it in the ' + 'Python environment of MadGraph.' % error) + return pyamplicol, umami_reweight + + +def process_name(procdef): + """A directory name for a process definition, e.g. 'pp_jjj'.""" + string = procdef.nice_string(prefix=False).split('@')[0] + initial, final = [side.split() for side in string.split('>')[:2]] + name = ''.join(initial) + '_' + ''.join(final) + return re.sub(r'[^A-Za-z0-9_]', '', + name.replace('~', 'x').replace('+', 'p').replace('-', 'm')) + + +def process_expression(procdef, model): + """The process in pyamplicol syntax, and its multiparticle definitions.""" + particles = model.get('particle_dict') + multiparticles = {} + labels = [] + for leg in procdef.get('legs'): + names = [particles[pid].get_name() for pid in leg.get('ids')] + if len(names) == 1: + labels.append(names[0]) + continue + label = next((key for key, value in multiparticles.items() + if value == names), None) + if label is None: + label = 'mp%d' % len(multiparticles) + multiparticles[label] = names + labels.append(label) + initial = [label for leg, label in zip(procdef.get('legs'), labels) + if not leg.get('state')] + final = [label for leg, label in zip(procdef.get('legs'), labels) + if leg.get('state')] + return '%s > %s' % (' '.join(initial), ' '.join(final)), multiparticles + + +def model_source(pyamplicol, model, cache_path): + """The UFO model of the output with its restriction, compiled for + pyamplicol once per output.""" + if os.path.exists(cache_path): + return pyamplicol.ModelSource.from_path(cache_path) + path = model.get('modelpath') + reference = os.path.basename(model.get('modelpath+restriction')) + restriction = 'default' + if '-' in reference: + suffix = reference.split('-', 1)[1] + if suffix == 'full': + restriction = 'none' + else: + restriction = pjoin(path, 'restrict_%s.dat' % suffix) + from pyamplicol.config import EvaluatorConfig, JITConfig + source = pyamplicol.ModelSource.from_path(path, restriction=restriction) + logger.info('Compiling the model %s for AmpliCol' % reference) + source.compile(prepared_output=cache_path, + evaluator=EvaluatorConfig(jit=JITConfig(optimization_level=2))) + return pyamplicol.ModelSource.from_path(cache_path) + + +def generate_artifact(pyamplicol, artifact, accuracy, expression, multiparticles, + model, workers): + from pyamplicol.config import resolve_config + config = resolve_config({ + 'schema_version': 1, + 'action': 'generate', + 'color': {'accuracy': accuracy}, + 'process': {'multiparticles': multiparticles}, + 'generation': {'workers': workers, 'emit_api_bundle': True, + 'relation_discovery': {'mode': 'off'}}, + 'evaluator': {'backend': 'jit', 'execution_mode': 'recurrence', + 'optimization': {'cores': workers}, + 'jit': {'optimization_level': 2}}, + }) + pyamplicol.Generator(config).generate(expression, artifact, model=model, + mode='replace') + + +def color_types(model, pdgs): + types = {} + for pdg in pdgs: + particle = model.get_particle(pdg) + types[pdg] = particle.get_color() + types[-pdg] = model.get_particle(particle.get_anti_pdg_code()).get_color() + return types + + +def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path): + """One subprocesses.json entry per channel of a UMAMI provider.""" + pdg_ids = metadata['pdg_ids'] + provider = metadata['provider'] + incoming = provider['incoming_count'] + entries = [] + for channel_index, channel in enumerate(metadata['channels']): + flavors = [] + for flavor_index, process in enumerate(channel['processes']): + if process['multichannels'] != [channel_index]: + raise madgraph.MadGraph5Error( + 'AmpliCol contributions shared by several channels are not ' + 'supported yet') + for term in process['matrix_elements']: + flavors.append({ + 'index': flavor_index, + 'options': [pdg_ids[term['pdg_ids']]], + 'factor': term['factor'], + 'mirror': False, + }) + first = flavors[0]['options'][0] + colored = sum(1 for pdg in first + if abs(model.get_particle(pdg).get_color()) > 1) + entries.append({ + 'backend': 'amplicol', + 'incoming': first[:incoming], + 'outgoing': first[incoming:], + 'channels': [], + 'path': rel_dir, + 'me_path': me_path, + 'ratio_path': ratio_path, + 'artifact': pjoin(rel_dir, 'lc'), + 'ratio_artifact': pjoin(rel_dir, 'reweight'), + 'umami_channel': channel_index, + 'phasespace_order': channel['phasespace_order'], + 'masses': provider['masses'], + 'flavors': flavors, + # tree level with the minimal coupling orders: every coloured + # parton beyond the first two comes with one power of alpha_s + 'qcd_power': max(colored - 2, 0), + 'color_flows': metadata['color_flows'], + 'pdg_color_types': color_types( + model, {abs(pdg) for pdgs in pdg_ids for pdg in pdgs}), + 'helicities': metadata['helicities'], + 'diagram_count': 0, + 'interference': False, + 'parameters': parameter_references(metadata), + }) + return entries + + +def parameter_references(metadata): + """The parameter card entry of every mutable runtime parameter.""" + references = {} + for process in metadata['runtime_processes']: + for parameter in process['parameters']: + if parameter.get('mutable') and parameter.get('lhablock'): + references[parameter['name']] = [parameter['lhablock'], + parameter['lhacode']] + return references + + +def export_amplicol_processes(dir_path, proc_defs, model, nb_core=None): + """Write SubProcesses/A_ for every '[amplicol]' process definition + and return the subprocesses.json entries.""" + pyamplicol, umami_reweight = import_pyamplicol() + workers = nb_core or 1 + subproc_dir = pjoin(dir_path, 'SubProcesses') + source = model_source(pyamplicol, model, pjoin(subproc_dir, MODEL_FILE)) + rusticol_config = shutil.which('rusticol-config') or 'rusticol-config' + entries = [] + for number, procdef in enumerate(proc_defs, start=1): + name = 'A%d_%s' % (number, process_name(procdef)) + path = pjoin(subproc_dir, name) + rel_dir = pjoin('SubProcesses', name) + expression, multiparticles = process_expression(procdef, model) + logger.info('Generating the AmpliCol matrix elements of %s in %s' % + (procdef.nice_string(prefix=False), path)) + os.makedirs(path, exist_ok=True) + for accuracy in ('lc', 'full'): + generate_artifact(pyamplicol, pjoin(path, accuracy), accuracy, + expression, multiparticles, source, workers) + provider_ids = umami_reweight.write_reweight_provider( + pjoin(path, 'lc'), pjoin(path, 'full'), pjoin(path, 'reweight')) + documents = umami_reweight.load_umami_metadata(pjoin(path, 'lc')) + copies = [] + for provider_id, metadata in zip(provider_ids, documents): + library = metadata['provider']['library'] + me_path = pjoin('lib', 'libamplicol_%s_%s.so' % (name, provider_id)) + ratio_path = pjoin('lib', 'libamplicol_%s_%s_ratio.so' % + (name, provider_id)) + copies.append('\tcp $(BUILD)/lc/%s ../../%s' % (library, me_path)) + copies.append('\tcp $(BUILD)/reweight/libreweight_%s.so ../../%s' % + (provider_id, ratio_path)) + entries.extend(subprocess_entries(metadata, model, rel_dir, me_path, + ratio_path)) + with open(pjoin(path, 'makefile'), 'w') as stream: + stream.write(MAKEFILE % {'rusticol_config': rusticol_config, + 'copies': '\n'.join(copies)}) + with open(pjoin(path, 'process.json'), 'w') as stream: + json.dump({'process': procdef.nice_string(prefix=False), + 'expression': expression, + 'multiparticles': multiparticles}, stream, indent=2) + return entries diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index e70cbfe98c..e791c79fa4 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3458,6 +3458,7 @@ def get_merged_info(self): return merged_subproc_info def finalize(self, matrix_elements=None, history='', *args, **kwargs): + self.export_amplicol() file_name = os.path.normpath(os.path.join( self.dir_path, "SubProcesses", "subprocesses.json" )) @@ -3504,8 +3505,25 @@ def finalize(self, matrix_elements=None, history='', *args, **kwargs): def pass_information_from_cmd(self, cmd): """Capture the process definitions from the command interface; needed to - generate the MadAnalysis5 default cards at output time.""" + generate the MadAnalysis5 default cards at output time, and the + [amplicol] processes.""" self.proc_defs = getattr(cmd, '_curr_proc_defs', None) + self.amplicol_defs = getattr(cmd, '_curr_amplicol_defs', None) or [] + self.amplicol_model = cmd._curr_model + self.amplicol_nb_core = cmd.options.get('nb_core') + + def export_amplicol(self): + """The [amplicol] processes, as subprocesses after the madmatrix ones.""" + if not getattr(self, 'amplicol_defs', None): + return + import madgraph.iolibs.export_amplicol as export_amplicol + for entry in export_amplicol.export_amplicol_processes( + self.dir_path, self.amplicol_defs, self.amplicol_model, + self.amplicol_nb_core): + # never merged with another subprocess + self.merged_subprocesses[('amplicol', len(self.process_info))].append( + (len(self.process_info), [])) + self.process_info.append(entry) def create_ma5_default_cards(self, matrix_elements, history): """Call MadAnalysis5 to write process-tailored default analysis cards From 40ce2e932feed303e1ee432774429172c5b9a433 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:01:25 +0200 Subject: [PATCH 08/22] madspace: weighting step matrix elements with a fixed channel A fixed channel is not an external input of a matrix element. The check that the random inputs of a step are passed in now only looks at them. --- madspace/src/phasespace/weighting_step.cpp | 11 +++++++- madspace/tests/test_integrand_me_outputs.py | 12 +++++---- madspace/tests/test_weighting_steps.py | 29 ++++++++++++++++++++- 3 files changed, 45 insertions(+), 7 deletions(-) diff --git a/madspace/src/phasespace/weighting_step.cpp b/madspace/src/phasespace/weighting_step.cpp index c0e75f7075..8982b1bf24 100644 --- a/madspace/src/phasespace/weighting_step.cpp +++ b/madspace/src/phasespace/weighting_step.cpp @@ -107,7 +107,16 @@ WeightingStepFunction::WeightingStepFunction( ); } } - if (matrix_element.external_inputs().size() != matrix_element.inputs().size()) { + bool has_random_inputs = std::any_of( + matrix_element.inputs().begin(), + matrix_element.inputs().end(), + [](auto input) { + return input == MatrixElement::random_color_in || + input == MatrixElement::random_helicity_in || + input == MatrixElement::random_diagram_in; + } + ); + if (has_random_inputs && matrix_element.sample_random_inputs()) { throw std::invalid_argument( "the random inputs of a weighting step matrix element are passed in " "and must not be sampled internally" diff --git a/madspace/tests/test_integrand_me_outputs.py b/madspace/tests/test_integrand_me_outputs.py index 82e400a5af..a5a9fc5c53 100644 --- a/madspace/tests/test_integrand_me_outputs.py +++ b/madspace/tests/test_integrand_me_outputs.py @@ -21,11 +21,11 @@ import madspace as ms -def python_matrix_element(context, function, outputs=("matrix_element",)): +def python_matrix_element(context, function, outputs=("matrix_element",), channel=-1): + """With a channel, the function also gets the fixed channel index.""" + inputs = ["momenta"] if channel < 0 else ["momenta", "channel_index"] api = context.add_matrix_element( - ms.PythonFunction( - function, ["momenta"], list(outputs), particle_count=len(PIDS) - ) + ms.PythonFunction(function, inputs, list(outputs), particle_count=len(PIDS)) ) output_enums = { "matrix_element": ms.MatrixElement.matrix_element_out, @@ -37,8 +37,10 @@ def python_matrix_element(context, function, outputs=("matrix_element",)): return ms.MatrixElement( api.index(), len(PIDS), - [ms.MatrixElement.momenta_in], + [ms.MatrixElement.momenta_in] + + ([] if channel < 0 else [ms.MatrixElement.channel_in]), [output_enums[output] for output in outputs], + channel=channel, ) diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py index fef4cbcf8b..1967cc74f8 100644 --- a/madspace/tests/test_weighting_steps.py +++ b/madspace/tests/test_weighting_steps.py @@ -76,6 +76,7 @@ def generate( step_outputs=("matrix_element",), override_outputs=(), integrand_outputs=("matrix_element",), + step_channel=-1, ): """Survey and generate, returns the status, the combined events and the channels. step_me and step_outputs can be lists, one entry per step.""" @@ -86,7 +87,7 @@ def generate( step_me = [step_me] * len(step_modes) step_outputs = [step_outputs] * len(step_modes) step_mes = [ - python_matrix_element(context, me, outputs) + python_matrix_element(context, me, outputs, step_channel) for me, outputs in zip(step_me, step_outputs) ] config = make_config() @@ -256,6 +257,32 @@ def test_ratio_step(tmp_path, step_modes): assert status_me.mean == status_ratio.mean +def test_ratio_with_fixed_channel(tmp_path): + """A step matrix element evaluated at a fixed channel gets it as an input, + which is not an input of the step.""" + + def ratio(momenta, channel): + return true_over_surrogate(momenta) * (channel == 3) + + results = [ + generate( + str(tmp_path / name), + surrogate_me, + ["reweight"], + step_me=step_me, + step_outputs=("matrix_element_ratio",), + step_channel=channel, + ) + for name, step_me, channel in [ + ("plain", true_over_surrogate, -1), + ("channel", ratio, 3), + ] + ] + (status_plain, events_plain, _), (status_channel, events_channel, _) = results + assert np.array_equal(events_plain, events_channel) + assert status_plain.mean == status_channel.mean + + def test_ratio_then_matrix_element(tmp_path, reference): """After a ratio step, the matrix element it implies is the reference of the next step, so the ratios still multiply to |M_true|^2 / |M_surrogate|^2.""" From 0e137ac8fafca8e9473d84cc63fd759f5a95c351 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:09:12 +0200 Subject: [PATCH 09/22] mg7: outputs with only [amplicol] processes An output without MG diagrams skips their export. The [amplicol] processes stand in for them where the output needs MG processes: the run card defaults, proc_characteristics and the model reference. --- madgraph/interface/madgraph_interface.py | 4 ++-- madgraph/iolibs/export_amplicol.py | 25 ++++++++++++++++++++++++ madgraph/iolibs/export_cpp.py | 18 +++++++++++++++-- 3 files changed, 43 insertions(+), 4 deletions(-) diff --git a/madgraph/interface/madgraph_interface.py b/madgraph/interface/madgraph_interface.py index 32f008229c..ce76124ce1 100755 --- a/madgraph/interface/madgraph_interface.py +++ b/madgraph/interface/madgraph_interface.py @@ -12544,7 +12544,7 @@ def do_output(self, line): # have been defined with multiparticle labels, because then # branching ratios necessitates to keep subprocesses independent. # That applies only if there is more than one subprocess of course. - if self._curr_amps[0].get_ninitial() == 1 and \ + if self._curr_amps and self._curr_amps[0].get_ninitial() == 1 and \ len(self._curr_amps)>1: processes = [amp.get('process') for amp in self._curr_amps if 'process' in list(amp.keys())] @@ -12940,7 +12940,7 @@ def generate_matrix_elements(self, group_processes=True): matrix_elements = self._curr_matrix_elements.get_matrix_elements() - if any(l.get('flavor') is None for l in matrix_elements[0].get('processes')[0].get('legs')): + if matrix_elements and any(l.get('flavor') is None for l in matrix_elements[0].get('processes')[0].get('legs')): [m.get_external_flavors() for m in matrix_elements] #precompute all flavors # Just the matrix.f files if self._export_format == 'matrix': diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index 585d4c9b7a..2f8a35e362 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -132,6 +132,30 @@ def generate_artifact(pyamplicol, artifact, accuracy, expression, multiparticles mode='replace') +def representative_processes(proc_defs, model): + """Processes standing for the [amplicol] definitions where the output needs + MG processes (run card defaults, proc_characteristics): every initial state, + with the first particle of each final-state leg.""" + import itertools + import madgraph.core.base_objects as base_objects + processes = [] + for procdef in proc_defs: + legs = procdef.get('legs') + initial = [leg.get('ids') for leg in legs if not leg.get('state')] + final = [leg.get('ids')[0] for leg in legs if leg.get('state')] + for ids in itertools.product(*initial): + pdgs = list(ids) + final + processes.append(base_objects.Process({ + 'legs': base_objects.LegList([ + base_objects.Leg({'id': pdg, 'number': number + 1, + 'state': number >= len(ids)}) + for number, pdg in enumerate(pdgs)]), + 'model': model, + 'id': procdef.get('id'), + })) + return processes + + def color_types(model, pdgs): types = {} for pdg in pdgs: @@ -177,6 +201,7 @@ def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path): 'umami_channel': channel_index, 'phasespace_order': channel['phasespace_order'], 'masses': provider['masses'], + 'supports_alpha_s': provider['supports_alpha_s'], 'flavors': flavors, # tree level with the minimal coupling orders: every coloured # parton beyond the first two comes with one power of alpha_s diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index e791c79fa4..7e34afcb35 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3512,6 +3512,14 @@ def pass_information_from_cmd(self, cmd): self.amplicol_model = cmd._curr_model self.amplicol_nb_core = cmd.options.get('nb_core') + def amplicol_processes(self): + """MG processes standing for the [amplicol] processes.""" + if not getattr(self, 'amplicol_defs', None): + return [] + import madgraph.iolibs.export_amplicol as export_amplicol + return export_amplicol.representative_processes( + self.amplicol_defs, self.amplicol_model) + def export_amplicol(self): """The [amplicol] processes, as subprocesses after the madmatrix ones.""" if not getattr(self, 'amplicol_defs', None): @@ -3613,6 +3621,8 @@ def create_run_card(self, matrix_elements, history): for me in matrix_elements['matrix_elements']] except (KeyError, TypeError): processes = None + if self.amplicol_processes(): + processes = (processes or []) + [self.amplicol_processes()] if processes: run_card.create_default_for_process(self.proc_characteristic, @@ -3649,8 +3659,12 @@ def create_proc_characteristics(self, matrix_elements): else: me_list = [] - procs = [] - qcd_orders = set() + procs = list(self.amplicol_processes()) + qcd_orders = set(entry['qcd_power'] for entry in self.process_info + if entry.get('backend') == 'amplicol') + if procs: + pc['nexternal'] = max(pc['nexternal'], len(procs[0].get('legs'))) + pc['ninitial'] = 2 for me in me_list: if not me.get('processes'): continue From 62cc98992efb30b2d75dacfb154a1fb4b3c2bb84 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:09:12 +0200 Subject: [PATCH 10/22] mg7: [amplicol] run card section The colour accuracy of the [amplicol] events, lc or full, whether the leading-colour events are re- or unweighted to it, and the settings of the leading-colour unweighting. --- madgraph/iolibs/template_files/mg7/run_card.toml | 9 +++++++++ madgraph/various/banner.py | 11 +++++++++++ 2 files changed, 20 insertions(+) diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index ea618c4862..0837f4258e 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -142,6 +142,15 @@ first_pass_max_overweight_truncation = %(surrogate.first_pass_max_overweight_tru first_pass_freeze_max_weight_after = %(surrogate.first_pass_freeze_max_weight_after)s param_card = %(surrogate.param_card)s # for a .so surrogate, "" uses Cards/param_card.dat +[amplicol] +# the [amplicol] processes are integrated at leading colour; their events are +# re- or unweighted to color_accuracy in a second pass (lc: no second pass) +color_accuracy = %(amplicol.color_accuracy)s +mode = %(amplicol.mode)s +# unweighting at leading colour (unweight mode); [generation] sets the second pass +first_pass_max_overweight_truncation = %(amplicol.first_pass_max_overweight_truncation)s +first_pass_freeze_max_weight_after = %(amplicol.first_pass_freeze_max_weight_after)s + [multiparticles] $multiparticles diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index c904f8691d..bcc6c6fb68 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6917,6 +6917,17 @@ def default_setup(self): comment="freeze_max_weight_after of the unweighting with the surrogate") self.add_toml_param('surrogate', 'param_card', "", comment="param card for a .so surrogate; \"\" uses Cards/param_card.dat") + # ---------------------------- [amplicol] ---------------------- + # the leading-colour [amplicol] subprocesses are re- or unweighted to + # the colour accuracy in a second pass + self.add_toml_param('amplicol', 'color_accuracy', "full", allowed=["full", "lc"], + comment="colour accuracy of the [amplicol] events, lc skips the second pass") + self.add_toml_param('amplicol', 'mode', "reweight", allowed=["unweight", "reweight"], + comment="unweight or reweight the events to the colour accuracy") + self.add_toml_param('amplicol', 'first_pass_max_overweight_truncation', 0.01, + comment="max_overweight_truncation of the leading-colour unweighting") + self.add_toml_param('amplicol', 'first_pass_freeze_max_weight_after', 10000, + comment="freeze_max_weight_after of the leading-colour unweighting") # ----------------------------- [madnis] ----------------------- self.add_toml_param('madnis', 'enable', False, allowed=["auto", True, False], auto=True, From b74dee3de7eb8fd6860dc4d00de999698e947d3b Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:09:13 +0200 Subject: [PATCH 11/22] mg7: subprocess hooks for other matrix-element backends The matrix elements of a subprocess are loaded by load_matrix_elements, and those of its integrand built by integrand_matrix_elements, which a subprocess of another backend overrides. A flavor can carry its own factor. --- .../mg7/subprocess_madmatrix.py | 75 +++++++++++-------- 1 file changed, 45 insertions(+), 30 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py index 98eb957a93..25ad94d7ea 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py +++ b/madgraph/iolibs/template_files/mg7/subprocess_madmatrix.py @@ -336,6 +336,10 @@ def cut_pids(observable_kwargs): **self.process.scale_kwargs ) + self.load_matrix_elements(all_api_paths, api_path_formats, member_metas) + + def load_matrix_elements(self, all_api_paths, api_path_formats, member_metas): + """The matrix element and [surrogate] of every member subprocess.""" if self.process.run_card["run"]["dummy_matrix_element"]: self.matrix_elements = [None] * len(all_api_paths) else: @@ -377,6 +381,9 @@ def load_surrogate(self, me_path_format: str, member_meta: dict): raise RuntimeError("surrogate loaded at different indices") return apis[0] + def has_weighting_steps(self) -> bool: + return bool(self.surrogates) + def weighting_steps(self) -> list: """The steps after the first unweighting: with a [surrogate], the second pass with the full matrix element.""" @@ -1004,35 +1011,10 @@ def activation(self, name: str) -> ms.MLP.Activation: else: raise ValueError(f"Invalid activation function '{name}'") - def build_integrands( - self, - phasespace: PhaseSpace, - madnis_training: bool = False, - drop_cuts_and_rescale: bool = False - ) -> list[ms.Integrand]: - flavors = [] - flavor_remap = [] - flavor_factors = [] - flavor_mirror = [] - flavor_diff_xs_indices = [] - flavor_subproc_indices = [] - flavor_per_subproc_remap = [] - - for flav in self.meta["flavors"]: - if self.unmerged_meta is not None: - diff_xs_index = flav["subprocess"] - subproc_index = self.meta["subprocesses"][diff_xs_index] - ps_flavor = flav["flavor"] - flavor_diff_xs_indices.append(diff_xs_index) - flavor_subproc_indices.append(subproc_index) - flavor_per_subproc_remap.append(ps_flavor) - flav = self.unmerged_meta[subproc_index]["flavors"][ps_flavor] - flavors.append(flav["options"][0]) - flavor_remap.append(flav["index"]) - flavor_factors.append(len(flav["options"])) - flavor_mirror.append(flav["mirror"]) - - cross_sections = [] + def integrand_matrix_elements(self) -> list: + """The matrix element of every member subprocess in the integrand. Sets + the step_matrix_elements of the second pass of a [surrogate].""" + matrix_elements = [] self.step_matrix_elements = [] for me_index, matrix_element in enumerate(self.matrix_elements): if self.surrogates: @@ -1084,6 +1066,39 @@ def build_integrands( self.process.me_frame, self.process.incoming_count, ) + matrix_elements.append(mat) + return matrix_elements + + def build_integrands( + self, + phasespace: PhaseSpace, + madnis_training: bool = False, + drop_cuts_and_rescale: bool = False + ) -> list[ms.Integrand]: + flavors = [] + flavor_remap = [] + flavor_factors = [] + flavor_mirror = [] + flavor_diff_xs_indices = [] + flavor_subproc_indices = [] + flavor_per_subproc_remap = [] + + for flav in self.meta["flavors"]: + if self.unmerged_meta is not None: + diff_xs_index = flav["subprocess"] + subproc_index = self.meta["subprocesses"][diff_xs_index] + ps_flavor = flav["flavor"] + flavor_diff_xs_indices.append(diff_xs_index) + flavor_subproc_indices.append(subproc_index) + flavor_per_subproc_remap.append(ps_flavor) + flav = self.unmerged_meta[subproc_index]["flavors"][ps_flavor] + flavors.append(flav["options"][0]) + flavor_remap.append(flav["index"]) + flavor_factors.append(flav.get("factor", len(flav["options"]))) + flavor_mirror.append(flav["mirror"]) + + cross_sections = [] + for mat in self.integrand_matrix_elements(): pdf_grid = None if self.process.leptonic else self.process.pdf_grid pdf_arg = None if self.process.leptonic else ms.CachedPdf() cross_sections.append( @@ -1100,7 +1115,7 @@ def build_integrands( pdf2=pdf_arg, input_momentum_fraction=True, decay=self.process.is_decay, - return_matrix_element=bool(self.surrogates), + return_matrix_element=self.has_weighting_steps(), ) ) partial_weights = self.process.systematics_enabled() From 6f2c76f9f88a33ccfafde1acb46a8f561e238c6b Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:09:13 +0200 Subject: [PATCH 12/22] mg7: AmpliCol subprocesses Every colour ordering of an [amplicol] process is integrated at leading colour with the colour-ordered phase-space mapping of its UMAMI channel, and its events are re- or unweighted to full colour with the ratio library. The parameters of the param card are passed to the AmpliCol libraries, which the launcher builds when they are missing. GPUs, gridpacks and dummy matrix elements are refused. --- .../template_files/mg7/madgraph_process.py | 64 ++++-- .../template_files/mg7/subprocess_amplicol.py | 184 ++++++++++++++++++ 2 files changed, 235 insertions(+), 13 deletions(-) create mode 100644 madgraph/iolibs/template_files/mg7/subprocess_amplicol.py diff --git a/madgraph/iolibs/template_files/mg7/madgraph_process.py b/madgraph/iolibs/template_files/mg7/madgraph_process.py index f1f70ae409..a7de60ee76 100644 --- a/madgraph/iolibs/template_files/mg7/madgraph_process.py +++ b/madgraph/iolibs/template_files/mg7/madgraph_process.py @@ -21,6 +21,7 @@ from madgraph.iolibs.template_files.mg7.custom_function_loader import load_custom_function from madgraph.iolibs.template_files.mg7.observables import CutItem, HistItem, build_observable, clean_pids, observable_pids from madgraph.iolibs.template_files.mg7.options import lhapdf_paths +from madgraph.iolibs.template_files.mg7.subprocess_amplicol import AmplicolSubprocess, amplicol_subproc_args, is_amplicol from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import MadgraphSubprocess, PhaseSpace, build_multi_channel_data class MadgraphProcess: @@ -1021,13 +1022,23 @@ def init_generator_config(self) -> None: cfg.vegas_damping = vegas_args["damping"] cfg.max_overweight_truncation = gen_args["max_overweight_truncation"] cfg.freeze_max_weight_after = gen_args["freeze_max_weight_after"] + first_pass = None if self.surrogate is not None and self.surrogate["mode"] == "unweight": + first_pass = self.surrogate + amplicol = self.run_card["amplicol"] + if ( + first_pass is None and amplicol["mode"] == "unweight" + and amplicol["color_accuracy"] != "lc" + and any(is_amplicol(meta) for meta in self.subprocess_data) + ): + first_pass = amplicol + if first_pass is not None: # [generation] applies to the second pass, which gives the events cfg.max_overweight_truncation = ( - self.surrogate["first_pass_max_overweight_truncation"] + first_pass["first_pass_max_overweight_truncation"] ) cfg.freeze_max_weight_after = ( - self.surrogate["first_pass_freeze_max_weight_after"] + first_pass["first_pass_freeze_max_weight_after"] ) cfg.start_batch_size = vegas_args["start_batch_size"] cfg.max_batch_size = vegas_args["max_batch_size"] @@ -1065,11 +1076,19 @@ def init_subprocesses(self) -> None: self.subprocesses = [] if self.merged_subprocess_data is None: for subproc_id, meta in enumerate(self.subprocess_data): - self.subprocesses.append(MadgraphSubprocess(self, meta, subproc_id)) + subproc_class = ( + AmplicolSubprocess if is_amplicol(meta) else MadgraphSubprocess + ) + self.subprocesses.append(subproc_class(self, meta, subproc_id)) else: for subproc_id, meta in enumerate(self.merged_subprocess_data): + # an AmpliCol subprocess is never merged with another one + first = self.subprocess_data[meta["subprocesses"][0]] + subproc_class = ( + AmplicolSubprocess if is_amplicol(first) else MadgraphSubprocess + ) self.subprocesses.append( - MadgraphSubprocess(self, meta, subproc_id, self.subprocess_data) + subproc_class(self, meta, subproc_id, self.subprocess_data) ) def compile_matrix_elements(self) -> list[str]: @@ -1091,10 +1110,14 @@ def compile_matrix_elements(self) -> list[str]: device_names = [device_names] cpu_mode = self.run_card["run"]["cpu_mode"] backends = [backend_of(name, cpu_mode) for name in device_names] - if not self.subprocess_data: + self.compile_amplicol() + madmatrix_data = [ + meta for meta in self.subprocess_data if not is_amplicol(meta) + ] + if not madmatrix_data: return backends - first_proc_path = self.subprocess_data[0]["path"] + first_proc_path = madmatrix_data[0]["path"] subproc_path = os.path.dirname(first_proc_path) fptype = fptype_of(self.run_card["run"]["precision"]) @@ -1117,7 +1140,7 @@ def compile_matrix_elements(self) -> list[str]: log_path = os.path.join(self.run_path, "compile_subprocesses.log") for backend in resolved: missing = [ - meta for meta in self.subprocess_data + meta for meta in madmatrix_data if not os.path.isfile(meta["me_path"].format(device=backend)) ] if not missing: @@ -1138,6 +1161,19 @@ def compile_matrix_elements(self) -> list[str]: ) return resolved + def compile_amplicol(self) -> None: + """Build the libraries of the [amplicol] processes that are missing.""" + paths = sorted({ + meta["path"] for meta in self.subprocess_data + if is_amplicol(meta) and not ( + os.path.isfile(meta["me_path"]) and os.path.isfile(meta["ratio_path"]) + ) + }) + log_path = os.path.join(self.run_path, "compile_amplicol.log") + for path in paths: + logger.info(f"Compiling the AmpliCol matrix elements in {path}, see {log_path}") + self.make_subprocesses(path, [], 1, log_path) + @staticmethod def make_subprocesses( subproc_path: str, args: list[str], nb_core: int, log_path: str @@ -1759,13 +1795,15 @@ def write_lhe_header(self, systematics) -> None: del writer # closes the file (writes the closing tag) def build_lhe_completer(self): - all_mcdata = ( - [subproc.build_multi_channel_data() for subproc in self.subprocesses] - if self.merged_subprocess_data is None else - [build_multi_channel_data(meta, self) for meta in self.subprocess_data] - ) + all_mcdata = [ + None if is_amplicol(meta) + else self.subprocesses[index].build_multi_channel_data() + if self.merged_subprocess_data is None + else build_multi_channel_data(meta, self) + for index, meta in enumerate(self.subprocess_data) + ] subproc_args = [ - ms.SubprocArgs( + amplicol_subproc_args(meta, self) if mcdata is None else ms.SubprocArgs( topologies = [topo[0] for topo in mcdata.topologies], permutations = mcdata.permutations, diagram_indices = mcdata.diagram_indices, diff --git a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py new file mode 100644 index 0000000000..0c66a8ce85 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py @@ -0,0 +1,184 @@ +"""Subprocesses with AmpliCol matrix elements ('[amplicol]' processes). + +Each subprocess is one colour ordering (UMAMI channel) of an AmpliCol provider, +integrated on its own at leading colour with a colour-ordered phase-space +mapping. A second pass re- or unweights its events to full colour with the +ratio of the full-colour to the leading-colour matrix element summed over all +colour orderings, at the helicity of the event. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +from madgraph.iolibs.template_files.mg7.common import lhe_helicities, ms +from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import ( + Channel, + MadgraphSubprocess, + PhaseSpace, +) + +if TYPE_CHECKING: + from madgraph.iolibs.template_files.mg7.madgraph_process import MadgraphProcess + + +def is_amplicol(meta: dict) -> bool: + return meta.get("backend") == "amplicol" + + +def amplicol_parameters(meta: dict, process: MadgraphProcess) -> dict[str, float]: + """The values of the mutable model parameters in the param card.""" + parameters = {} + for name, (block, code) in meta["parameters"].items(): + try: + parameters[name] = float( + process.param_card.get_value(block.lower(), tuple(code)) + ) + except (KeyError, ValueError): + continue + return parameters + + +def chain_topology(incoming_masses: list[float], outgoing_masses: list[float]): + """A t-channel chain without propagators, for the LHE completion.""" + n_out = len(outgoing_masses) + vertices = [["i0", "p0", "o0"]] + vertices += [[f"p{i - 1}", f"p{i}", f"o{i}"] for i in range(1, n_out - 1)] + vertices.append(["i1", f"p{n_out - 2}", f"o{n_out - 1}"]) + propagators = [ms.Propagator() for _ in range(n_out - 1)] + return ms.Topology( + ms.Diagram(incoming_masses, outgoing_masses, propagators, vertices) + ) + + +def amplicol_subproc_args(meta: dict, process: MadgraphProcess) -> ms.SubprocArgs: + """LHE completion of an AmpliCol subprocess: its events have no diagram.""" + return ms.SubprocArgs( + topologies=[chain_topology( + [process.get_mass(pid) for pid in meta["incoming"]], + [process.get_mass(pid) for pid in meta["outgoing"]], + )], + permutations=[[]], + diagram_indices=[[]], + diagram_color_indices=[[]], + color_flows=meta["color_flows"], + pdg_color_types={ + int(key): value for key, value in meta["pdg_color_types"].items() + }, + helicities=lhe_helicities(meta, process.run_card), + pdg_ids=[flavor["options"] for flavor in meta["flavors"]], + ) + + +class AmplicolSubprocess(MadgraphSubprocess): + def load_matrix_elements(self, all_api_paths, api_path_formats, member_metas): + run_args = self.process.run_card["run"] + if any(device != "cpu" for device in self.process.device_types): + raise RuntimeError("[amplicol] processes are only supported on CPUs") + if run_args["dummy_matrix_element"]: + raise RuntimeError("[amplicol] processes need their matrix elements") + if self.process.run_card["gridpack"]["save_gridpack"]: + raise RuntimeError("[amplicol] processes do not support gridpacks yet") + if not self.meta.get("supports_alpha_s", True): + raise RuntimeError( + "the AmpliCol matrix elements of this model cannot take the running " + "alpha_s, as the model has no SMINPUTS 3 parameter" + ) + parameters = amplicol_parameters(self.meta, self.process) + self.channel_index = self.meta["umami_channel"] + self.matrix_elements = [ + self.load(self.meta["me_path"], self.meta["artifact"], parameters) + ] + self.surrogates = [] + self.ratio_matrix_element = None + if self.process.run_card["amplicol"]["color_accuracy"] != "lc": + self.ratio_matrix_element = self.load( + self.meta["ratio_path"], self.meta["ratio_artifact"], parameters + ) + + def load(self, path: str, artifact: str, parameters: dict): + # the libraries find their pyamplicol artifacts from an absolute path + for context in self.process.contexts: + api = context.load_matrix_element(path, os.path.abspath(artifact), parameters) + return api + + def build_multi_channel_data(self): + return None + + def build_multichannel_phasespace(self) -> PhaseSpace: + mapping = ms.PhaseSpaceMapping( + self.incoming_masses + self.outgoing_masses, + self.process.e_cm, + leptonic=self.process.leptonic, + invariant_power=self.process.run_card["phasespace"]["invariant_power"], + mode=ms.PhaseSpaceMapping.color_ordered, + cuts=self.cuts, + color_order=self.meta["phasespace_order"], + ) + prefix = f"subproc{self.subproc_id}.amplicol" + discrete_sym, discrete_flavor = self.build_discrete( + 1, len(self.meta["flavors"]), prefix + ) + channel = Channel( + phasespace_mapping=mapping, + adaptive_mapping=self.build_vegas(mapping, prefix), + discrete_sym=discrete_sym, + discrete_flavor=discrete_flavor, + channel_weight_indices=[0], + name="C", + active_flavors=[], + ) + return PhaseSpace(mode="multichannel", channels=[channel], symfact=[None]) + + def build_flat_phasespace(self) -> PhaseSpace: + return self.build_multichannel_phasespace() + + def simplify_phasespace(self, multi_phasespace, cross_sections) -> PhaseSpace: + return multi_phasespace + + def integrand_matrix_elements(self) -> list: + inputs = [ + ms.MatrixElement.momenta_in, + ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in, + ms.MatrixElement.channel_in, + ms.MatrixElement.random_helicity_in, + ] + outputs = [ + ms.MatrixElement.matrix_element_out, + ms.MatrixElement.color_index_out, + ms.MatrixElement.helicity_index_out, + ] + return [ms.MatrixElement( + self.matrix_elements[0].index(), self.particle_count, inputs, outputs, 1, + True, self.process.me_frame, self.process.incoming_count, + self.channel_index, + )] + + def has_weighting_steps(self) -> bool: + return self.ratio_matrix_element is not None + + def weighting_steps(self) -> list: + """The second pass to full colour.""" + if self.ratio_matrix_element is None: + return [] + ratio = ms.MatrixElement( + self.ratio_matrix_element.index(), + self.particle_count, + [ms.MatrixElement.momenta_in, ms.MatrixElement.flavor_in, + ms.MatrixElement.channel_in, ms.MatrixElement.helicity_in], + [ms.MatrixElement.matrix_element_ratio_out], + 1, + False, + self.process.me_frame, + self.process.incoming_count, + self.channel_index, + ) + gen_args = self.process.run_card["generation"] + return [ms.WeightingStep( + [ratio], + mode=self.process.run_card["amplicol"]["mode"], + max_overweight_truncation=gen_args["max_overweight_truncation"], + freeze_max_weight_after=gen_args["freeze_max_weight_after"], + )] From a8bfeee108fade4f06d4df445f1fcaf381960607 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:11:00 +0200 Subject: [PATCH 13/22] mg7: [amplicol] processes with flavour-grouped legs MG replaces the light quarks and leptons of a process definition by flavour groups, which pyamplicol does not know: the export passes the original flavours of a leg, or the members of its groups. The directory is named with the multiparticle labels of the definition. --- madgraph/iolibs/export_amplicol.py | 45 ++++++++++++++----- madgraph/iolibs/export_cpp.py | 3 +- .../interface/test_amplicol_syntax.py | 42 +++++++++++++++++ 3 files changed, 79 insertions(+), 11 deletions(-) diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index 2f8a35e362..ec8d9e5f0a 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -60,11 +60,35 @@ def import_pyamplicol(): return pyamplicol, umami_reweight -def process_name(procdef): - """A directory name for a process definition, e.g. 'pp_jjj'.""" - string = procdef.nice_string(prefix=False).split('@')[0] - initial, final = [side.split() for side in string.split('>')[:2]] - name = ''.join(initial) + '_' + ''.join(final) +def leg_pdgs(leg, model): + """The pdg ids of a leg: the flavours MG merged into a group id, or the ids + with the group ids expanded.""" + if leg.get('flavor'): + return list(leg.get('flavor')) + merged = model.get('merged_particles') + pdgs = [] + for pid in leg.get('ids'): + members = merged.get(abs(pid), [abs(pid)]) + pdgs.extend(member if pid > 0 else -member for member in members) + return pdgs + + +def process_name(procdef, model, multiparticles): + """A directory name for a process definition, e.g. 'pp_jjj', with the + multiparticle labels of the interface.""" + particles = model.get('particle_dict') + labels = [] + for leg in procdef.get('legs'): + pdgs = leg_pdgs(leg, model) + matches = [label for label, ids in multiparticles.items() + if sorted(ids) == sorted(pdgs)] + preferred = 'j' if leg.get('state') else 'p' + label = preferred if preferred in matches else next(iter(matches), None) + if label is None: + label = ''.join(particles[pid].get_name() for pid in pdgs) + labels.append((leg.get('state'), label)) + name = ''.join(l for state, l in labels if not state) + '_' + \ + ''.join(l for state, l in labels if state) return re.sub(r'[^A-Za-z0-9_]', '', name.replace('~', 'x').replace('+', 'p').replace('-', 'm')) @@ -75,7 +99,7 @@ def process_expression(procdef, model): multiparticles = {} labels = [] for leg in procdef.get('legs'): - names = [particles[pid].get_name() for pid in leg.get('ids')] + names = [particles[pid].get_name() for pid in leg_pdgs(leg, model)] if len(names) == 1: labels.append(names[0]) continue @@ -141,8 +165,8 @@ def representative_processes(proc_defs, model): processes = [] for procdef in proc_defs: legs = procdef.get('legs') - initial = [leg.get('ids') for leg in legs if not leg.get('state')] - final = [leg.get('ids')[0] for leg in legs if leg.get('state')] + initial = [leg_pdgs(leg, model) for leg in legs if not leg.get('state')] + final = [leg_pdgs(leg, model)[0] for leg in legs if leg.get('state')] for ids in itertools.product(*initial): pdgs = list(ids) + final processes.append(base_objects.Process({ @@ -228,7 +252,8 @@ def parameter_references(metadata): return references -def export_amplicol_processes(dir_path, proc_defs, model, nb_core=None): +def export_amplicol_processes(dir_path, proc_defs, model, nb_core=None, + multiparticles=None): """Write SubProcesses/A_ for every '[amplicol]' process definition and return the subprocesses.json entries.""" pyamplicol, umami_reweight = import_pyamplicol() @@ -238,7 +263,7 @@ def export_amplicol_processes(dir_path, proc_defs, model, nb_core=None): rusticol_config = shutil.which('rusticol-config') or 'rusticol-config' entries = [] for number, procdef in enumerate(proc_defs, start=1): - name = 'A%d_%s' % (number, process_name(procdef)) + name = 'A%d_%s' % (number, process_name(procdef, model, multiparticles or {})) path = pjoin(subproc_dir, name) rel_dir = pjoin('SubProcesses', name) expression, multiparticles = process_expression(procdef, model) diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index 7e34afcb35..dcd3d0be43 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3511,6 +3511,7 @@ def pass_information_from_cmd(self, cmd): self.amplicol_defs = getattr(cmd, '_curr_amplicol_defs', None) or [] self.amplicol_model = cmd._curr_model self.amplicol_nb_core = cmd.options.get('nb_core') + self.amplicol_multiparticles = dict(getattr(cmd, '_multiparticles', {})) def amplicol_processes(self): """MG processes standing for the [amplicol] processes.""" @@ -3527,7 +3528,7 @@ def export_amplicol(self): import madgraph.iolibs.export_amplicol as export_amplicol for entry in export_amplicol.export_amplicol_processes( self.dir_path, self.amplicol_defs, self.amplicol_model, - self.amplicol_nb_core): + self.amplicol_nb_core, self.amplicol_multiparticles): # never merged with another subprocess self.merged_subprocesses[('amplicol', len(self.process_info))].append( (len(self.process_info), [])) diff --git a/tests/unit_tests/interface/test_amplicol_syntax.py b/tests/unit_tests/interface/test_amplicol_syntax.py index 1e20f0d35e..493801d065 100644 --- a/tests/unit_tests/interface/test_amplicol_syntax.py +++ b/tests/unit_tests/interface/test_amplicol_syntax.py @@ -83,3 +83,45 @@ def test_unsupported_options_are_refused(self): 'p p > j j / b [amplicol]', 'z > e+ e- [amplicol]']: self.assertRaises(InvalidCmd, self.cmd.do_generate, line) + + +class AmplicolExportTest(unittest.TestCase): + """The process definitions as pyamplicol sees them, without pyamplicol.""" + + def setUp(self): + self.cmd = MGCmd.MasterCmd() + self.cmd.exec_cmd('import model sm', printcmd=False, precmd=True) + self.cmd.exec_cmd('define q = u d', printcmd=False, precmd=True) + self.cmd.exec_cmd('define qx = u~ d~', printcmd=False, precmd=True) + + def export(self, line): + import madgraph.iolibs.export_amplicol as export_amplicol + self.cmd.exec_cmd('generate %s [amplicol]' % line, printcmd=False, + precmd=True) + procdef = self.cmd._curr_amplicol_defs[0] + model = self.cmd._curr_model + return (export_amplicol.process_name(procdef, model, + self.cmd._multiparticles), + export_amplicol.process_expression(procdef, model), + export_amplicol.representative_processes([procdef], model)) + + def test_flavour_groups_are_expanded(self): + """MG replaces the light quarks by a flavour group, which pyamplicol + does not know: the original flavours are passed on.""" + name, expression, processes = self.export('u u~ > e+ e- g') + self.assertEqual(name, 'uux_epemg') + self.assertEqual(expression, ('u u~ > e+ e- g', {})) + self.assertEqual(len(processes), 1) + name, expression, processes = self.export('q qx > e+ e- g') + self.assertEqual(name, 'qqx_epemg') + self.assertEqual(expression, ('mp0 mp1 > e+ e- g', + {'mp0': ['u', 'd'], 'mp1': ['u~', 'd~']})) + self.assertEqual(len(processes), 4) + + def test_multiparticles(self): + name, (expression, multiparticles), processes = self.export('p p > j j') + self.assertEqual(name, 'pp_jj') + self.assertEqual(expression, 'mp0 mp0 > mp0 mp0') + self.assertEqual(sorted(multiparticles['mp0']), + sorted(['g', 'u', 'd', 's', 'c', 'u~', 'd~', 's~', 'c~'])) + self.assertEqual(len(processes), 81) From 4e6d7164ab778598ba89d7b4053889eb93883ea4 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:22:04 +0200 Subject: [PATCH 14/22] mg7: coupling orders of the [amplicol] processes pyamplicol gets the largest orders of the couplings that MG selects by default, the minimal WEIGHTED order of the process definition, instead of its minimal orders per process, which differ for the processes of a multiparticle definition. The power of alpha_s follows from them. pyamplicol runs with one worker, as Symbolica's restricted mode allows a single thread. --- madgraph/iolibs/export_amplicol.py | 48 ++++++++++++------- madgraph/iolibs/export_cpp.py | 3 +- .../interface/test_amplicol_syntax.py | 11 +++++ 3 files changed, 44 insertions(+), 18 deletions(-) diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index ec8d9e5f0a..fe3315fb6c 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -116,6 +116,20 @@ def process_expression(procdef, model): return '%s > %s' % (' '.join(initial), ' '.join(final)), multiparticles +def coupling_orders(procdef, model): + """The largest orders of the couplings above the lowest hierarchy, as MG + selects them by default: the minimal WEIGHTED order of the definition, + with n - 2 vertices at tree level.""" + import madgraph.core.diagram_generation as diagram_generation + weighted = diagram_generation.MultiProcess.find_optimal_process_orders( + procdef)['WEIGHTED'] + hierarchy = model.get('order_hierarchy') + lowest = min(hierarchy.values()) + vertices = len(procdef.get('legs')) - 2 + return {order: max((weighted - vertices * lowest) // (weight - lowest), 0) + for order, weight in hierarchy.items() if weight > lowest} + + def model_source(pyamplicol, model, cache_path): """The UFO model of the output with its restriction, compiled for pyamplicol once per output.""" @@ -139,17 +153,20 @@ def model_source(pyamplicol, model, cache_path): def generate_artifact(pyamplicol, artifact, accuracy, expression, multiparticles, - model, workers): + orders, model): from pyamplicol.config import resolve_config + # one worker: Symbolica's restricted mode allows a single thread config = resolve_config({ 'schema_version': 1, 'action': 'generate', 'color': {'accuracy': accuracy}, - 'process': {'multiparticles': multiparticles}, - 'generation': {'workers': workers, 'emit_api_bundle': True, + 'process': {'multiparticles': multiparticles, + 'coupling_order_policy': 'explicit', + 'max_coupling_orders': orders}, + 'generation': {'workers': 1, 'emit_api_bundle': True, 'relation_discovery': {'mode': 'off'}}, 'evaluator': {'backend': 'jit', 'execution_mode': 'recurrence', - 'optimization': {'cores': workers}, + 'optimization': {'cores': 1}, 'jit': {'optimization_level': 2}}, }) pyamplicol.Generator(config).generate(expression, artifact, model=model, @@ -189,7 +206,7 @@ def color_types(model, pdgs): return types -def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path): +def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power): """One subprocesses.json entry per channel of a UMAMI provider.""" pdg_ids = metadata['pdg_ids'] provider = metadata['provider'] @@ -210,8 +227,6 @@ def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path): 'mirror': False, }) first = flavors[0]['options'][0] - colored = sum(1 for pdg in first - if abs(model.get_particle(pdg).get_color()) > 1) entries.append({ 'backend': 'amplicol', 'incoming': first[:incoming], @@ -227,9 +242,7 @@ def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path): 'masses': provider['masses'], 'supports_alpha_s': provider['supports_alpha_s'], 'flavors': flavors, - # tree level with the minimal coupling orders: every coloured - # parton beyond the first two comes with one power of alpha_s - 'qcd_power': max(colored - 2, 0), + 'qcd_power': qcd_power, 'color_flows': metadata['color_flows'], 'pdg_color_types': color_types( model, {abs(pdg) for pdgs in pdg_ids for pdg in pdgs}), @@ -252,12 +265,10 @@ def parameter_references(metadata): return references -def export_amplicol_processes(dir_path, proc_defs, model, nb_core=None, - multiparticles=None): +def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None): """Write SubProcesses/A_ for every '[amplicol]' process definition and return the subprocesses.json entries.""" pyamplicol, umami_reweight = import_pyamplicol() - workers = nb_core or 1 subproc_dir = pjoin(dir_path, 'SubProcesses') source = model_source(pyamplicol, model, pjoin(subproc_dir, MODEL_FILE)) rusticol_config = shutil.which('rusticol-config') or 'rusticol-config' @@ -267,12 +278,16 @@ def export_amplicol_processes(dir_path, proc_defs, model, nb_core=None, path = pjoin(subproc_dir, name) rel_dir = pjoin('SubProcesses', name) expression, multiparticles = process_expression(procdef, model) + orders = coupling_orders(procdef, model) + # every vertex carries one coupling: the rest are powers of g_s + qcd_power = (len(procdef.get('legs')) - 2 - sum(orders.values()) + if len(orders) == 1 else -1) logger.info('Generating the AmpliCol matrix elements of %s in %s' % (procdef.nice_string(prefix=False), path)) os.makedirs(path, exist_ok=True) for accuracy in ('lc', 'full'): generate_artifact(pyamplicol, pjoin(path, accuracy), accuracy, - expression, multiparticles, source, workers) + expression, multiparticles, orders, source) provider_ids = umami_reweight.write_reweight_provider( pjoin(path, 'lc'), pjoin(path, 'full'), pjoin(path, 'reweight')) documents = umami_reweight.load_umami_metadata(pjoin(path, 'lc')) @@ -286,12 +301,13 @@ def export_amplicol_processes(dir_path, proc_defs, model, nb_core=None, copies.append('\tcp $(BUILD)/reweight/libreweight_%s.so ../../%s' % (provider_id, ratio_path)) entries.extend(subprocess_entries(metadata, model, rel_dir, me_path, - ratio_path)) + ratio_path, qcd_power)) with open(pjoin(path, 'makefile'), 'w') as stream: stream.write(MAKEFILE % {'rusticol_config': rusticol_config, 'copies': '\n'.join(copies)}) with open(pjoin(path, 'process.json'), 'w') as stream: json.dump({'process': procdef.nice_string(prefix=False), 'expression': expression, - 'multiparticles': multiparticles}, stream, indent=2) + 'multiparticles': multiparticles, + 'max_coupling_orders': orders}, stream, indent=2) return entries diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index dcd3d0be43..de2d2f116f 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3510,7 +3510,6 @@ def pass_information_from_cmd(self, cmd): self.proc_defs = getattr(cmd, '_curr_proc_defs', None) self.amplicol_defs = getattr(cmd, '_curr_amplicol_defs', None) or [] self.amplicol_model = cmd._curr_model - self.amplicol_nb_core = cmd.options.get('nb_core') self.amplicol_multiparticles = dict(getattr(cmd, '_multiparticles', {})) def amplicol_processes(self): @@ -3528,7 +3527,7 @@ def export_amplicol(self): import madgraph.iolibs.export_amplicol as export_amplicol for entry in export_amplicol.export_amplicol_processes( self.dir_path, self.amplicol_defs, self.amplicol_model, - self.amplicol_nb_core, self.amplicol_multiparticles): + self.amplicol_multiparticles): # never merged with another subprocess self.merged_subprocesses[('amplicol', len(self.process_info))].append( (len(self.process_info), [])) diff --git a/tests/unit_tests/interface/test_amplicol_syntax.py b/tests/unit_tests/interface/test_amplicol_syntax.py index 493801d065..d52a290ad6 100644 --- a/tests/unit_tests/interface/test_amplicol_syntax.py +++ b/tests/unit_tests/interface/test_amplicol_syntax.py @@ -125,3 +125,14 @@ def test_multiparticles(self): self.assertEqual(sorted(multiparticles['mp0']), sorted(['g', 'u', 'd', 's', 'c', 'u~', 'd~', 's~', 'c~'])) self.assertEqual(len(processes), 81) + + def test_coupling_orders(self): + """The orders MG selects by default, as the largest non-QCD orders.""" + import madgraph.iolibs.export_amplicol as export_amplicol + for line, orders in [('p p > j j j', {'QED': 0}), + ('p p > e+ e- j', {'QED': 2}), + ('p p > w+ j j', {'QED': 1})]: + self.cmd.exec_cmd('generate %s [amplicol]' % line, printcmd=False, + precmd=True) + self.assertEqual(export_amplicol.coupling_orders( + self.cmd._curr_amplicol_defs[0], self.cmd._curr_model), orders) From 327e2b5dc7191bc0750f483278b73b600b1cd237 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:22:04 +0200 Subject: [PATCH 15/22] mg7: refuse cuts that are not flavour blind for [amplicol] The flavours of an AmpliCol channel share its momenta and its cuts, built for the first flavour, so every cut has to select the same particles for all of them. --- .../template_files/mg7/subprocess_amplicol.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py index 0c66a8ce85..eedabb5118 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py +++ b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING from madgraph.iolibs.template_files.mg7.common import lhe_helicities, ms +from madgraph.iolibs.template_files.mg7.observables import clean_pids from madgraph.iolibs.template_files.mg7.subprocess_madmatrix import ( Channel, MadgraphSubprocess, @@ -72,6 +73,27 @@ def amplicol_subproc_args(meta: dict, process: MadgraphProcess) -> ms.SubprocArg class AmplicolSubprocess(MadgraphSubprocess): + def __init__(self, process, meta, subproc_id, unmerged_meta=None): + super().__init__(process, meta, subproc_id, unmerged_meta) + self.check_flavor_blind_cuts() + + def check_flavor_blind_cuts(self): + """AmpliCol groups the flavours and identical-particle orderings that + share a matrix element into one channel, so its cuts, built for the + first flavour, must select the same particles for all of them.""" + options = [ + clean_pids(pdgs) for flavor in self.meta["flavors"] + for pdgs in flavor["options"] + ] + for item in self.process.cut_data: + for group in item.observable_kwargs["select_pids"]: + if len({tuple(pid in group for pid in pdgs) for pdgs in options}) > 1: + raise RuntimeError( + f"the cut {item.observable_kwargs['name']} does not select " + f"the same particles in all flavours of the [amplicol] " + f"subprocess {self.meta['incoming']} > {self.meta['outgoing']}" + ) + def load_matrix_elements(self, all_api_paths, api_path_formats, member_metas): run_args = self.process.run_card["run"] if any(device != "cpu" for device in self.process.device_types): From b38793b27c1c9dea3b2f73939b3b4101bfabc1a6 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:22:05 +0200 Subject: [PATCH 16/22] tests: mg7 [amplicol] cross sections against madmatrix gg > ggg in both colour modes and at leading colour, p p > j j j, and an output mixing madmatrix and [amplicol] processes. --- tests/acceptance_tests/test_mg7_amplicol.py | 158 ++++++++++++++++++++ 1 file changed, 158 insertions(+) create mode 100644 tests/acceptance_tests/test_mg7_amplicol.py diff --git a/tests/acceptance_tests/test_mg7_amplicol.py b/tests/acceptance_tests/test_mg7_amplicol.py new file mode 100644 index 0000000000..539f31ffbd --- /dev/null +++ b/tests/acceptance_tests/test_mg7_amplicol.py @@ -0,0 +1,158 @@ +################################################################################ +# +# Copyright (c) 2026 The MadGraph5_aMC@NLO Development team and Contributors +# +# This file is a part of the MadGraph5_aMC@NLO project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph5_aMC@NLO license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""mg7 runs of '[amplicol]' processes against the same processes with madmatrix. + +The AmpliCol subprocesses are integrated at leading colour and re- or +unweighted to full colour, so their cross sections must agree with the +madmatrix ones within the Monte Carlo errors, alone and mixed with madmatrix +processes in one output. Needs madspace, pyamplicol and LHAPDF; the tests skip +otherwise. Knobs: MG7_AMPLICOL_EVENTS (default 10000) and MG7_AMPLICOL_NSIGMA +(default 4). + + ./tests/test_manager.py test_mg7_amplicol -pA -t0 +""" + +import glob +import gzip +import json +import os +import re +import shutil +import subprocess +import sys +import tempfile +import unittest + +import madgraph.interface.master_interface as MGCmd +import madgraph.various.misc as misc + +pjoin = os.path.join + +_EVENTS = int(os.environ.get('MG7_AMPLICOL_EVENTS', 10000)) +_NSIGMA = float(os.environ.get('MG7_AMPLICOL_NSIGMA', 4)) +_PDF = 'NNPDF23_lo_as_0130_qed' + + +def _edit_run_card(path, **amplicol): + card = open(path).read() + card = card.replace('fixed_ren_scale = false', 'fixed_ren_scale = true') + card = card.replace('fixed_fact_scale = false', 'fixed_fact_scale = true') + card = re.sub(r'(?m)^pdf = ".*"$', 'pdf = "%s"' % _PDF, card) + card = re.sub(r'(?m)^events = \d+', 'events = %d' % _EVENTS, card) + card = re.sub(r'(?m)^seed = -1', 'seed = 7', card) + card = re.sub(r'(?m)^make_plots = true', 'make_plots = false', card) + card = re.sub(r'(\[systematics\]\nenable = )true', r'\1false', card) + card = re.sub(r'(\[madnis\]\n(?:#.*\n)*enable = )"auto"', r'\1false', card) + head, section = card.split('\n[amplicol]\n') + for key, value in amplicol.items(): + section, count = re.subn(r'(?m)^%s = .*$' % key, + '%s = "%s"' % (key, value), section, count=1) + assert count == 1, key + card = head + '\n[amplicol]\n' + section + open(path, 'w').write(card) + + +def _color_balanced(lhe_path): + """Whether every colour line of every event starts and ends once.""" + with gzip.open(lhe_path, 'rt') as stream: + text = stream.read() + for event in re.findall(r'\n(.*?)', text, re.S): + lines = [line.split() for line in event.strip().splitlines()[1:] + if line.strip() and not line.startswith('<')] + tags = {} + for particle in lines: + status, color, anticolor = int(particle[1]), int(particle[4]), int(particle[5]) + if status == -1: + color, anticolor = anticolor, color + for tag, sign in ((color, 1), (anticolor, -1)): + if tag: + tags[tag] = tags.get(tag, 0) + sign + if any(tags.values()): + return False + return True + + +class MG7AmplicolTest(unittest.TestCase): + + def setUp(self): + try: + import madspace # noqa: F401 + import pyamplicol # noqa: F401 + except ImportError as error: + self.skipTest('madspace or pyamplicol unavailable: %s' % error) + if not misc.resolve_lhapdf().find_set(_PDF): + self.skipTest('%s LHAPDF data not found' % _PDF) + self.path = tempfile.mkdtemp(prefix='mg7_amplicol_') + + def tearDown(self): + shutil.rmtree(self.path, ignore_errors=True) + + def run_mg7(self, name, processes, **amplicol): + """Output and run the processes, returns the cross section, its error + and the run directory.""" + run_dir = pjoin(self.path, name) + mg = MGCmd.MasterCmd() + mg.no_notification() + mg.exec_cmd('set automatic_html_opening False --no_save') + for index, process in enumerate(processes): + mg.exec_cmd('%s %s' % ('generate' if index == 0 else 'add process', + process)) + mg.exec_cmd('output mg7 %s' % run_dir) + _edit_run_card(pjoin(run_dir, 'Cards', 'run_card.toml'), **amplicol) + log = pjoin(run_dir, 'generate_events.log') + with open(log, 'w') as stream: + status = subprocess.call( + [sys.executable, pjoin(run_dir, 'bin', 'generate_events'), '-f'], + cwd=run_dir, stdout=stream, stderr=subprocess.STDOUT) + self.assertEqual(status, 0, open(log).read()[-3000:]) + info = sorted(glob.glob(pjoin(run_dir, 'Events', '*', 'info.json')))[-1] + with open(info) as stream: + result = json.load(stream)['process'] + return float(result['mean']), float(result['error']), run_dir + + def assert_compatible(self, first, second, factor=1.): + (mean1, error1), (mean2, error2) = first, second + error = (error1**2 + (factor * error2)**2)**0.5 + self.assertLess(abs(mean1 - factor * mean2), _NSIGMA * error, + '%g +- %g vs %g +- %g' % (mean1, error1, factor * mean2, + factor * error2)) + + def test_gg_ggg(self): + """Full colour in both modes, and leading colour, which differs by + (N^2 - 1) / N^2 for five gluons.""" + madmatrix = self.run_mg7('mm', ['g g > g g g'])[:2] + reweight = self.run_mg7('rw', ['g g > g g g [amplicol]']) + self.assert_compatible(reweight[:2], madmatrix) + events = glob.glob(pjoin(reweight[2], 'Events', '*', 'events.lhe.gz')) + if events: + self.assertTrue(_color_balanced(events[-1])) + unweight = self.run_mg7('uw', ['g g > g g g [amplicol]'], mode='unweight') + self.assert_compatible(unweight[:2], madmatrix) + leading = self.run_mg7('lc', ['g g > g g g [amplicol]'], + color_accuracy='lc') + self.assert_compatible(madmatrix, leading[:2], 8. / 9.) + + def test_pp_jjj(self): + madmatrix = self.run_mg7('mm', ['p p > j j j'])[:2] + amplicol = self.run_mg7('amp', ['p p > j j j [amplicol]'])[:2] + self.assert_compatible(amplicol, madmatrix) + + def test_mixed(self): + """[amplicol] and madmatrix processes in one output.""" + jj = self.run_mg7('jj', ['p p > j j'])[:2] + ggg = self.run_mg7('ggg', ['g g > g g g'])[:2] + mixed = self.run_mg7('mixed', ['p p > j j', 'g g > g g g [amplicol]'])[:2] + self.assert_compatible( + mixed, (jj[0] + ggg[0], (jj[1]**2 + ggg[1]**2)**0.5)) From 3ace55773a20e3430792f604a6e19b9292162b04 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:28:13 +0200 Subject: [PATCH 17/22] mg7: [amplicol] fixes for quark processes and models without alpha_s The LHE completion takes the masses of antiparticles too. A model without an alpha_s parameter can be used with a fixed renormalization scale, at the alpha_s of the model. The p p > j j j test is skipped until pyamplicol drops the processes without amplitudes from its multiparticle expansion and its UFO quark processes agree with madmatrix. --- .../template_files/mg7/subprocess_amplicol.py | 13 ++++++++----- tests/acceptance_tests/test_mg7_amplicol.py | 3 +++ 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py index eedabb5118..b75ca6b990 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py +++ b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py @@ -57,8 +57,8 @@ def amplicol_subproc_args(meta: dict, process: MadgraphProcess) -> ms.SubprocArg """LHE completion of an AmpliCol subprocess: its events have no diagram.""" return ms.SubprocArgs( topologies=[chain_topology( - [process.get_mass(pid) for pid in meta["incoming"]], - [process.get_mass(pid) for pid in meta["outgoing"]], + [process.get_mass(pid) for pid in clean_pids(meta["incoming"])], + [process.get_mass(pid) for pid in clean_pids(meta["outgoing"])], )], permutations=[[]], diagram_indices=[[]], @@ -102,10 +102,12 @@ def load_matrix_elements(self, all_api_paths, api_path_formats, member_metas): raise RuntimeError("[amplicol] processes need their matrix elements") if self.process.run_card["gridpack"]["save_gridpack"]: raise RuntimeError("[amplicol] processes do not support gridpacks yet") - if not self.meta.get("supports_alpha_s", True): + self.supports_alpha_s = self.meta.get("supports_alpha_s", True) + if not self.supports_alpha_s and not self.process.run_card["beam"]["fixed_ren_scale"]: raise RuntimeError( "the AmpliCol matrix elements of this model cannot take the running " - "alpha_s, as the model has no SMINPUTS 3 parameter" + "alpha_s, as the model has no SMINPUTS 3 parameter: use a fixed " + "renormalization scale" ) parameters = amplicol_parameters(self.meta, self.process) self.channel_index = self.meta["umami_channel"] @@ -160,9 +162,10 @@ def simplify_phasespace(self, multi_phasespace, cross_sections) -> PhaseSpace: return multi_phasespace def integrand_matrix_elements(self) -> list: + # without an alpha_s input, the model value of the fixed scale is used inputs = [ ms.MatrixElement.momenta_in, - ms.MatrixElement.alpha_s_in, + *([ms.MatrixElement.alpha_s_in] if self.supports_alpha_s else []), ms.MatrixElement.flavor_in, ms.MatrixElement.channel_in, ms.MatrixElement.random_helicity_in, diff --git a/tests/acceptance_tests/test_mg7_amplicol.py b/tests/acceptance_tests/test_mg7_amplicol.py index 539f31ffbd..d8fee330c8 100644 --- a/tests/acceptance_tests/test_mg7_amplicol.py +++ b/tests/acceptance_tests/test_mg7_amplicol.py @@ -144,6 +144,9 @@ def test_gg_ggg(self): color_accuracy='lc') self.assert_compatible(madmatrix, leading[:2], 8. / 9.) + @unittest.skip('pyamplicol expands p p > j j j into processes without ' + 'amplitudes, and its quark processes from the sm UFO model ' + 'differ from madmatrix (the built-in pyamplicol model agrees)') def test_pp_jjj(self): madmatrix = self.run_mg7('mm', ['p p > j j j'])[:2] amplicol = self.run_mg7('amp', ['p p > j j j [amplicol]'])[:2] From f0dd992defb9682e6e43d54d3769d1de2c3ee410 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 14 Oct 2026 23:47:17 +0200 Subject: [PATCH 18/22] mg7: no madmatrix files in outputs with only [amplicol] processes Without MG diagrams the output gets no src/, backend/, madmatrix makefiles or model code, which nothing would build: only the cards, the launcher and the [amplicol] subprocesses. --- madgraph/iolibs/export_cpp.py | 2 ++ madmatrix/output.py | 8 ++++++++ 2 files changed, 10 insertions(+) diff --git a/madgraph/iolibs/export_cpp.py b/madgraph/iolibs/export_cpp.py index de2d2f116f..e5d766e9c0 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3509,6 +3509,8 @@ def pass_information_from_cmd(self, cmd): [amplicol] processes.""" self.proc_defs = getattr(cmd, '_curr_proc_defs', None) self.amplicol_defs = getattr(cmd, '_curr_amplicol_defs', None) or [] + # without MG diagrams there is no madmatrix code to write or build + self.has_madmatrix = bool(cmd._curr_amps) or not self.amplicol_defs self.amplicol_model = cmd._curr_model self.amplicol_multiparticles = dict(getattr(cmd, '_multiparticles', {})) diff --git a/madmatrix/output.py b/madmatrix/output.py index 6e9cc06a88..7b299e5576 100644 --- a/madmatrix/output.py +++ b/madmatrix/output.py @@ -216,6 +216,12 @@ def get_makefile_replace_dict(self, model): # AV - overload the default version: create CMake directory, do not create lib directory def copy_template(self, model): + if not getattr(self, 'has_madmatrix', True): + # only the cards and the launcher: no src/, backend/ or makefiles + self.dirs_to_create = ['bin', 'lib', 'Cards', 'SubProcesses'] + self.from_template = {key: self.from_template[key] for key in ('.', 'Cards')} + self.template_src_make = self.template_Sub_make = None + return super().copy_template(model) super().copy_template(model) # Copy Arithmetics headers for the double-word expansion (FPTYPE=e) arithmetics_src = pjoin(self.madmatrix_templates, 'Arithmetics') @@ -293,6 +299,8 @@ def generate_subprocess_directory(self, matrix_element, cpp_helas_call_writer, p def convert_model(self, model, wanted_lorentz=[], wanted_couplings=[], **opts): # **opts: an option meant for one exporter (npwave, for the dual HELAS # libraries of P-wave bound states) is passed by keyword to all of them + if not getattr(self, 'has_madmatrix', True): + return if hasattr(model , 'cudacpp_wanted_ordered_couplings'): wanted_couplings = model.cudacpp_wanted_ordered_couplings del model.cudacpp_wanted_ordered_couplings From 4de6aedd46852650ee4f41517dd399d564a7a480 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 15 Oct 2026 14:59:07 +0200 Subject: [PATCH 19/22] mg7: load the [amplicol] libraries once for all their channels Each colour-ordered channel is its own subprocess, and each loaded the same two libraries again, with an instance per thread that reads the whole pyamplicol artifact. --- .../iolibs/template_files/mg7/madgraph_process.py | 2 ++ .../iolibs/template_files/mg7/subprocess_amplicol.py | 12 ++++++++---- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/madgraph_process.py b/madgraph/iolibs/template_files/mg7/madgraph_process.py index a7de60ee76..0f910aba49 100644 --- a/madgraph/iolibs/template_files/mg7/madgraph_process.py +++ b/madgraph/iolibs/template_files/mg7/madgraph_process.py @@ -251,6 +251,8 @@ def init_event_dir(self) -> None: def init_context(self) -> None: device_names = self.run_card["run"]["device"] self.contexts = [] + # AmpliCol libraries, shared by all channels of their process + self.amplicol_apis = {} self.device_types = [] self.devices = [] self.pool_sizes = [] diff --git a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py index b75ca6b990..807edefd08 100644 --- a/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py +++ b/madgraph/iolibs/template_files/mg7/subprocess_amplicol.py @@ -122,10 +122,14 @@ def load_matrix_elements(self, all_api_paths, api_path_formats, member_metas): ) def load(self, path: str, artifact: str, parameters: dict): - # the libraries find their pyamplicol artifacts from an absolute path - for context in self.process.contexts: - api = context.load_matrix_element(path, os.path.abspath(artifact), parameters) - return api + # loaded once for all channels; each load reads the whole pyamplicol artifact + apis = self.process.amplicol_apis + if path not in apis: + for context in self.process.contexts: + apis[path] = context.load_matrix_element( + path, os.path.abspath(artifact), parameters + ) + return apis[path] def build_multi_channel_data(self): return None From d0567581811ccfcb9f7082a9d793fa1c1eb47356 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 15 Oct 2026 15:41:41 +0200 Subject: [PATCH 20/22] mg7: mirrored beams of the [amplicol] processes pyamplicol keeps one process per unordered initial state, so p p > j j j had g u > ... but not u g > ..., and its cross section was 21% low. Such flavours are integrated with the mirrored beams, as madmatrix does, when the process definition allows the other order. --- madgraph/iolibs/export_amplicol.py | 27 +++++++++++++++---- .../interface/test_amplicol_syntax.py | 15 +++++++++++ 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index fe3315fb6c..f8b919587a 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -206,11 +206,25 @@ def color_types(model, pdgs): return types -def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power): - """One subprocesses.json entry per channel of a UMAMI provider.""" +def mirrored(pdgs, incoming, initial_legs, exported): + """Whether the process with swapped beams is requested but not exported: + pyamplicol keeps one process per unordered initial state.""" + if incoming != 2 or pdgs[0] == pdgs[1]: + return False + swapped = ((pdgs[1], pdgs[0]), tuple(sorted(pdgs[2:]))) + return (pdgs[1] in initial_legs[0] and pdgs[0] in initial_legs[1] + and swapped not in exported) + + +def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power, + initial_legs): + """One subprocesses.json entry per channel of a UMAMI provider; initial_legs + are the pdg ids allowed for each initial leg.""" pdg_ids = metadata['pdg_ids'] provider = metadata['provider'] incoming = provider['incoming_count'] + exported = {(tuple(pdgs[:incoming]), tuple(sorted(pdgs[incoming:]))) + for pdgs in pdg_ids} entries = [] for channel_index, channel in enumerate(metadata['channels']): flavors = [] @@ -220,11 +234,12 @@ def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power) 'AmpliCol contributions shared by several channels are not ' 'supported yet') for term in process['matrix_elements']: + pdgs = pdg_ids[term['pdg_ids']] flavors.append({ 'index': flavor_index, - 'options': [pdg_ids[term['pdg_ids']]], + 'options': [pdgs], 'factor': term['factor'], - 'mirror': False, + 'mirror': mirrored(pdgs, incoming, initial_legs, exported), }) first = flavors[0]['options'][0] entries.append({ @@ -300,8 +315,10 @@ def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None): copies.append('\tcp $(BUILD)/lc/%s ../../%s' % (library, me_path)) copies.append('\tcp $(BUILD)/reweight/libreweight_%s.so ../../%s' % (provider_id, ratio_path)) + initial_legs = [set(leg_pdgs(leg, model)) + for leg in procdef.get('legs') if not leg.get('state')] entries.extend(subprocess_entries(metadata, model, rel_dir, me_path, - ratio_path, qcd_power)) + ratio_path, qcd_power, initial_legs)) with open(pjoin(path, 'makefile'), 'w') as stream: stream.write(MAKEFILE % {'rusticol_config': rusticol_config, 'copies': '\n'.join(copies)}) diff --git a/tests/unit_tests/interface/test_amplicol_syntax.py b/tests/unit_tests/interface/test_amplicol_syntax.py index d52a290ad6..854b7bb3ff 100644 --- a/tests/unit_tests/interface/test_amplicol_syntax.py +++ b/tests/unit_tests/interface/test_amplicol_syntax.py @@ -136,3 +136,18 @@ def test_coupling_orders(self): precmd=True) self.assertEqual(export_amplicol.coupling_orders( self.cmd._curr_amplicol_defs[0], self.cmd._curr_model), orders) + + def test_mirrored_beams(self): + """pyamplicol keeps one process per unordered initial state; the other + beam order is integrated as the mirror where the definition allows it.""" + import madgraph.iolibs.export_amplicol as export_amplicol + mirrored = export_amplicol.mirrored + partons = {21, 1, 2, -1, -2} + exported = {((21, 2), (2, 21, 21)), ((2, 2), (2, 2, 21)), + ((2, -2), (-2, 2, 21)), ((-2, 2), (-2, 2, 21))} + self.assertTrue(mirrored([21, 2, 21, 21, 2], 2, [partons, partons], exported)) + self.assertFalse(mirrored([2, 2, 2, 2, 21], 2, [partons, partons], exported)) + # both beam orders are exported + self.assertFalse(mirrored([2, -2, 2, -2, 21], 2, [partons, partons], exported)) + # g u > u g g only + self.assertFalse(mirrored([21, 2, 21, 21, 2], 2, [{21}, {2}], exported)) From 8f045497cecd727c5eb60c15f5136c1f73fbe9cd Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 15 Oct 2026 15:41:43 +0200 Subject: [PATCH 21/22] tests: [amplicol] p p > j j j and g g > t t~ g against madmatrix --- tests/acceptance_tests/test_mg7_amplicol.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/acceptance_tests/test_mg7_amplicol.py b/tests/acceptance_tests/test_mg7_amplicol.py index d8fee330c8..880de0ab8c 100644 --- a/tests/acceptance_tests/test_mg7_amplicol.py +++ b/tests/acceptance_tests/test_mg7_amplicol.py @@ -144,14 +144,18 @@ def test_gg_ggg(self): color_accuracy='lc') self.assert_compatible(madmatrix, leading[:2], 8. / 9.) - @unittest.skip('pyamplicol expands p p > j j j into processes without ' - 'amplitudes, and its quark processes from the sm UFO model ' - 'differ from madmatrix (the built-in pyamplicol model agrees)') def test_pp_jjj(self): + """Quark lines and the expansion of the multiparticle labels.""" madmatrix = self.run_mg7('mm', ['p p > j j j'])[:2] amplicol = self.run_mg7('amp', ['p p > j j j [amplicol]'])[:2] self.assert_compatible(amplicol, madmatrix) + def test_gg_ttxg(self): + """Massive quarks in the colour-ordered phase space.""" + madmatrix = self.run_mg7('mm', ['g g > t t~ g'])[:2] + amplicol = self.run_mg7('amp', ['g g > t t~ g [amplicol]'])[:2] + self.assert_compatible(amplicol, madmatrix) + def test_mixed(self): """[amplicol] and madmatrix processes in one output.""" jj = self.run_mg7('jj', ['p p > j j'])[:2] From 2d8f6039365891c130efdc41ccb7675685901f62 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 15 Oct 2026 20:14:20 +0200 Subject: [PATCH 22/22] mg7: generate the [amplicol] artifacts in a worker process Symbolica, imported by pyamplicol, holds a per-user lock until its process exits, so a MadGraph session that had output [amplicol] processes blocked all others. pyamplicol now only runs in amplicol_worker.py, which waits while another process holds the lock, and its workers and cores are left to pyamplicol, licensed through its Symbolica library key. --- madgraph/iolibs/amplicol_worker.py | 81 +++++++++++ madgraph/iolibs/export_amplicol.py | 146 ++++++++++++-------- tests/acceptance_tests/test_mg7_amplicol.py | 11 +- 3 files changed, 176 insertions(+), 62 deletions(-) create mode 100644 madgraph/iolibs/amplicol_worker.py diff --git a/madgraph/iolibs/amplicol_worker.py b/madgraph/iolibs/amplicol_worker.py new file mode 100644 index 0000000000..74001176c6 --- /dev/null +++ b/madgraph/iolibs/amplicol_worker.py @@ -0,0 +1,81 @@ +################################################################################ +# +# Copyright (c) 2026 The MadGraph5_aMC@NLO Development team and Contributors +# +# This file is a part of the MadGraph5_aMC@NLO project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph5_aMC@NLO license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""Generation of the '[amplicol]' artifacts with pyamplicol, run in its own +process by export_amplicol. + +Symbolica, which pyamplicol imports, holds a per-user lock until the process +exits, so it is never imported by MadGraph itself: sessions do not block each +other while they are open. Usage: python amplicol_worker.py JOB.json +""" + +import json +import logging +import sys + + +def model_source(pyamplicol, model): + """The model compiled for pyamplicol, cached in the output.""" + import os + if not os.path.exists(model['cache']): + from pyamplicol.config import EvaluatorConfig, JITConfig + source = pyamplicol.ModelSource.from_path(model['path'], + restriction=model['restriction']) + source.compile(prepared_output=model['cache'], + evaluator=EvaluatorConfig(jit=JITConfig(optimization_level=2))) + return pyamplicol.ModelSource.from_path(model['cache']) + + +def generate_artifact(pyamplicol, artifact, accuracy, process, model): + from pyamplicol.config import resolve_config + # workers and cores are left to pyamplicol, which clamps them without a license + config = resolve_config({ + 'schema_version': 1, + 'action': 'generate', + 'color': {'accuracy': accuracy}, + 'process': {'multiparticles': process['multiparticles'], + 'coupling_order_policy': 'explicit', + 'max_coupling_orders': process['orders']}, + 'generation': {'emit_api_bundle': True, + 'relation_discovery': {'mode': 'off'}}, + 'evaluator': {'backend': 'jit', 'execution_mode': 'recurrence', + 'jit': {'optimization_level': 2}}, + }) + pyamplicol.Generator(config).generate(process['expression'], artifact, + model=model, mode='replace') + + +def main(job_path): + import os + logging.basicConfig(level=logging.WARNING, format='%(name)s: %(message)s') + with open(job_path) as stream: + job = json.load(stream) + import pyamplicol + from pyamplicol.artifacts import umami_reweight + model = model_source(pyamplicol, job['model']) + providers = [] + for process in job['processes']: + path = process['path'] + for accuracy in ('lc', 'full'): + generate_artifact(pyamplicol, os.path.join(path, accuracy), accuracy, + process, model) + providers.append(umami_reweight.write_reweight_provider( + os.path.join(path, 'lc'), os.path.join(path, 'full'), + os.path.join(path, 'reweight'))) + with open(job['result'], 'w') as stream: + json.dump(providers, stream) + + +if __name__ == '__main__': + main(sys.argv[1]) diff --git a/madgraph/iolibs/export_amplicol.py b/madgraph/iolibs/export_amplicol.py index f8b919587a..c278c40453 100644 --- a/madgraph/iolibs/export_amplicol.py +++ b/madgraph/iolibs/export_amplicol.py @@ -22,11 +22,16 @@ on its own with a colour-ordered phase-space mapping. """ +import importlib.util import json import logging import os import re import shutil +import subprocess +import sys +import tempfile +import time import madgraph @@ -49,15 +54,61 @@ """ -def import_pyamplicol(): - try: - import pyamplicol - from pyamplicol.artifacts import umami_reweight - except ImportError as error: +WORKER = pjoin(os.path.dirname(os.path.abspath(__file__)), 'amplicol_worker.py') +# printed by Symbolica when another process holds its per-user thread +SYMBOLICA_BUSY = 'Cannot start another restricted Symbolica thread' + + +def check_pyamplicol(): + """pyamplicol is only imported by the worker, see amplicol_worker.""" + if importlib.util.find_spec('pyamplicol') is None: raise madgraph.InvalidCmd( - 'The [amplicol] processes need pyamplicol (%s). Install it in the ' - 'Python environment of MadGraph.' % error) - return pyamplicol, umami_reweight + 'The [amplicol] processes need pyamplicol. Install it in the Python ' + 'environment of MadGraph.') + + +def run_worker(job): + """Run amplicol_worker on a job, waiting while Symbolica is busy elsewhere.""" + env = dict(os.environ, SYMBOLICA_HIDE_BANNER='1', + PYTHONPATH=os.pathsep.join(p for p in sys.path if p)) + waiting = False + while True: + worker = subprocess.Popen([sys.executable, WORKER, job], env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True) + busy = False + for line in worker.stdout: + if SYMBOLICA_BUSY in line: + busy = True + elif not line.startswith(('\u250c', '\u2502', '\u2514', '\u251c')): + sys.stdout.write(line) + if worker.wait() == 0: + return + if not busy: + raise madgraph.InvalidCmd('pyamplicol failed to generate the [amplicol] ' + 'matrix elements, see the messages above') + if 'symbolica' in sys.modules: + raise madgraph.InvalidCmd( + 'Symbolica is imported in this MadGraph process, so the pyamplicol ' + 'worker cannot start: do not import pyamplicol or symbolica here') + if not waiting: + logger.info('Symbolica is used by another process: waiting for it') + waiting = True + time.sleep(10) + + +def load_metadata(artifact): + """The UMAMI metadata documents of an artifact, one per provider.""" + root = pjoin(artifact, 'API', 'umami') + with open(pjoin(root, 'metadata.json')) as stream: + data = json.load(stream) + if data.get('kind') != 'pyamplicol-umami-provider-index': + return [data] + documents = [] + for provider in data['providers']: + with open(pjoin(root, provider['metadata'])) as stream: + documents.append(json.load(stream)) + return documents def leg_pdgs(leg, model): @@ -130,11 +181,9 @@ def coupling_orders(procdef, model): for order, weight in hierarchy.items() if weight > lowest} -def model_source(pyamplicol, model, cache_path): +def model_job(model, cache_path): """The UFO model of the output with its restriction, compiled for pyamplicol once per output.""" - if os.path.exists(cache_path): - return pyamplicol.ModelSource.from_path(cache_path) path = model.get('modelpath') reference = os.path.basename(model.get('modelpath+restriction')) restriction = 'default' @@ -144,33 +193,9 @@ def model_source(pyamplicol, model, cache_path): restriction = 'none' else: restriction = pjoin(path, 'restrict_%s.dat' % suffix) - from pyamplicol.config import EvaluatorConfig, JITConfig - source = pyamplicol.ModelSource.from_path(path, restriction=restriction) - logger.info('Compiling the model %s for AmpliCol' % reference) - source.compile(prepared_output=cache_path, - evaluator=EvaluatorConfig(jit=JITConfig(optimization_level=2))) - return pyamplicol.ModelSource.from_path(cache_path) - - -def generate_artifact(pyamplicol, artifact, accuracy, expression, multiparticles, - orders, model): - from pyamplicol.config import resolve_config - # one worker: Symbolica's restricted mode allows a single thread - config = resolve_config({ - 'schema_version': 1, - 'action': 'generate', - 'color': {'accuracy': accuracy}, - 'process': {'multiparticles': multiparticles, - 'coupling_order_policy': 'explicit', - 'max_coupling_orders': orders}, - 'generation': {'workers': 1, 'emit_api_bundle': True, - 'relation_discovery': {'mode': 'off'}}, - 'evaluator': {'backend': 'jit', 'execution_mode': 'recurrence', - 'optimization': {'cores': 1}, - 'jit': {'optimization_level': 2}}, - }) - pyamplicol.Generator(config).generate(expression, artifact, model=model, - mode='replace') + if not os.path.exists(cache_path): + logger.info('Compiling the model %s for AmpliCol' % reference) + return {'path': path, 'restriction': restriction, 'cache': cache_path} def representative_processes(proc_defs, model): @@ -283,31 +308,40 @@ def parameter_references(metadata): def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None): """Write SubProcesses/A_ for every '[amplicol]' process definition and return the subprocesses.json entries.""" - pyamplicol, umami_reweight = import_pyamplicol() + check_pyamplicol() subproc_dir = pjoin(dir_path, 'SubProcesses') - source = model_source(pyamplicol, model, pjoin(subproc_dir, MODEL_FILE)) rusticol_config = shutil.which('rusticol-config') or 'rusticol-config' - entries = [] + processes = [] for number, procdef in enumerate(proc_defs, start=1): name = 'A%d_%s' % (number, process_name(procdef, model, multiparticles or {})) path = pjoin(subproc_dir, name) + expression, process_multiparticles = process_expression(procdef, model) + os.makedirs(path, exist_ok=True) + logger.info('Generating the AmpliCol matrix elements of %s in %s' % + (procdef.nice_string(prefix=False), path)) + processes.append({'name': name, 'path': path, 'expression': expression, + 'multiparticles': process_multiparticles, + 'orders': coupling_orders(procdef, model)}) + with tempfile.TemporaryDirectory(prefix='amplicol_job_') as scratch: + job = pjoin(scratch, 'job.json') + with open(job, 'w') as stream: + json.dump({'model': model_job(model, pjoin(subproc_dir, MODEL_FILE)), + 'processes': processes, + 'result': pjoin(scratch, 'providers.json')}, stream) + run_worker(job) + with open(pjoin(scratch, 'providers.json')) as stream: + providers = json.load(stream) + entries = [] + for procdef, process, provider_ids in zip(proc_defs, processes, providers): + name, path, orders = process['name'], process['path'], process['orders'] rel_dir = pjoin('SubProcesses', name) - expression, multiparticles = process_expression(procdef, model) - orders = coupling_orders(procdef, model) # every vertex carries one coupling: the rest are powers of g_s qcd_power = (len(procdef.get('legs')) - 2 - sum(orders.values()) if len(orders) == 1 else -1) - logger.info('Generating the AmpliCol matrix elements of %s in %s' % - (procdef.nice_string(prefix=False), path)) - os.makedirs(path, exist_ok=True) - for accuracy in ('lc', 'full'): - generate_artifact(pyamplicol, pjoin(path, accuracy), accuracy, - expression, multiparticles, orders, source) - provider_ids = umami_reweight.write_reweight_provider( - pjoin(path, 'lc'), pjoin(path, 'full'), pjoin(path, 'reweight')) - documents = umami_reweight.load_umami_metadata(pjoin(path, 'lc')) + initial_legs = [set(leg_pdgs(leg, model)) + for leg in procdef.get('legs') if not leg.get('state')] copies = [] - for provider_id, metadata in zip(provider_ids, documents): + for provider_id, metadata in zip(provider_ids, load_metadata(pjoin(path, 'lc'))): library = metadata['provider']['library'] me_path = pjoin('lib', 'libamplicol_%s_%s.so' % (name, provider_id)) ratio_path = pjoin('lib', 'libamplicol_%s_%s_ratio.so' % @@ -315,8 +349,6 @@ def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None): copies.append('\tcp $(BUILD)/lc/%s ../../%s' % (library, me_path)) copies.append('\tcp $(BUILD)/reweight/libreweight_%s.so ../../%s' % (provider_id, ratio_path)) - initial_legs = [set(leg_pdgs(leg, model)) - for leg in procdef.get('legs') if not leg.get('state')] entries.extend(subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power, initial_legs)) with open(pjoin(path, 'makefile'), 'w') as stream: @@ -324,7 +356,7 @@ def export_amplicol_processes(dir_path, proc_defs, model, multiparticles=None): 'copies': '\n'.join(copies)}) with open(pjoin(path, 'process.json'), 'w') as stream: json.dump({'process': procdef.nice_string(prefix=False), - 'expression': expression, - 'multiparticles': multiparticles, + 'expression': process['expression'], + 'multiparticles': process['multiparticles'], 'max_coupling_orders': orders}, stream, indent=2) return entries diff --git a/tests/acceptance_tests/test_mg7_amplicol.py b/tests/acceptance_tests/test_mg7_amplicol.py index 880de0ab8c..802aa2ca72 100644 --- a/tests/acceptance_tests/test_mg7_amplicol.py +++ b/tests/acceptance_tests/test_mg7_amplicol.py @@ -26,6 +26,7 @@ import glob import gzip +import importlib.util import json import os import re @@ -87,11 +88,11 @@ def _color_balanced(lhe_path): class MG7AmplicolTest(unittest.TestCase): def setUp(self): - try: - import madspace # noqa: F401 - import pyamplicol # noqa: F401 - except ImportError as error: - self.skipTest('madspace or pyamplicol unavailable: %s' % error) + # not imported: Symbolica would hold its lock against the pyamplicol worker + missing = [name for name in ('madspace', 'pyamplicol') + if importlib.util.find_spec(name) is None] + if missing: + self.skipTest('%s unavailable' % ' and '.join(missing)) if not misc.resolve_lhapdf().find_set(_PDF): self.skipTest('%s LHAPDF data not found' % _PDF) self.path = tempfile.mkdtemp(prefix='mg7_amplicol_')