diff --git a/.github/workflows/docs_examples.yml b/.github/workflows/docs_examples.yml index 3259fe375b..5afbaa3343 100644 --- a/.github/workflows/docs_examples.yml +++ b/.github/workflows/docs_examples.yml @@ -3,13 +3,17 @@ name: MadSpace docs examples # The other MG7 acceptance workflows all carry `paths-ignore: ['docs/**']`, so a # docs-only commit currently triggers none of them. This workflow does the # opposite: it triggers ONLY on the paths that can affect the documentation -# examples, so editing docs/source/madspace/examples/*.rst is actually tested. +# examples, so editing docs/source/madspace/examples/*.rst (or the event +# streaming example in docs/source/event_stream.rst) is actually tested. on: push: branches: - '**' paths: - 'docs/source/madspace/examples/**' + - 'docs/source/event_stream.rst' + - 'madgraph/iolibs/template_files/mg7/**' + - 'madgraph/mg7/**' - 'madspace/**' - 'tests/acceptance_tests/test_madspace_examples.py' - '.github/workflows/docs_examples.yml' @@ -18,6 +22,9 @@ on: types: [opened, synchronize, reopened] paths: - 'docs/source/madspace/examples/**' + - 'docs/source/event_stream.rst' + - 'madgraph/iolibs/template_files/mg7/**' + - 'madgraph/mg7/**' - 'madspace/**' - 'tests/acceptance_tests/test_madspace_examples.py' - '.github/workflows/docs_examples.yml' @@ -55,4 +62,4 @@ jobs: - name: run the madspace doc examples run: | cd $GITHUB_WORKSPACE - ./tests/test_manager.py test_madspace_example_simple_mappings test_madspace_example_cuts test_madspace_example_diagram_mapping test_madspace_example_integration_order test_madspace_example_matrix_element test_madspace_example_pdf test_madspace_example_integrator test_madspace_example_integrator_madnis test_madspace_example_flow_training -pA -t0 -l INFO + ./tests/test_manager.py test_madspace_example_simple_mappings test_madspace_example_cuts test_madspace_example_diagram_mapping test_madspace_example_integration_order test_madspace_example_matrix_element test_madspace_example_pdf test_madspace_example_integrator test_madspace_example_integrator_madnis test_madspace_example_flow_training test_gridpack_example_event_stream -pA -t0 -l INFO diff --git a/.github/workflows/pyamplicol.yml b/.github/workflows/pyamplicol.yml index b847a4b4e9..12e82d7d1b 100644 --- a/.github/workflows/pyamplicol.yml +++ b/.github/workflows/pyamplicol.yml @@ -73,7 +73,7 @@ jobs: if [[ -z "$base" || "$base" =~ ^0+$ ]] || ! git cat-file -e "$base^{commit}" 2>/dev/null; then changed=true elif git diff --name-only "$base" HEAD \ - | grep -qE '^(pyamplicol/|\.github/workflows/pyamplicol|\.github/actions/(pyamplicol|build_pyamplicol|install_pyamplicol)|VERSION$)'; then + | grep -qE '^(pyamplicol/|\.github/workflows/pyamplicol|\.github/actions/(pyamplicol|build_pyamplicol|install_pyamplicol)|madspace/include/madspace/umami\.h$|madgraph/mg7/bootstrap\.py$|VERSION$)'; then changed=true else changed=false @@ -173,7 +173,7 @@ jobs: python - <<'EOF' import sys sys.path.insert(0, ".") - from madgraph.iolibs.template_files.mg7 import bootstrap + from madgraph.mg7 import bootstrap bootstrap.ensure_pyamplicol(interactive=False) bootstrap.check_pyamplicol_license() assert bootstrap.check_pyamplicol_source_hash() diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 25160316e2..ad98defd71 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -227,12 +227,16 @@ jobs: done test -d madspace/install/madspace - # Also checks that Symbolica accepts pyamplicol's library key. + # Also checks that Symbolica accepts pyamplicol's library key; only the + # wheel download is retried, so a rejected key fails the step. for attempt in $(seq 10); do - python3 pyamplicol/install.py --bin -y && break + if python3 pyamplicol/install.py --bin -y; then status=0; break; else status=$?; fi + # The wheel installed, so the license check failed: retrying cannot help. + test -d pyamplicol/install/pyamplicol && break echo "attempt $attempt: pyamplicol==${{ needs.check_version.outputs.version }} not installable yet, retrying..." sleep 15 done + test "$status" -eq 0 test -d pyamplicol/install/pyamplicol - name: Sign the tarball with Sigstore diff --git a/docs/source/conf.py b/docs/source/conf.py index e2963634d0..80c12fbfb2 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -14,6 +14,17 @@ if os.path.isdir(os.path.join(_local_install, "madspace")): sys.path.insert(0, _local_install) +# bin/event_stream.py of the gridpacks, for the API reference in event_stream.rst +sys.path.insert( + 0, + os.path.abspath( + os.path.join( + os.path.dirname(__file__), + "..", "..", "madgraph", "iolibs", "template_files", "mg7", + ) + ), +) + # Configuration file for the Sphinx documentation builder. # # For the full list of built-in configuration values, see the documentation: diff --git a/docs/source/custom_functions.rst b/docs/source/custom_functions.rst index a8eb54166e..78bd6f7a14 100644 --- a/docs/source/custom_functions.rst +++ b/docs/source/custom_functions.rst @@ -290,6 +290,27 @@ matrix products. If your function relies on it, set ``OMP_NUM_THREADS=1``. The Python API behind these functions is ``madspace.PythonFunction``. It can also be used directly with madspace. +.. _custom-function-surrogate: + +Python surrogates +----------------- + +A Python surrogate of the matrix element (see :ref:`run-card-surrogate`) uses the same +mechanism. Its file defines ``FRAMEWORK`` and ``custom_surrogate(momenta, pdg_ids, +alpha_s)``, which returns :math:`|M|^2` as a float64 array of shape ``(batch,)``. It takes +no parameters. ``momenta`` and ``pdg_ids`` have the shapes and types given above, with the +momenta that the full matrix element gets, and ``alpha_s`` is a float64 array of shape +``(batch,)``. The function is called for all subprocesses, which ``pdg_ids`` identifies: + +.. code-block:: python + + import numpy as np + + FRAMEWORK = "numpy" + + def custom_surrogate(momenta, pdg_ids, alpha_s): + return np.ones(len(momenta)) + Examples -------- diff --git a/docs/source/event_stream.rst b/docs/source/event_stream.rst new file mode 100644 index 0000000000..c3ce6391c8 --- /dev/null +++ b/docs/source/event_stream.rst @@ -0,0 +1,159 @@ +Streaming events from Python +============================ + +A gridpack can generate events directly into a Python program, without writing any files. +This is useful when the events are consumed right away, for example by an analysis written +in Python or as training data for a neural network. The ``EventStream`` class in the +gridpack's ``bin/event_stream.py`` returns the events either as LHE event objects or as +batches of NumPy arrays or PyTorch tensors. + +Event streaming needs a gridpack created with ``run_mode = "fix_max_weight"`` (see +:ref:`gridpack-run-mode`). In this mode, the maximum weight of every channel is fixed when +the gridpack is created. An event is final as soon as it is unweighted, so it can be returned +right away without a survey or a combination step. + +Example: the e+ e- invariant mass +--------------------------------- + +This example fills a histogram of the invariant mass of the electron pair in +``p p > e+ e-``. First, create the gridpack at the MadGraph7 prompt: + +.. code-block:: text + :name: event-stream-mg5-commands + + generate p p > e+ e- + output mg7 PROC_ppee + launch PROC_ppee + set save_gridpack true + set run_mode fix_max_weight + done + +The gridpack is written to ``PROC_ppee/Events/run_01/gridpack``. Its ``bin`` directory +contains ``event_stream.py``, so add it to the Python path and import ``EventStream``. The +script can run from any directory. + +.. code-block:: python + + import sys + + import numpy as np + + sys.path.insert(0, "PROC_ppee/Events/run_01/gridpack/bin") + from event_stream import EventStream + +Creating an ``EventStream`` starts the event generation in background threads. Its +arguments, like the device, the seed and the batch sizes, default to the values in the +gridpack's ``Cards/grid_run_card.toml``. The stream is endless, so ``max_events`` limits +the number of events. Each batch is an :py:class:`EventBatch ` +with one array per LHE field. The particle fields have the shape +``(batch_size, max_particle_count)``. + +.. code-block:: python + + bins = np.linspace(60.0, 120.0, 31) + hist = np.zeros(len(bins) - 1) + weight_sum = 0.0 + + with EventStream(seed=42) as stream: + for batch in stream.batches(10000, max_events=100000): + # four-momentum of the e+ e- pair: sum over the outgoing electrons + is_electron = (np.abs(batch.pdg_id) == 11) & (batch.status_code == 1) + e, px, py, pz = ( + np.sum(np.where(is_electron, p, 0.0), axis=1) + for p in (batch.energy, batch.px, batch.py, batch.pz) + ) + m_ee = np.sqrt(np.maximum(e**2 - px**2 - py**2 - pz**2, 0.0)) + hist += np.histogram(m_ee, bins, weights=batch.weight)[0] + weight_sum += np.sum(batch.weight) + cross_section = stream.integral + print(f"cross section: {cross_section:.1f} +- {stream.integral_error:.1f} pb") + +The event weights are not normalized to the cross section (see :ref:`event-stream-weights`). +Scaling them to the cross section gives the differential cross section in pb/GeV: + +.. code-block:: python + + dsigma = hist * cross_section / weight_sum / np.diff(bins) + for low, value in zip(bins[:-1], dsigma): + print(f"{low:5.0f} GeV {value:8.3f} pb/GeV " + "#" * int(50 * value / dsigma.max())) + + peak = np.argmax(dsigma) + assert bins[peak] < 91.2 < bins[peak + 1], "the Z peak is missing" + +The output shows the Z peak (shortened):: + + cross section: 880.2 +- 0.6 pb + ... + 84 GeV 7.569 pb/GeV ## + 86 GeV 16.543 pb/GeV ##### + 88 GeV 51.287 pb/GeV ################ + 90 GeV 158.376 pb/GeV ################################################## + 92 GeV 66.743 pb/GeV ##################### + 94 GeV 18.457 pb/GeV ##### + ... + +Output formats +-------------- + +There are three ways to get events from a stream: + +- Iterating over the stream yields one ``madspace.LHEEvent`` at a time. The events are + completed in chunks of 1000 internally. If the loop is left early, the rest of the chunk + is returned first by the next call, so no event is skipped. +- :py:meth:`next_events ` returns a list of a given + number of ``LHEEvent`` objects. +- :py:meth:`batches ` and + :py:meth:`next_batch ` return + :py:class:`EventBatch ` objects. With ``output_type="torch"``, + their fields are PyTorch tensors on the CPU instead of NumPy arrays. + +Events have different numbers of particles if the process has several subprocesses or +includes decays. The particle fields are then padded with zeros. ``particle_count`` holds +the number of particles of every event and ``particle_mask`` marks the valid entries. + +.. _event-stream-weights: + +Weights and cross section +------------------------- + +After the unweighting, every event weight is ``1``, or ``-1`` for an event with a negative +weight. The exception are overweight events, whose weight exceeded the fixed maximum weight +of their channel. They keep the ratio of their weight to the maximum weight, so their weight +is larger than ``1``. Treating them as ``1`` would bias the distributions in the phase-space +regions where they occur. + +The events of ``[amplicol]`` subprocesses are then reweighted to full colour in the default +``mode = "reweight"``: their weight is multiplied by the ratio of the full-colour to the +leading-colour matrix element, so it is not ``1``. With ``mode = "unweight"``, they are +unweighted a second time and have unit weights apart from overweights. + +In all cases, to normalize a sample to the cross section, scale its weights by +``stream.integral / sum(weights)``. The cross section estimate improves while events are +generated: it combines the integrals stored in the gridpack with the weighted events +generated by the stream so far. Read it after generating the sample, as in the example +above. + +How it works +------------ + +Each channel of the gridpack keeps a buffer of unweighted events, which background threads +refill whenever it runs low. The channel of every event is drawn at random, with +probabilities given by the current cross section estimates of the channels. Generating +the phase-space points and evaluating the matrix elements happen in separate steps, as in a +normal gridpack run. + +The events are reproducible: for a given seed, the stream always returns the same sequence of +events. It does not depend on the number of threads or on how the events are split into +batches. + +API reference +------------- + +.. autoclass:: event_stream.EventStream + :members: + :special-members: __iter__ + :member-order: bysource + +.. autoclass:: event_stream.EventBatch + :members: + :member-order: bysource diff --git a/docs/source/gridpacks.rst b/docs/source/gridpacks.rst index 747a746cda..315d33ee4b 100644 --- a/docs/source/gridpacks.rst +++ b/docs/source/gridpacks.rst @@ -70,6 +70,7 @@ ends: longer, especially for channels with a low unweighting efficiency, but the gridpack runs do not have to estimate the maximum weights first. This helps if each run only generates a few events. + The gridpack also gets ``bin/event_stream.py``, see :ref:`gridpack-event-stream`. ``regular`` The events of the run are generated as usual, then the gridpack is stored. @@ -96,6 +97,15 @@ Make sure every job uses a different seed. The default seed ``-1`` draws a fresh seed for each call, which is fine for independent jobs. The seed that was used is recorded in the output. +.. _gridpack-event-stream: + +Streaming events from Python +---------------------------- + +A gridpack created with ``run_mode = "fix_max_weight"`` also contains +``bin/event_stream.py``. Its ``EventStream`` class generates events in memory and returns +them to a Python program, without writing any files. See :doc:`event_stream`. + Run card settings ----------------- @@ -147,6 +157,9 @@ gridpack is created. * - - ``cut_efficiency_threshold``, ``max_cut_repetitions`` - Card only. + * - ``[amplicol]`` + - ``first_pass_max_overweight_truncation``, ``first_pass_freeze_max_weight_after`` + - ``--first_pass_max_overweight_truncation``, ``--first_pass_freeze_max_weight_after`` * - ``[gridpack]`` - ``output_dir``, ``temp_output_dir`` - ``--output_dir``, ``--temp_output_dir`` @@ -160,6 +173,14 @@ the main run card, see :doc:`run_card`. ``./bin/generate_events --help`` lists a Relative ``output_dir`` and ``temp_output_dir`` paths in the card refer to the gridpack directory. Both are reset to ``""`` when the gridpack is created. +With ``[amplicol]`` processes unweighted to full colour, ``max_overweight_truncation`` and +``freeze_max_weight_after`` set the second pass, which gives the events, and the +``first_pass_`` options the leading-colour unweighting of these processes, as in the main +run. The other processes are unweighted once, as set by ``[generation]``. The colour +accuracy and the mode are fixed when the gridpack is created. The pyamplicol artifacts +the libraries read are copied to ``SubProcesses/``, and ``include_source = true`` is not +supported for these processes. + The card names the ``cpu_mode`` that the gridpack was built with, which can differ from the ``"auto"`` of the original run card. The matrix element libraries carry this name. If you change ``cpu_mode`` or ``precision``, you have to rebuild them. diff --git a/docs/source/index.rst b/docs/source/index.rst index d891915be8..acc499b537 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -30,6 +30,7 @@ If you need a stable release, use the custom_functions diagram_filters gridpacks + event_stream madevent_to_mg7 .. toctree:: diff --git a/docs/source/madspace/cpp-api.rst b/docs/source/madspace/cpp-api.rst index 29789c1b58..e4c49d0f6d 100644 --- a/docs/source/madspace/cpp-api.rst +++ b/docs/source/madspace/cpp-api.rst @@ -70,6 +70,7 @@ Builders that assemble compute-graph functions for channel weights, energy scale cpp/SubchannelWeights cpp/Unweighter cpp/VegasHistogram + cpp/WeightingStepFunction Phase-space building blocks --------------------------- @@ -85,6 +86,7 @@ Supporting types shared by the mappings: topologies, cuts, observables and densi cpp/PdfGrid cpp/Propagator cpp/Topology + cpp/WeightingStep cpp/initialize_vegas_grid Compute graph @@ -127,6 +129,7 @@ The high-level driver plus the tensor and device layer it runs on: integration, cpp/EventHistogramSpec cpp/EventHistograms cpp/EventRecord + cpp/EventStream cpp/FieldLayout cpp/GeneratorBatchJob cpp/GeneratorConfig @@ -157,6 +160,10 @@ The high-level driver plus the tensor and device layer it runs on: integration, cpp/Runtime cpp/Sizes cpp/StatusFile + cpp/StepChainResult + cpp/StepEvents + cpp/StepIntegral + cpp/StepSettings cpp/SubprocessObservables cpp/SubprocessSystArgs cpp/SystematicsCalculator @@ -171,6 +178,8 @@ The high-level driver plus the tensor and device layer it runs on: integration, cpp/UnalignedRef cpp/Variation cpp/VegasGridOptimizer + cpp/WeightingGroup + cpp/WeightingPipeline cpp/available_backends cpp/compute_generation_batch_event_count cpp/cpu_device @@ -186,6 +195,8 @@ The high-level driver plus the tensor and device layer it runs on: integration, cpp/set_simd_mode cpp/supported_simd_modes cpp/set_lib_path + cpp/weighting_groups_from_json + cpp/weighting_groups_to_json Utilities --------- diff --git a/docs/source/madspace/python-api.rst b/docs/source/madspace/python-api.rst index 9744f3f539..a5ba7b7ea5 100644 --- a/docs/source/madspace/python-api.rst +++ b/docs/source/madspace/python-api.rst @@ -86,6 +86,7 @@ Supporting types shared by the mappings: topologies, cuts, observables and densi python/PdfGrid python/Propagator python/Topology + python/WeightingStep python/initialize_vegas_grid Compute graph @@ -125,6 +126,7 @@ The high-level driver plus the tensor and device layer it runs on: integration, python/EventGenerator python/EventHistogramSpec python/EventHistograms + python/EventStream python/GeneratorConfig python/GeneratorStatus python/Histogram @@ -155,6 +157,7 @@ The high-level driver plus the tensor and device layer it runs on: integration, python/Variation python/VegasGridOptimizer python/Verbosity + python/WeightingGroup python/available_backends python/compute_generation_batch_event_count python/cpu_device @@ -170,6 +173,8 @@ The high-level driver plus the tensor and device layer it runs on: integration, python/set_simd_mode python/supported_simd_modes python/set_lib_path + python/weighting_groups_from_json + python/weighting_groups_to_json Utilities --------- diff --git a/docs/source/run_card.rst b/docs/source/run_card.rst index b310007ac8..af7d8ba1c4 100644 --- a/docs/source/run_card.rst +++ b/docs/source/run_card.rst @@ -129,7 +129,8 @@ See :doc:`gridpacks` for details. - ``fix_max_weight``: no events. Each channel runs until ``freeze_max_weight_after`` unweighted events are reached. Its maximum weight is then fixed in the gridpack, and its integral is used to share the events between the channels. This can be - slow for channels with a low unweighting efficiency. + slow for channels with a low unweighting efficiency. The gridpack also gets + ``bin/event_stream.py`` to stream events in memory from Python. See :ref:`gridpack-run-mode`. @@ -369,6 +370,52 @@ An empty string disables the function. See :doc:`custom_functions` for details. A parameter of a custom function, for example ``cut_parameters.pt_min = 30.0``. It replaces the default value declared by the function. See :ref:`custom-function-parameters`. +.. _run-card-surrogate: + +[surrogate] +----------- + +A fast surrogate of the matrix element, for example a neural network. The integrand and +its optimization use the surrogate. The events are then reweighted or unweighted a second +time with the full matrix element, which is only evaluated for the events that pass the +first unweighting. The cross section and the events are those of the full matrix element. +The surrogate does not return channel weights, so it needs ``sde_strategy = +"denominators"``. It is not supported for outputs with ``[amplicol]`` subprocesses. + +``function`` (``""``) + The surrogate. An empty string disables it. + + - A ``.so`` library implementing the UMAMI interface. It gets the momenta, + :math:`\alpha_s` and the flavor index of the madmatrix libraries as inputs, and + returns the matrix element. It is loaded once per subprocess: ``{subprocess}`` and + ``{device}`` in the path are replaced by the subprocess name (e.g. ``P1_gg_ttx``) + and the device, like in ``lib/libmadmatrix_{subprocess}_{device}.so``. It has to + be built for the device of the run. If it implements ``umami_set_parameter``, it + gets the run-time parameters of the madmatrix libraries, such as ``bwcutoff``. + - A Python file defining ``custom_surrogate(momenta, pdg_ids, alpha_s)``, which returns + :math:`|M|^2` for a batch of events. ``momenta`` has the shape ``(batch, n, 4)``, + ``pdg_ids`` the shape ``(batch, n)`` and identifies subprocess and flavor, and + ``alpha_s`` the shape ``(batch,)``. ``FRAMEWORK = "torch"`` passes torch tensors, + which also run on GPUs, otherwise numpy arrays are passed. See + :ref:`custom-function-surrogate`. + +``mode`` (``"unweight"``) + ``"unweight"`` unweights the events a second time with respect to the ratio of the full + matrix element to the surrogate. The events then have unit weights apart from + overweights. ``"reweight"`` keeps the ratio in the event weights. + +``first_pass_max_overweight_truncation`` (``0.01``), ``first_pass_freeze_max_weight_after`` (``10000``) + ``max_overweight_truncation`` and ``freeze_max_weight_after`` of the first + unweighting, with the surrogate, in unweight mode. The ones of ``[generation]`` set + the second unweighting with the full matrix element, which gives the final events, + and the unweighting of the subprocesses without a surrogate. In reweight mode, the + events are unweighted once, with the surrogate, as set by ``[generation]``. + +``param_card`` (``""``) + The param card a ``.so`` surrogate is initialized with. An empty string uses + ``Cards/param_card.dat``. A madmatrix library with a different param card, for example, + is a surrogate for testing. + .. _run-card-multiparticles: [multiparticles] diff --git a/madgraph/interface/madgraph_interface.py b/madgraph/interface/madgraph_interface.py index 901b77434a..f959108ca5 100755 --- a/madgraph/interface/madgraph_interface.py +++ b/madgraph/interface/madgraph_interface.py @@ -180,6 +180,27 @@ 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 = next((m for m in _BACKEND_BRACKET.finditer(line) + if m.group('backend').lower() in ME_BACKENDS), None) + if not match: + return line, '' + rest = line[:match.start()] + ' ' + line[match.end():] + if '[' in rest or ']' in rest: + 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 @@ -2410,7 +2431,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) @@ -3966,6 +3991,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 @@ -4054,6 +4081,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', @@ -4099,6 +4127,30 @@ def extract_diagram_filter(self, args, option='--diagram_filter'): # 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 @@ -4197,7 +4249,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) @@ -6138,6 +6193,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 = "" @@ -7660,6 +7716,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() @@ -9734,11 +9791,11 @@ def do_launch(self, line): # what needs madspace). The bootstrap has to be told whether it may # take over the terminal: in a subprocess it could read that off # sys.argv/sys.stdin, but in process those describe MG5, not the run. - from madgraph.iolibs.template_files.mg7 import bootstrap as mg7_bootstrap + from madgraph.mg7 import bootstrap as mg7_bootstrap mg7_bootstrap.ensure_madspace( interactive=bool(self.use_rawinput) and not options['force'], jobs=self.options['nb_core']) - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch MG7 = mg7_launch.MG7Cmd(me_dir=me_dir, options=self.options) if options['interactive']: @@ -11128,6 +11185,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' @@ -11303,6 +11361,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 @@ -11786,6 +11845,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 @@ -12682,7 +12742,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())] @@ -13078,7 +13138,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': @@ -13709,7 +13769,7 @@ def integrate_widths_madspace(self, opts, decay_info): after others succeeded, the caller still has the decay_info it would hand to MadEvent. """ - from madgraph.iolibs.template_files.mg7 import bootstrap as mg7_bootstrap + from madgraph.mg7 import bootstrap as mg7_bootstrap # the survey of a decay uses no PDF: only the build needs the environment self.setup_mg7_environment(lhapdf=False) mg7_bootstrap.ensure_madspace(interactive=False, @@ -14167,6 +14227,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 57c4458924..4daa6efe6b 100755 --- a/madgraph/interface/master_interface.py +++ b/madgraph/interface/master_interface.py @@ -178,6 +178,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/madgraph/interface/tutorials/lo.py b/madgraph/interface/tutorials/lo.py index 2514029a7d..300e5d1251 100644 --- a/madgraph/interface/tutorials/lo.py +++ b/madgraph/interface/tutorials/lo.py @@ -45,11 +45,11 @@ def _mg7_bootstrap(): Everything the tutorial wants to know about madspace goes through it, so the tutorial cannot disagree with the launcher about where madspace is or - whether it is there (see iolibs/template_files/mg7/bootstrap.py). + whether it is there (see madgraph/mg7/bootstrap.py). """ try: - from madgraph.iolibs.template_files.mg7 import bootstrap + from madgraph.mg7 import bootstrap except Exception: return None return bootstrap @@ -233,7 +233,7 @@ def _systematics_rows(info): if not info: return illustrative try: - from madgraph.iolibs.template_files.mg7 import systematics_summary + from madgraph.mg7 import systematics_summary except Exception: return illustrative diff --git a/madgraph/interface/tutorials/mg7.py b/madgraph/interface/tutorials/mg7.py index 5874066426..dc66d78cea 100644 --- a/madgraph/interface/tutorials/mg7.py +++ b/madgraph/interface/tutorials/mg7.py @@ -168,7 +168,8 @@ documentation maps the old names onto the new sections. **Gridpacks** work here too, via the `[gridpack]` section. The "Gridpacks" -page of the documentation describes them. +page of the documentation describes them, including `bin/event_stream.py`, a +Python API that streams events in memory. %(see_also)s diff --git a/madgraph/iolibs/amplicol_worker.py b/madgraph/iolibs/amplicol_worker.py new file mode 100644 index 0000000000..ee3237a7c9 --- /dev/null +++ b/madgraph/iolibs/amplicol_worker.py @@ -0,0 +1,103 @@ +################################################################################ +# +# 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 and compiled + again unless the key of the cache is model['key'].""" + import os + import shutil + try: + with open(model['key_file']) as stream: + cached = json.load(stream).get('key') + except (OSError, ValueError, AttributeError): + cached = None + if cached != model['key'] or not os.path.exists(model['cache']): + from pyamplicol.config import EvaluatorConfig, JITConfig + if os.path.exists(model['key_file']): + os.remove(model['key_file']) + if os.path.isdir(model['cache']): + shutil.rmtree(model['cache']) + elif os.path.exists(model['cache']): + os.remove(model['cache']) + source = pyamplicol.ModelSource.from_path(model['path'], + restriction=model['restriction']) + source.compile(prepared_output=model['cache'], + evaluator=EvaluatorConfig(jit=JITConfig(optimization_level=2))) + with open(model['key_file'], 'w') as stream: + json.dump({'key': model['key'], 'sources': model['sources']}, stream, + indent=2) + return pyamplicol.ModelSource.from_path(model['cache']) + + +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'], + 'flavor_blind_strict': True, + 'coupling_order_policy': 'explicit', + 'max_coupling_orders': process['orders']}, + # pack_payloads: one payloads.pacbin instead of files per process + 'generation': {'emit_api_bundle': True, 'pack_payloads': True, + 'relation_discovery': {'mode': 'off'}}, + 'evaluator': {'backend': 'jit', 'execution_mode': 'recurrence', + 'jit': {'optimization_level': 2}}, + }) + 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 + from pyamplicol.licensing import detect_symbolica_license + if not detect_symbolica_license(suggest=False).licensed: + logging.warning('Symbolica does not accept the license key of pyamplicol ' + '(%s): generating in its restricted mode', pyamplicol.__file__) + model = model_source(pyamplicol, job['model']) + providers = [] + for process in job['processes']: + 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 new file mode 100644 index 0000000000..1e091e3877 --- /dev/null +++ b/madgraph/iolibs/export_amplicol.py @@ -0,0 +1,489 @@ +################################################################################ +# +# 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 UMAMI channel of every +provider becomes one entry of subprocesses.json: colour orderings integrated +over a common set of colour-ordered phase-space channels, which differ in where +the colour singlets attach and in the s-channel resonances they follow, and are +weighted by their mapping densities. Subprocesses related by a flavour +relabelling that is a symmetry of the allowed vertices are generated once +(pyamplicol's strict flavour-blind grouping) and evaluated through their +representative. +""" + +import hashlib +import json +import logging +import os +import re +import shlex +import subprocess +import sys +import tempfile +import time + +import madgraph + +logger = logging.getLogger('madgraph.export_amplicol') +pjoin = os.path.join + +MODEL_FILE = 'amplicol_model.pyamplicol-model' +# the pyamplicol that generated the output, checked before compiling +PYAMPLICOL_RECORD = 'amplicol_pyamplicol.json' + +MAKEFILE = """\ +# Builds the AmpliCol matrix elements of this process into ../../lib. +RUSTICOL_CONFIG ?= $(CURDIR)/../rusticol-config +BUILD := $(CURDIR)/build + +all: +\t$(CURDIR)/../check-pyamplicol +\t$(MAKE) -C lc/API/umami shared BUILD_DIR=$(BUILD)/lc RUSTICOL_CONFIG=$(RUSTICOL_CONFIG) +\t$(MAKE) -C reweight shared BUILD_DIR=$(BUILD)/reweight RUSTICOL_CONFIG=$(RUSTICOL_CONFIG) +%(copies)s +clean: +\trm -rf $(BUILD) +""" + + +WORKER = pjoin(os.path.dirname(os.path.abspath(__file__)), 'amplicol_worker.py') +# printed by Symbolica when another process holds its per-user thread +SYMBOLICA_BUSY = 'Cannot start another restricted Symbolica thread' +# seconds between the attempts while Symbolica is busy, and before giving up +SYMBOLICA_RETRY = 10 +SYMBOLICA_MAX_WAIT = 3600 + + +# rusticol-config of the pyamplicol in pyamplicol/install, for the makefiles +RUSTICOL_CONFIG = """\ +#!/bin/sh +PYTHONPATH=%(install)s${PYTHONPATH:+:$PYTHONPATH} exec %(python)s -m rusticol_config "$@" +""" + +# fails unless that pyamplicol is the one recorded in PYAMPLICOL_RECORD +CHECK_PYAMPLICOL = """\ +#!/bin/sh +PYTHONPATH=%(install)s${PYTHONPATH:+:$PYTHONPATH} exec %(python)s -c \\ + 'import sys; sys.path.insert(0, sys.argv[1]); from madgraph.mg7.bootstrap import check_pyamplicol_build; check_pyamplicol_build(sys.argv[2])' \\ + %(mg_root)s "$(dirname "$0")/%(record)s" +""" + + +def check_pyamplicol(interactive=None, jobs=None): + """Install pyamplicol in pyamplicol/install if needed and put it first on + sys.path, which the worker inherits. Never imported here, see amplicol_worker. + Returns its version and source hash.""" + from madgraph.mg7 import bootstrap + try: + bootstrap.ensure_pyamplicol(interactive=interactive, jobs=jobs) + return bootstrap.check_pyamplicol_source_hash() + except RuntimeError as error: + raise madgraph.InvalidCmd('The [amplicol] processes need pyamplicol: %s' % error) + + +def write_rusticol_config(subproc_dir, pyamplicol): + """SubProcesses/rusticol-config, the RUSTICOL_CONFIG of the makefiles, and + check-pyamplicol with the pyamplicol it has to be.""" + from madgraph.mg7 import bootstrap + values = {'install': shlex.quote(str(bootstrap.PYAMPLICOL_INSTALL_DIR)), + 'python': shlex.quote(sys.executable), + 'mg_root': shlex.quote(str(bootstrap.MG_ROOT)), + 'record': PYAMPLICOL_RECORD} + for name, template in (('rusticol-config', RUSTICOL_CONFIG), + ('check-pyamplicol', CHECK_PYAMPLICOL)): + path = pjoin(subproc_dir, name) + with open(path, 'w') as stream: + stream.write(template % values) + os.chmod(path, 0o755) + with open(pjoin(subproc_dir, PYAMPLICOL_RECORD), 'w') as stream: + json.dump(pyamplicol, stream, indent=2) + + +def run_worker(job): + """Run amplicol_worker on a job, waiting while Symbolica is busy elsewhere.""" + env = dict(os.environ, SYMBOLICA_HIDE_BANNER='1', + PYTHONPATH=os.pathsep.join(p for p in sys.path if p)) + start = time.monotonic() + waiting = False + while True: + returncode, busy = run_worker_once(job, env) + if returncode == 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 time.monotonic() - start > SYMBOLICA_MAX_WAIT: + raise madgraph.InvalidCmd( + 'Symbolica has been used by another process for %d minutes: only one ' + 'pyamplicol process may run per user, stop the other one and output ' + 'again' % (SYMBOLICA_MAX_WAIT // 60)) + if not waiting: + logger.info('Symbolica is used by another process: waiting for it') + waiting = True + time.sleep(SYMBOLICA_RETRY) + + +def run_worker_once(job, env): + """The exit code of one worker run, and whether Symbolica was busy. The + worker never outlives this call, so it cannot keep holding Symbolica's lock.""" + worker = subprocess.Popen([sys.executable, WORKER, job], env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True) + busy = False + try: + for line in worker.stdout: + if SYMBOLICA_BUSY in line: + busy = True + elif not line.startswith(('\u250c', '\u2502', '\u2514', '\u251c')): + sys.stdout.write(line) + return worker.wait(), busy + finally: + if worker.poll() is None: + worker.terminate() + try: + worker.wait(timeout=10) + except subprocess.TimeoutExpired: + worker.kill() + worker.wait() + worker.stdout.close() + + +def load_metadata(artifact): + """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): + """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')) + + +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_pdgs(leg, model)] + 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 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_job(model, cache_path, pyamplicol=None): + """The UFO model of the output with its restriction, compiled for + pyamplicol once per output and again when the model, its restriction or + pyamplicol change: the cache is keyed by a sidecar json file.""" + path = model.get('modelpath') + reference = model.get('name') + # the suffix after the UFO directory name, which may itself contain '-' + suffix = model.get('restrict_name') + if not suffix: + restriction = 'default' + elif suffix == 'full': + restriction = 'none' + else: + restriction = pjoin(path, 'restrict_%s.dat' % suffix) + sources = {'model': os.path.realpath(path), 'restriction': restriction, + 'restriction_card': file_hash(pjoin(path, 'restrict_default.dat') + if restriction == 'default' else restriction), + 'pyamplicol': pyamplicol} + key = hashlib.sha256(json.dumps(sources, sort_keys=True).encode()).hexdigest() + key_file = os.path.splitext(cache_path)[0] + '.json' + if not (os.path.exists(cache_path) and cached_model_key(key_file) == key): + logger.info('Compiling the model %s for AmpliCol' % reference) + return {'path': path, 'restriction': restriction, 'cache': cache_path, + 'key': key, 'key_file': key_file, 'sources': sources} + + +def file_hash(path): + if not os.path.isfile(path): + return None + with open(path, 'rb') as stream: + return hashlib.sha256(stream.read()).hexdigest() + + +def cached_model_key(key_file): + """The key of the cached model, see amplicol_worker.model_source.""" + try: + with open(key_file) as stream: + return json.load(stream).get('key') + except (OSError, ValueError, AttributeError): + return None + + +def representative_processes(proc_defs, model): + """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, and every particle of each + final-state leg once, so that the default histograms see the whole group + (a jet, not a gluon).""" + import itertools + import madgraph.core.base_objects as base_objects + processes = [] + for procdef in proc_defs: + legs = procdef.get('legs') + initial = [leg_pdgs(leg, model) for leg in legs if not leg.get('state')] + finals = [leg_pdgs(leg, model) for leg in legs if leg.get('state')] + first = [pdgs[0] for pdgs in finals] + states = [list(ids) + first for ids in itertools.product(*initial)] + states += [[pdgs[0] for pdgs in initial] + first[:i] + [pdg] + first[i + 1:] + for i, options in enumerate(finals) for pdg in options[1:]] + for pdgs in states: + ids = pdgs[:len(initial)] + processes.append(base_objects.Process({ + 'legs': base_objects.LegList([ + base_objects.Leg({'id': pdg, 'number': number + 1, + '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: + 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 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 resonance_pdg(model, pdg, decay_pdgs): + """The resonance of the given species carrying the charge of its outgoing + decay products: W+ or W-, say, for the rows of one selector.""" + particle = model.get_particle(abs(pdg)) + if particle.get('self_antipart'): + return abs(pdg) + charge = sum(model.get_particle(p).get_charge() for p in decay_pdgs) + return abs(pdg) if abs(particle.get_charge() - charge) < 1e-6 else -abs(pdg) + + +def subprocess_entries(metadata, model, rel_dir, me_path, ratio_path, qcd_power, + initial_legs): + """One subprocesses.json entry per channel of a UMAMI provider; initial_legs + 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']): + phasespace = channel['phasespace_channels'] + flavors = [] + for flavor_index, process in enumerate(channel['processes']): + # the rows of a selector share its evaluation and incoming legs: one + # flavour per mirror flag and resonance charges, its options written + # in proportion to factor + by_kind = {} + for term in process['matrix_elements']: + pdgs = pdg_ids[term['pdg_ids']] + mirror = mirrored(pdgs, incoming, initial_legs, exported) + resonances = [ + [resonance_pdg(model, r['pdg'], [pdgs[leg] for leg in r['legs']]) + for r in phasespace[index]['resonances']] + for index in process['multichannels']] + flavor = by_kind.setdefault((mirror, str(resonances)), { + 'index': flavor_index, 'options': [], 'weights': [], + 'factor': 0, 'mirror': mirror, + 'color_index': process['color_flows'] or 0, + 'channels': process['multichannels'], + 'resonance_pdgs': resonances}) + flavor['options'].append(pdgs) + flavor['weights'].append(term['factor']) + flavor['factor'] += term['factor'] + flavors.extend(by_kind.values()) + first = flavors[0]['options'][0] + 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_channels': [ + {'order': ps['phasespace_order'], 'resonances': ps['resonances']} + for ps in phasespace], + 'colored_resonances': channel.get('colored_resonances', []), + 'masses': provider['masses'], + 'supports_alpha_s': provider['supports_alpha_s'], + 'flavors': flavors, + '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}), + '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, multiparticles=None, + interactive=None, jobs=None): + """Write SubProcesses/A_ for every '[amplicol]' process definition + and return the subprocesses.json entries. ``interactive`` and ``jobs`` are + those of the pyamplicol installation, see bootstrap.ensure_pyamplicol.""" + pyamplicol = check_pyamplicol(interactive, jobs) + subproc_dir = pjoin(dir_path, 'SubProcesses') + write_rusticol_config(subproc_dir, pyamplicol) + processes = [] + for number, procdef in enumerate(proc_defs, start=1): + name = 'A%d_%s' % (number, process_name(procdef, model, multiparticles or {})) + 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), + pyamplicol), + '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) + # 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) + 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, 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' % + (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, qcd_power, initial_legs)) + with open(pjoin(path, 'makefile'), 'w') as stream: + stream.write(MAKEFILE % {'copies': '\n'.join(copies)}) + with open(pjoin(path, 'process.json'), 'w') as stream: + json.dump({'process': procdef.nice_string(prefix=False), + 'expression': process['expression'], + 'multiparticles': process['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 debabcae47..dad9e27728 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3223,7 +3223,7 @@ def mg7_launcher_source(interpreter, mg5_dir): " os.execv(_INTERPRETER, [_INTERPRETER,\n" " os.path.abspath(__file__)] + sys.argv[1:])\n" + "sys.path.append(%r)\n" % mg5_dir - + "from madgraph.iolibs.template_files.mg7.launch import main\n" + + "from madgraph.mg7.launch import main\n" "if __name__ == '__main__':\n" " os.chdir(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))\n" " try:\n" @@ -3319,7 +3319,7 @@ def copy_template(self, model): "#! /usr/bin/env python3\n" "import sys, os\n" f"sys.path.append('{MG5DIR}')\n" - "from madgraph.iolibs.template_files.mg7.npy_to_lhe import main\n" + "from madgraph.mg7.npy_to_lhe import main\n" "if __name__ == '__main__':\n" " main(me_dir=os.path.dirname(os.path.dirname(os.path.realpath(__file__))))\n" ) @@ -3331,7 +3331,7 @@ def copy_template(self, model): "#! /usr/bin/env python3\n" "import sys, os\n" f"sys.path.append('{MG5DIR}')\n" - "from madgraph.iolibs.template_files.mg7.init_custom_func import main\n" + "from madgraph.mg7.init_custom_func import main\n" "if __name__ == '__main__':\n" " sys.exit(main(me_dir=os.path.dirname(os.path.dirname(os.path.realpath(__file__)))))\n" ) @@ -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,41 @@ 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 [] + # 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', {})) + # how pyamplicol is installed if it is missing + self.amplicol_install = { + 'interactive': bool(getattr(cmd, 'use_rawinput', False)), + 'jobs': getattr(cmd, 'options', {}).get('nb_core'), + } + + def amplicol_processes(self): + """MG processes standing for the [amplicol] processes.""" + 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): + 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_multiparticles, + **getattr(self, 'amplicol_install', {})): + # 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 @@ -3595,10 +3629,14 @@ def create_run_card(self, matrix_elements, history): for me in matrix_elements['matrix_elements']] except (KeyError, TypeError): processes = None + amplicol = bool(self.amplicol_processes()) + if amplicol: + processes = (processes or []) + [self.amplicol_processes()] if processes: run_card.create_default_for_process(self.proc_characteristic, - history, processes) + history, processes, + amplicol=amplicol) # persist the model so the runtime can reload it: to compute the # widths set to 'auto' in the param_card (and recompute them at # each scan point) and to reset the parameters the model derives @@ -3631,8 +3669,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 diff --git a/madgraph/iolibs/template_files/mg7/custom_function_loader.py b/madgraph/iolibs/template_files/mg7/custom_function_loader.py index 66918b9316..6fcc99c176 100644 --- a/madgraph/iolibs/template_files/mg7/custom_function_loader.py +++ b/madgraph/iolibs/template_files/mg7/custom_function_loader.py @@ -73,6 +73,32 @@ def scale(*args): return scale +def _python_function(path, name, device_types): + """The function ``name`` of the Python module at ``path``, its FRAMEWORK, + checked against the devices of the run, and the module.""" + module = _import_module(path) + function = getattr(module, name, None) + if not callable(function): + raise RuntimeError(f"{path} does not define a function {name}") + framework = getattr(module, "FRAMEWORK", "numpy") + if framework not in ("numpy", "torch"): + raise RuntimeError(f"{path}: FRAMEWORK must be 'numpy' or 'torch', not {framework!r}") + if framework == "numpy" and any(device != "cpu" for device in device_types): + raise RuntimeError( + f"{path}: numpy functions run on the CPU only, use FRAMEWORK = 'torch' " + "for GPU runs" + ) + if framework == "torch" and "cpu" in device_types: + import torch + + # madspace already splits CPU batches over its threads + if torch.get_num_threads() != 1: + logger.info("Limiting torch to one thread per madspace thread") + torch.set_num_threads(1) + + return function, framework, module + + def load_custom_function(ms, contexts, device_types, kind, path, parameters=None): """Load the custom ``kind`` function at ``path`` into each context, in the same order everywhere, with the values of ``parameters`` ({name: value}) @@ -102,27 +128,8 @@ def load_custom_function(ms, contexts, device_types, kind, path, parameters=None f"[custom_functions] {kind} = {path!r}: expected a .so library or a .py file" ) - module = _import_module(path) name = f"custom_{kind}" - function = getattr(module, name, None) - if not callable(function): - raise RuntimeError(f"{path} does not define a function {name}") - framework = getattr(module, "FRAMEWORK", "numpy") - if framework not in ("numpy", "torch"): - raise RuntimeError(f"{path}: FRAMEWORK must be 'numpy' or 'torch', not {framework!r}") - if framework == "numpy" and any(device != "cpu" for device in device_types): - raise RuntimeError( - f"{path}: numpy functions run on the CPU only, use FRAMEWORK = 'torch' " - "for GPU runs" - ) - if framework == "torch" and "cpu" in device_types: - import torch - - # madspace already splits CPU batches over its threads - if torch.get_num_threads() != 1: - logger.info("Limiting torch to one thread per madspace thread") - torch.set_num_threads(1) - + function, framework, module = _python_function(path, name, device_types) declared = getattr(module, "PARAMETERS", None) unknown = sorted(set(parameters) - set(declared or {})) if unknown: @@ -146,3 +153,86 @@ def load_custom_function(ms, contexts, device_types, kind, path, parameters=None context.add_matrix_element(python_function, parameters=parameters) for context in contexts ] + + +def _exports(path, symbol): + """Whether the shared library at ``path`` exports ``symbol``.""" + import ctypes + + return hasattr(ctypes.CDLL(os.path.abspath(path)), symbol) + + +def load_surrogate( + ms, contexts, backends, device_types, path, subprocess, pdg_ids, param_card, + parameters=None, +): + """Load the [surrogate] function for one subprocess into each context, and + return the ``MatrixElementApi`` of each context. + + A .so path is a UMAMI library with the momenta, alpha_s and flavor index as + inputs, where ``{subprocess}`` and ``{device}`` are replaced, initialized + with ``param_card``. It gets the run-time ``parameters`` ({name: value}) of + the madmatrix libraries if it implements umami_set_parameter. A .py file + defines ``custom_surrogate(momenta, pdg_ids, alpha_s)``, which returns + |M|^2; ``pdg_ids[flavor_index]`` are the PDG ids of the subprocess for each + flavor index of its matrix element.""" + if path.endswith(".so"): + apis = [] + for context, backend, device in zip(contexts, backends, device_types): + library = path.format(subprocess=subprocess, device=backend) + if not os.path.isfile(library): + raise RuntimeError(f"[surrogate] function = {path!r}: {library} not found") + library_parameters = {} + if parameters and _exports(library, "umami_set_parameter"): + library_parameters = {str(k): float(v) for k, v in parameters.items()} + try: + api = context.load_matrix_element(library, param_card, library_parameters) + except RuntimeError as error: + if "could not set parameter" not in str(error): + raise + raise RuntimeError( + f"[surrogate] function = {path!r}: {error}. The surrogate gets the " + f"run-time parameters of the madmatrix libraries " + f"({', '.join(library_parameters)}) and has to accept them." + ) from error + if api.device_type() != device: + raise RuntimeError( + f"[surrogate] function = {path!r}: {library} was built for " + f"{api.device_type()}, but the run uses {device}" + ) + apis.append(api) + return apis + if not path.endswith(".py"): + raise RuntimeError( + f"[surrogate] function = {path!r}: expected a .so library or a .py file" + ) + function, framework, _ = _python_function(path, "custom_surrogate", device_types) + particle_count = len(pdg_ids[0]) + if framework == "torch": + import torch + + tables = {} + + def surrogate(momenta, flavor_index, alpha_s): + device = momenta.device + if device not in tables: + tables[device] = torch.tensor(pdg_ids, dtype=torch.int32, device=device) + return function(momenta, tables[device][flavor_index.long()], alpha_s) + + else: + import numpy as np + + table = np.array(pdg_ids, dtype=np.int32) + + def surrogate(momenta, flavor_index, alpha_s): + return function(momenta, table[flavor_index], alpha_s) + + python_function = ms.PythonFunction( + surrogate, + ["momenta", "flavor_index", "alpha_s"], + ["matrix_element"], + framework=framework, + particle_count=particle_count, + name=f"custom_surrogate ({subprocess})", + ) + return [context.add_matrix_element(python_function) for context in contexts] diff --git a/madgraph/iolibs/template_files/mg7/event_stream.py b/madgraph/iolibs/template_files/mg7/event_stream.py new file mode 100644 index 0000000000..ac0f59ce73 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/event_stream.py @@ -0,0 +1,347 @@ +"""In-memory event generation from a gridpack, as a Python API. + +Only available in gridpacks made with ``[gridpack] run_mode = "fix_max_weight"``: +their channels carry a fixed maximum weight, so every event is final as soon as +it is unweighted and the events can be streamed without a survey or a combine +pass. Nothing is written to disk. + +Usage:: + + import sys + sys.path.insert(0, "/bin") + from event_stream import EventStream + + with EventStream(device="cpu", seed=42) as stream: + for event in stream: # ms.LHEEvent objects, endless + ... + for batch in stream.batches(10000, output_type="torch"): + batch.px # shape (10000, max_particle_count) + +Event weights are in units of the maximum weight: +-1, except for overweight +events and for events reweighted in a later step (the default [amplicol] +mode = "reweight"), which carry the ratio of the matrix elements. To normalize +a sample to the cross section, scale its weights by +``stream.integral / sum(weights)``. The event sequence is reproducible from the +seed, independent of the thread count and of how it is split into batches. +""" + +import collections +import dataclasses +import json +import sys +from typing import Any, Iterator, Optional + +import numpy as np + +from gridpack_setup import ( + SOURCE_HASH_MESSAGE, + build_lhe_meta, + gridpack_path, + load_channels, + load_context_data, + load_lhe_completer, + load_madspace_data, + load_run_card, + load_weighting_groups, + make_context, + resolve_seed, + source_hash_matches, +) +import madspace as ms + +__all__ = ["EventStream", "EventBatch"] + +# events completed per call into madspace while iterating over single events +_ITER_CHUNK = 1000 + + +@dataclasses.dataclass +class EventBatch: + """A batch of LHE events as a structure of arrays: event fields have shape + (batch_size,), particle fields (batch_size, max_particle_count). Particle + slots beyond ``particle_count`` are padded with zeros, see + ``particle_mask``. numpy arrays or CPU torch tensors, depending on the + ``output_type`` of :meth:`EventStream.batches`.""" + + #: Process ID of the event (IDPRUP). + process_id: Any + #: Event weight (XWGTUP), +-1 or larger in magnitude for overweight events, times + #: the matrix-element ratio of a reweighting step. Not normalized. + weight: Any + #: Renormalization scale in GeV (SCALUP). + scale: Any + #: QED coupling (AQEDUP). + alpha_qed: Any + #: QCD coupling (AQCDUP). + alpha_qcd: Any + #: Number of particles of the event, including resonances (NUP). + particle_count: Any + #: True for the particle slots below ``particle_count``. + particle_mask: Any + #: PDG ID of the particle (IDUP). + pdg_id: Any + #: -1 for incoming, 1 for outgoing particles, 2 for resonances (ISTUP). + status_code: Any + #: Index of the first mother, counting from 1; 0 if there is none (MOTHUP). + mother1: Any + #: Index of the second mother, counting from 1; 0 if there is none (MOTHUP). + mother2: Any + #: Color line, 0 for a color singlet (ICOLUP). + color: Any + #: Anti-color line (ICOLUP). + anti_color: Any + #: x component of the momentum in GeV. + px: Any + #: y component of the momentum in GeV. + py: Any + #: z component of the momentum in GeV. + pz: Any + #: Energy in GeV. + energy: Any + #: Mass in GeV. + mass: Any + #: Proper lifetime c*tau in mm, 0 if not tracked (VTIMUP). + lifetime: Any + #: Cosine of the spin angle, or 9 if the spin is not specified (SPINUP). + spin: Any + + def __len__(self) -> int: + return len(self.weight) + + +def _events_to_arrays(events: list, template: dict) -> dict: + """The fields of ms.LHEEvent objects as arrays like those of + ms.EventStream.next_batch, with the dtypes and padding of ``template``.""" + arrays = {} + for key, value in template.items(): + if key == "particle_count": + data = [len(event.particles) for event in events] + elif value.ndim == 1: + data = [getattr(event, key) for event in events] + else: + data = np.zeros((len(events), value.shape[1]), value.dtype) + for row, event in zip(data, events): + row[: len(event.particles)] = [getattr(p, key) for p in event.particles] + arrays[key] = np.asarray(data, value.dtype) + return arrays + + +def _concatenate(first: dict, second: dict) -> dict: + return {key: np.concatenate([first[key], second[key]]) for key in first} + + +class EventStream: + """Endless stream of unweighted events from this gridpack. + + Constructing it starts the generation in background threads, which keep + a buffer of events for every channel. The channel of every event is picked + at random according to the current cross-section estimates, which are + refined as events are generated. + + Arguments left as None take their value from Cards/grid_run_card.toml. + + Args: + device: device to generate on, like "cpu", "cuda" or "cuda:1". + seed: run seed; -1 draws a random one (see :attr:`seed`). + cpu_batch_size: generation batch size on a CPU device. + gpu_batch_size: generation batch size on a GPU device. + thread_pool_size: number of worker threads. + cut_efficiency_threshold: fraction of a batch that has to pass the cuts. + max_cut_repetitions: how often a batch is repeated to pass the cuts. + ignore_source_hash: run even if madspace differs from the version the + gridpack was made with (can lead to errors or incorrect results). + """ + + #: The seed of the stream; a seed of -1 is replaced by the random seed drawn. + seed: int + + def __init__( + self, + device: Optional[str] = None, + seed: Optional[int] = None, + cpu_batch_size: Optional[int] = None, + gpu_batch_size: Optional[int] = None, + thread_pool_size: Optional[int] = None, + cut_efficiency_threshold: Optional[float] = None, + max_cut_repetitions: Optional[int] = None, + ignore_source_hash: bool = False, + ): + self._stream = None + # events fetched for the iterator but not returned yet + self._pending = collections.deque() + run_card = load_run_card() + run_args = run_card["run"] + gen_args = run_card["generation"] + madspace_data = load_madspace_data() + if not source_hash_matches(madspace_data): + if not ignore_source_hash: + raise RuntimeError( + f"{SOURCE_HASH_MESSAGE}. Pass ignore_source_hash=True to run anyway." + ) + print(f"WARNING: {SOURCE_HASH_MESSAGE}", file=sys.stderr) + + with open(gridpack_path("data", "channels.json")) as f: + channels = json.load(f) + if not any("max_weight" in channel for channel in channels.values()): + raise RuntimeError( + "EventStream needs a gridpack made with [gridpack] run_mode = " + '"fix_max_weight"' + ) + + def pick(value, default): + return default if value is None else value + + if device is None: + devices = run_args["device"] + device = devices[0] if isinstance(devices, list) else devices + if thread_pool_size is None: + is_gpu = device.split(":")[0] in ("cuda", "hip") + thread_pool_size = run_args[ + "gpu_thread_pool_size" if is_gpu else "cpu_thread_pool_size" + ] + self._context, backend = make_context( + device, run_args["cpu_mode"], thread_pool_size, thread_pool_size + ) + + config = ms.GeneratorConfig() + config.cpu_batch_size = pick(cpu_batch_size, gen_args["cpu_batch_size"]) + config.gpu_batch_size = pick(gpu_batch_size, gen_args["gpu_batch_size"]) + config.cut_efficiency_threshold = pick( + cut_efficiency_threshold, gen_args["cut_efficiency_threshold"] + ) + config.max_cut_repetitions = pick( + max_cut_repetitions, gen_args["max_cut_repetitions"] + ) + + self.seed = resolve_seed(pick(seed, run_args.get("seed", -1))) + load_context_data([self._context], [backend], madspace_data) + channel_generators = load_channels([self._context], config) + # the maxima of all passes are fixed, so their settings do not matter + groups = load_weighting_groups( + gen_args["max_overweight_truncation"], gen_args["freeze_max_weight_after"], + gen_args["max_overweight_truncation"], gen_args["freeze_max_weight_after"], + ) + self._stream = ms.EventStream( + self._context, channel_generators, self.seed, load_lhe_completer(), config, + groups=groups, + ) + + def __iter__(self) -> Iterator["ms.LHEEvent"]: + """Endless iterator over the events as ms.LHEEvent objects. It takes + events from the stream in chunks of 1000. The events of a chunk that are + not iterated over are kept: the next calls return them first, so no + event is lost when the loop is left early.""" + while True: + if not self._pending: + self._pending.extend(self._stream.next_events(_ITER_CHUNK)) + yield self._pending.popleft() + + def _take_pending(self, count: int) -> list: + return [self._pending.popleft() for _ in range(min(count, len(self._pending)))] + + def next_events(self, count: int) -> list: + """The next ``count`` events as ms.LHEEvent objects.""" + events = self._take_pending(count) + if len(events) < count: + events.extend(self._stream.next_events(count - len(events))) + return events + + def next_batch(self, batch_size: int, output_type: str = "numpy") -> EventBatch: + """The next ``batch_size`` events as an :class:`EventBatch`.""" + if output_type not in ("numpy", "torch"): + raise ValueError('output_type must be "numpy" or "torch"') + pending = self._take_pending(batch_size) + arrays = self._stream.next_batch(batch_size - len(pending)) + if pending: + arrays = _concatenate(_events_to_arrays(pending, arrays), arrays) + counts = arrays["particle_count"] + max_count = arrays["pdg_id"].shape[1] + arrays["particle_mask"] = np.arange(max_count)[None, :] < counts[:, None] + if output_type == "torch": + import torch + arrays = {key: torch.from_numpy(value) for key, value in arrays.items()} + return EventBatch(**arrays) + + def batches( + self, + batch_size: int, + output_type: str = "numpy", + max_events: Optional[int] = None, + ) -> Iterator[EventBatch]: + """Iterator over batches of ``batch_size`` events, endless unless + ``max_events`` is given (the last batch may then be smaller).""" + if batch_size <= 0: + raise ValueError("batch_size must be positive") + remaining = max_events + while remaining is None or remaining > 0: + size = batch_size if remaining is None else min(batch_size, remaining) + yield self.next_batch(size, output_type) + if remaining is not None: + remaining -= size + + @property + def integral(self) -> float: + """Current estimate of the cross section, in pb.""" + return self._stream.status().mean + + @property + def integral_error(self) -> float: + """Standard error of :attr:`integral`.""" + return self._stream.status().error + + @property + def abs_integral(self) -> float: + """Current estimate of the integral of the absolute weights.""" + return self._stream.status().mean_abs + + @property + def n_weighted(self) -> int: + """Number of weighted events the estimates are based on.""" + return self._stream.status().count + + @property + def n_weighted_after_cuts(self) -> int: + """Number of those weighted events that passed the cuts.""" + return self._stream.status().count_after_cuts + + @property + def n_events(self) -> int: + """Number of events returned so far.""" + return self._stream.event_count() + + @property + def max_particle_count(self) -> int: + """Largest number of particles per event, including resonances.""" + return self._stream.max_particle_count() + + def status(self): + """Combined ms.GeneratorStatus over all channels.""" + return self._stream.status() + + def channel_status(self) -> list: + """ms.GeneratorStatus of every channel; ``count_unweighted`` is the + number of its events returned so far.""" + return self._stream.channel_status() + + def channel_probabilities(self) -> list: + """Current probability of every channel to be picked for an event.""" + return self._stream.channel_probabilities() + + def lhe_meta(self): + """ms.LHEMeta for writing the events to an LHE file, with the current + cross section estimate.""" + return build_lhe_meta(self._stream.status(), self.seed) + + def close(self) -> None: + """Stop the background generation.""" + if self._stream is not None: + self._stream.close() + + def __enter__(self) -> "EventStream": + return self + + def __exit__(self, *exc) -> None: + self.close() + + def __del__(self) -> None: + self.close() diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index 2f1ab4f6cd..1046b928e3 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -1,40 +1,33 @@ #! /usr/bin/env python3 +import argparse import atexit +import glob import gzip +import json import os import shutil import subprocess import sys import tempfile -from pathlib import Path - -# search order for madspace package: -# 1. a precompiled install (madspace/install/madspace) -# 2. bundled source that still needs to be built (madspace/install.py) -# 3. otherwise fall back to madspace is available in the environment -_GRIDPACK_DIR = Path(os.path.realpath(__file__)).parent.parent -_LOCAL_MADSPACE_DIR = _GRIDPACK_DIR / "madspace" -_LOCAL_INSTALL_DIR = _LOCAL_MADSPACE_DIR / "install" -if (_LOCAL_INSTALL_DIR / "madspace").is_dir(): - sys.path.insert(0, str(_LOCAL_INSTALL_DIR)) -elif (_LOCAL_MADSPACE_DIR / "install.py").is_file(): - print() - print("You don't have madspace installed for this gridpack") - print("Running interactive madspace installation script") - print() - - _result = subprocess.run([sys.executable, str(_LOCAL_MADSPACE_DIR / "install.py")]) - if _result.returncode != 0: - raise RuntimeError("madspace installation failed — see output above") - sys.path.insert(0, str(_LOCAL_INSTALL_DIR)) +# bin/gridpack_setup.py, next to this script: also locates madspace +from gridpack_setup import ( + SOURCE_HASH_MESSAGE, + build_lhe_meta, + load_channels, + load_context_data, + load_lhe_completer, + load_madspace_data, + load_run_card, + load_systematics, + load_weighting_groups, + make_context, + resolve_seed, + source_hash_matches, +) import madspace as ms -from custom_function_loader import load_custom_function -import glob -import json -import tomllib -import argparse + def resolve_verbosity(verbosity: str) -> str: """Resolve the run_card "auto" verbosity to "pretty"/"log" depending on @@ -45,14 +38,6 @@ def resolve_verbosity(verbosity: str) -> str: return verbosity -def resolve_seed(seed: int) -> int: - """Resolve the run_card "seed": -1 draws a fresh 64-bit seed via - os.urandom, any other value is used as-is.""" - if seed == -1: - return int.from_bytes(os.urandom(8), "big") - return seed - - def compress_lhe(path: str, use_pigz="auto", cpu_threads: int = -1) -> None: """Compress the completed LHE with the same size/tool policy as misc.gzip. @@ -76,39 +61,6 @@ def compress_lhe(path: str, use_pigz="auto", cpu_threads: int = -1) -> None: os.remove(path) -def build_lhe_meta(event_generator, seed: int, systematics=None): - """The LHE header/ metadata of this gridpack run: the cards, beams - and PDF the gridpack was made with (data/lhe_meta.json), with the cross - section and seed of this run. ``systematics`` adds the block, - for the writers that do not inject it themselves (header.lhe); leave it - unset for combine_to_lhe, which does.""" - headers = [] - data = {} - meta_path = os.path.join("data", "lhe_meta.json") - if os.path.exists(meta_path): - with open(meta_path) as f: - data = json.load(f) - headers = [ms.LHEHeader(name=h["name"], content=h["content"], - escape_content=h["escape_content"]) - for h in data.pop("headers")] - headers.append(ms.LHEHeader(name="MG7Seed", content=str(seed))) - if systematics is not None and systematics.weight_ids: - headers.append(ms.LHEHeader( - name="initrwgt", content=systematics.initrwgt(), escape_content=False - )) - status = event_generator.status() - # IDWTUP and XMAXUP of this run, as launch.lhe_weight_info: negative when - # the sample can hold negative weights, and the unit weight sigma_abs - negative = status.mean < 0 or status.mean_abs > abs(status.mean) * (1 + 1e-12) - data["weight_mode"] = -4 if negative else 3 - return ms.LHEMeta( - # positional: the pybind arg name for max_weight is non-kwarg-safe - processes=[ms.LHEProcess(status.mean, status.error, status.mean_abs, 1)], - headers=headers, - **data, - ) - - def write_lhe_header(run_path: str, meta) -> None: """Write header.lhe next to events.npy: the
/ blocks the npy formats otherwise drop, with no events.""" @@ -121,15 +73,12 @@ def write_lhe_header(run_path: str, meta) -> None: def main() -> None: - # load run card and metadata - with open(os.path.join("Cards", "grid_run_card.toml"), "rb") as f: - run_card = tomllib.load(f) + run_card = load_run_card() run_args = run_card["run"] gen_args = run_card["generation"] grid_args = run_card.get("gridpack", {}) - param_card_path = os.path.join("Cards", "param_card.dat") - with open(os.path.join("data", "data.json")) as f: - madspace_data = json.load(f) + amplicol_args = run_card.get("amplicol", {}) + madspace_data = load_madspace_data() # parse command line arguments parser = argparse.ArgumentParser() @@ -183,17 +132,27 @@ def main() -> None: choices=["lhe", "lhe_npy", "compact_npy"] ) parser.add_argument("--events", type=int, default=gen_args["events"]) + # these two set the final pass, which gives the events parser.add_argument("--max_overweight_truncation", type=float, default=gen_args["max_overweight_truncation"]) parser.add_argument("--freeze_max_weight_after", type=int, default=gen_args["freeze_max_weight_after"]) + parser.add_argument( + "--first_pass_max_overweight_truncation", type=float, + default=amplicol_args.get("first_pass_max_overweight_truncation", 0.01), + help="max_overweight_truncation of the leading-colour unweighting of " + "[amplicol] processes unweighted to full colour" + ) + parser.add_argument( + "--first_pass_freeze_max_weight_after", type=int, + default=amplicol_args.get("first_pass_freeze_max_weight_after", 10000), + help="freeze_max_weight_after of the leading-colour unweighting of " + "[amplicol] processes unweighted to full colour" + ) parser.add_argument("--cpu_batch_size", type=int, default=gen_args["cpu_batch_size"]) parser.add_argument("--gpu_batch_size", type=int, default=gen_args["gpu_batch_size"]) args = parser.parse_args() - if madspace_data["source_hash"] != ms.SOURCE_HASH: - message = ( - "The madspace version is not identical to the one used to generate " - "the gridpack. This can lead to errors or incorrect results" - ) + if not source_hash_matches(madspace_data): + message = SOURCE_HASH_MESSAGE if not args.ignore_source_hash: sys.exit( f"\033[1m\033[31mERROR\033[39m: {message}.\n" @@ -248,33 +207,25 @@ def resolve_dir(cli_value, card_value): # initialize context device_names = args.device if args.device else run_args["device"] - cpu_mode = run_args["cpu_mode"] contexts = [] backends = [] for device_name in device_names: - if ":" in device_name: - device_type, device_index_str = device_name.split(":") - device_index = int(device_index_str) - else: - device_type = device_name - device_index = 0 - # cpu_mode names the SIMD width of the CPU code, so it applies to the - # 'cpu' devices only: cuda/hip build the backend named after the device. - backends.append(cpu_mode if device_type == "cpu" else device_type) - if device_type == "cuda": - device = ms.cuda_device(device_index) - pool_size = args.gpu_thread_pool_size - elif device_type == "hip": - device = ms.hip_device(device_index) - pool_size = args.gpu_thread_pool_size - else: - device = ms.cpu_device() - pool_size = args.cpu_thread_pool_size - contexts.append(ms.Context(device=device, thread_count=pool_size)) + context, backend = make_context( + device_name, + run_args["cpu_mode"], + args.cpu_thread_pool_size, + args.gpu_thread_pool_size, + ) + contexts.append(context) + backends.append(backend) # set up generator configuration config = ms.GeneratorConfig() config.target_count = args.events + groups = load_weighting_groups( + args.max_overweight_truncation, args.freeze_max_weight_after, + args.first_pass_max_overweight_truncation, args.first_pass_freeze_max_weight_after, + ) config.max_overweight_truncation = args.max_overweight_truncation config.freeze_max_weight_after = args.freeze_max_weight_after config.cpu_batch_size = args.cpu_batch_size @@ -284,50 +235,21 @@ def resolve_dir(cli_value, card_value): config.cut_efficiency_threshold = gen_args["cut_efficiency_threshold"] config.max_cut_repetitions = gen_args["max_cut_repetitions"] - # set up contexts - # run-time matrix-element parameters (bwcutoff) of the run that made the - # gridpack; gridpacks written before they were recorded used the default - me_parameters = madspace_data.get("me_parameters", {}) - globals_path = os.path.join("data", "globals.tar") - for context in contexts: - context.load_globals(globals_path) - # custom functions come first, as in the run that made the gridpack - device_types = [b if b in ("cuda", "hip") else "cpu" for b in backends] - for entry in madspace_data.get("custom_functions", []): - # [kind, path, parameters]; older gridpacks store fewer entries - entry = entry if isinstance(entry, list) else ["", entry] - kind, path, parameters = (entry + [{}])[:3] - load_custom_function(ms, contexts, device_types, kind, path, parameters) - for context, backend in zip(contexts, backends): - for me_path in madspace_data["matrix_elements"]: - context.load_matrix_element( - me_path.format(device=backend), param_card_path, me_parameters - ) - - # set up generators - with open(os.path.join("data", "channels.json")) as f: - channels = json.load(f) - channel_generators = [ - ms.ChannelEventGenerator.load_json( - json.dumps(channel), - contexts, - event_file=os.path.join(temp_dir, f"events.{name}.npy"), - weight_file=os.path.join(temp_dir, f"weights.{name}.npy"), - config=config, - ) - for name, channel in channels.items() - ] + # set up contexts and generators + load_context_data(contexts, backends, madspace_data) + channel_generators = load_channels(contexts, config, temp_dir) event_generator = ms.EventGenerator( contexts=contexts, channels=channel_generators, status_file=ms.StatusFile(os.path.join(run_path, "info.json")), config=config, seed=seed, + groups=groups, ) # scale/PDF systematics (as configured when the gridpack was made) - systematics = load_systematics(run_card, backends, param_card_path, - me_parameters) + systematics = load_systematics(run_card, backends, + madspace_data.get("me_parameters", {})) # run generation event_generator.generate() @@ -337,24 +259,24 @@ def resolve_dir(cli_value, card_value): os.path.join(run_path, "events.npy"), systematics ) write_lhe_header(run_path, - build_lhe_meta(event_generator, seed, systematics)) + build_lhe_meta(event_generator.status(), seed, systematics)) # what npy_to_lhe needs to complete these events into LHE later shutil.copy(os.path.join("data", "lhe.json"), os.path.join(run_path, "lhe_completer.json")) elif output_format == "lhe_npy": - lhe_completer = ms.LHECompleter.load(os.path.join("data", "lhe.json")) + lhe_completer = load_lhe_completer() event_generator.combine_to_lhe_npy( os.path.join(run_path, "events.npy"), lhe_completer, systematics ) write_lhe_header(run_path, - build_lhe_meta(event_generator, seed, systematics)) + build_lhe_meta(event_generator.status(), seed, systematics)) elif output_format == "lhe": - lhe_completer = ms.LHECompleter.load(os.path.join("data", "lhe.json")) + lhe_completer = load_lhe_completer() lhe_path = os.path.join(run_path, "events.lhe") # the cross section is known once generate() has converged, which is # before combine_to_lhe writes the block event_generator.combine_to_lhe( - lhe_path, lhe_completer, build_lhe_meta(event_generator, seed), + lhe_path, lhe_completer, build_lhe_meta(event_generator.status(), seed), systematics ) compress_lhe(lhe_path, run_args.get("use_pigz", "auto"), @@ -369,121 +291,6 @@ def resolve_dir(cli_value, card_value): json.dump(data, f, indent=1) -def _pdf_search_paths(stored_path): - """Directories in which the LHAPDF grids may live: the LHAPDF_DATA_PATH - entries, the paths known to the lhapdf module, then the directory the - gridpack was created with.""" - paths = [] - env = os.environ.get("LHAPDF_DATA_PATH") - if env: - paths += env.split(os.pathsep) - try: - import lhapdf - paths += list(lhapdf.paths()) - except Exception: - pass - if stored_path: - paths.append(stored_path) - return paths - - -def _locate_pdf_file(stored_file): - """Find a PDF grid/info file: the stored absolute path, else the same - / relative to one of the PDF search paths.""" - if os.path.exists(stored_file): - return stored_file - set_dir, name = os.path.split(stored_file) - set_root, set_name = os.path.split(set_dir) - for base in _pdf_search_paths(set_root): - candidate = os.path.join(base, set_name, name) - if os.path.exists(candidate): - return candidate - raise FileNotFoundError( - "PDF file %s not found (set LHAPDF_DATA_PATH to the LHAPDF data directory)" - % stored_file) - - -def load_systematics(run_card, backends=(), param_card_path=None, - me_parameters=None): - """Rebuild the ms.SystematicsCalculator saved with the gridpack - (data/systematics.json) when [systematics] enable is set; None otherwise. - The matrix elements of the mixed-order subprocesses are reloaded into a CPU - context (when a CPU library is available) for the mu_R variations.""" - path = os.path.join("data", "systematics.json") - try: - enabled = bool(run_card["systematics"]["enable"]) - except Exception: - enabled = False - if not enabled or not os.path.exists(path): - return None - with open(path) as f: - data = json.load(f) - config = ms.SystematicsConfig.from_json(json.dumps(data["config"])) - for spec in config.pdf_members: - spec.grid_file = _locate_pdf_file(spec.grid_file) - spec.info_file = _locate_pdf_file(spec.info_file) - subproc_args = [ - ms.SubprocessSystArgs.from_json(json.dumps(a)) for a in data["subproc_args"] - ] - nominal_pdf = nominal_pdf2 = None - if config.has_pdf: - nominal_pdf = ms.PdfGrid(_locate_pdf_file(data["nominal_grid_file"])) - # the second beam's set, recorded only when it differs from the first - if data.get("nominal_grid_file2"): - nominal_pdf2 = ms.PdfGrid(_locate_pdf_file(data["nominal_grid_file2"])) - # The mu_R variations always need an alpha_s grid, so a gridpack that does - # not name one cannot reweight anything -- drop the systematics rather than - # die here (gridpacks written before the launcher recorded the file for a - # run without parton luminosity). - if not data.get("nominal_info_file"): - print("WARNING systematics: the gridpack records no alpha_s .info file, " - "the scale/PDF weights are dropped") - return None - nominal_alpha_s = ms.AlphaSGrid(_locate_pdf_file(data["nominal_info_file"])) - # PDFs, alpha_s and matrix elements are evaluated on this CPU context - context = ms.Context(device=ms.cpu_device(), thread_count=1) - matrix_elements, flavor_remap = [], [] - need = [i for i, a in enumerate(subproc_args) if a.qcd_power < 0] - # the CPU backend of this run (cpu_mode, resolved when the gridpack was - # saved), else the one the launcher used - backend = data.get("me_backend") - if need and data.get("me_paths"): - cpu = [b for b in backends if not str(b).startswith(("cuda", "hip")) and b != "auto"] - if cpu: - backend = cpu[0] - if need and backend and data.get("me_paths"): - flavor_remap = data.get("flavor_remap", []) - for i, me_path in enumerate(data["me_paths"]): - if i not in need: - matrix_elements.append(None) - continue - lib = me_path.format(device=backend) - if not os.path.exists(lib): - print("WARNING systematics: %s not found, mu_R variations of the " - "mixed-order subprocesses are dropped" % lib) - matrix_elements, flavor_remap = [], [] - break - api = context.load_matrix_element(lib, param_card_path, - me_parameters or {}) - matrix_elements.append(ms.MatrixElement( - api, - [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, - ms.MatrixElement.flavor_in], - [ms.MatrixElement.matrix_element_out], - False, - # the frame the events were generated with, not whatever the - # gridpack's run card says now: it has to match the matrix - # element inside the saved channel generators - data.get("me_frame", []), - data.get("incoming_count", 2), - )) - systematics = ms.SystematicsCalculator( - config, subproc_args, nominal_pdf, nominal_alpha_s, - context, matrix_elements, flavor_remap, nominal_pdf2=nominal_pdf2) - for warning in systematics.warnings: - print("WARNING systematics: %s" % warning) - return systematics - if __name__ == '__main__': os.chdir(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) try: diff --git a/madgraph/iolibs/template_files/mg7/gridpack_setup.py b/madgraph/iolibs/template_files/mg7/gridpack_setup.py new file mode 100644 index 0000000000..89fce519f6 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/gridpack_setup.py @@ -0,0 +1,324 @@ +"""Gridpack setup shared by bin/generate_events and bin/event_stream.py: locating +madspace, reading the cards and data files, and building the contexts, channels, +LHE completer and systematics of the gridpack. Paths are resolved relative to +the gridpack directory, so nothing depends on the current directory.""" + +import json +import os +import subprocess +import sys +from pathlib import Path + +GRIDPACK_DIR = Path(os.path.realpath(__file__)).parent.parent + +# search order for madspace package: +# 1. a precompiled install (madspace/install/madspace) +# 2. bundled source that still needs to be built (madspace/install.py) +# 3. otherwise fall back to madspace is available in the environment +_LOCAL_MADSPACE_DIR = GRIDPACK_DIR / "madspace" +_LOCAL_INSTALL_DIR = _LOCAL_MADSPACE_DIR / "install" +if (_LOCAL_INSTALL_DIR / "madspace").is_dir(): + sys.path.insert(0, str(_LOCAL_INSTALL_DIR)) +elif (_LOCAL_MADSPACE_DIR / "install.py").is_file(): + print() + print("You don't have madspace installed for this gridpack") + print("Running interactive madspace installation script") + print() + + _result = subprocess.run([sys.executable, str(_LOCAL_MADSPACE_DIR / "install.py")]) + if _result.returncode != 0: + raise RuntimeError("madspace installation failed — see output above") + sys.path.insert(0, str(_LOCAL_INSTALL_DIR)) + +import madspace as ms + + +def gridpack_path(*parts) -> str: + """A path inside the gridpack directory.""" + return str(GRIDPACK_DIR.joinpath(*parts)) + + +def load_run_card(): + """The gridpack run card (Cards/grid_run_card.toml), as plain TOML: a + gridpack must run without a MadGraph installation.""" + import tomllib + with open(gridpack_path("Cards", "grid_run_card.toml"), "rb") as f: + return tomllib.load(f) + + +def load_madspace_data() -> dict: + """data/data.json: the matrix elements, their run-time parameters and the + madspace source hash the gridpack was made with.""" + with open(gridpack_path("data", "data.json")) as f: + return json.load(f) + + +def source_hash_matches(madspace_data: dict) -> bool: + return madspace_data["source_hash"] == ms.SOURCE_HASH + + +SOURCE_HASH_MESSAGE = ( + "The madspace version is not identical to the one used to generate " + "the gridpack. This can lead to errors or incorrect results" +) + + +def resolve_seed(seed: int) -> int: + """Resolve the run_card "seed": -1 draws a fresh 64-bit seed via + os.urandom, any other value is used as-is.""" + if seed == -1: + return int.from_bytes(os.urandom(8), "big") + return seed + + +def make_context(device_name: str, cpu_mode: str, cpu_thread_pool_size: int, + gpu_thread_pool_size: int): + """The context for a device name like "cpu", "cuda:1" or "hip", with the + backend its matrix elements are built for.""" + if ":" in device_name: + device_type, device_index_str = device_name.split(":") + device_index = int(device_index_str) + else: + device_type = device_name + device_index = 0 + # cpu_mode names the SIMD width of the CPU code, so it applies to the + # 'cpu' devices only: cuda/hip build the backend named after the device. + backend = cpu_mode if device_type == "cpu" else device_type + if device_type == "cuda": + device = ms.cuda_device(device_index) + pool_size = gpu_thread_pool_size + elif device_type == "hip": + device = ms.hip_device(device_index) + pool_size = gpu_thread_pool_size + else: + device = ms.cpu_device() + pool_size = cpu_thread_pool_size + return ms.Context(device=device, thread_count=pool_size), backend + + +def load_context_data(contexts, backends, madspace_data: dict) -> None: + """Load the globals and the matrix elements of the gridpack into every + context.""" + for context in contexts: + context.load_globals(gridpack_path("data", "globals.tar")) + # custom functions come first, as in the run that made the gridpack + from custom_function_loader import load_custom_function + device_types = [b if b in ("cuda", "hip") else "cpu" for b in backends] + for entry in madspace_data.get("custom_functions", []): + # [kind, path, parameters]; older gridpacks store fewer entries + entry = entry if isinstance(entry, list) else ["", entry] + kind, path, parameters = (entry + [{}])[:3] + load_custom_function( + ms, contexts, device_types, kind, _gridpack_relative(path), parameters + ) + # in the order of the run, which gave them their indices + for context, backend in zip(contexts, backends): + for path_format, init_path, parameters in madspace_data["matrix_elements"]: + context.load_matrix_element( + _gridpack_relative(path_format.format(device=backend)), + _gridpack_relative(init_path), + parameters, + ) + + +def load_channels(contexts, config, file_dir=None): + """The channel generators stored in data/channels.json. With `file_dir`, + their temporary event files are written there; without, events are only + kept in memory.""" + with open(gridpack_path("data", "channels.json")) as f: + channels = json.load(f) + return [ + ms.ChannelEventGenerator.load_json( + json.dumps(channel), + contexts, + event_file=os.path.join(file_dir, f"events.{name}.npy") if file_dir else "", + weight_file=os.path.join(file_dir, f"weights.{name}.npy") if file_dir else "", + config=config, + ) + for name, channel in channels.items() + ] + + +def load_weighting_groups( + max_overweight_truncation, freeze_max_weight_after, + first_pass_max_overweight_truncation, first_pass_freeze_max_weight_after, +): + """The re- and unweighting steps after the first unweighting stored in + data/groups.json, with the unweighting settings of the final pass. The + first unweighting of a group with an unweighting step is a first pass.""" + with open(gridpack_path("data", "groups.json")) as f: + groups = json.load(f) + for group in groups: + unweights = False + for step in group["steps"]: + step["max_overweight_truncation"] = max_overweight_truncation + step["freeze_max_weight_after"] = freeze_max_weight_after + unweights |= step["mode"] == "unweight" + # the card replaces the stored values, which older gridpacks lack + group.pop("first_pass_max_overweight_truncation", None) + group.pop("first_pass_freeze_max_weight_after", None) + if unweights: + group["first_pass_max_overweight_truncation"] = ( + first_pass_max_overweight_truncation + ) + group["first_pass_freeze_max_weight_after"] = first_pass_freeze_max_weight_after + return ms.weighting_groups_from_json(json.dumps(groups)) + + +def load_lhe_completer(): + return ms.LHECompleter.load(gridpack_path("data", "lhe.json")) + + +def build_lhe_meta(status, seed: int, systematics=None): + """The LHE header/ metadata of this gridpack run: the cards, beams + and PDF the gridpack was made with (data/lhe_meta.json), with the cross + section (from the generator `status`) and seed of this run. + ``systematics`` adds the block, for the writers that do not + inject it themselves (header.lhe); leave it unset for combine_to_lhe, + which does.""" + headers = [] + data = {} + meta_path = gridpack_path("data", "lhe_meta.json") + if os.path.exists(meta_path): + with open(meta_path) as f: + data = json.load(f) + headers = [ms.LHEHeader(name=h["name"], content=h["content"], + escape_content=h["escape_content"]) + for h in data.pop("headers")] + headers.append(ms.LHEHeader(name="MG7Seed", content=str(seed))) + if systematics is not None and systematics.weight_ids: + headers.append(ms.LHEHeader( + name="initrwgt", content=systematics.initrwgt(), escape_content=False + )) + # IDWTUP and XMAXUP of this run, as launch.lhe_weight_info: negative when + # the sample can hold negative weights, and the unit weight sigma_abs + negative = status.mean < 0 or status.mean_abs > abs(status.mean) * (1 + 1e-12) + data["weight_mode"] = -4 if negative else 3 + return ms.LHEMeta( + # positional: the pybind arg name for max_weight is non-kwarg-safe + processes=[ms.LHEProcess(status.mean, status.error, status.mean_abs, 1)], + headers=headers, + **data, + ) + + +def _gridpack_relative(path: str) -> str: + """Paths stored relative to the gridpack refer to the gridpack directory.""" + return path if os.path.isabs(path) else gridpack_path(path) + + +def _pdf_search_paths(stored_path): + """Directories in which the LHAPDF grids may live: the LHAPDF_DATA_PATH + entries, the paths known to the lhapdf module, then the directory the + gridpack was created with.""" + paths = [] + env = os.environ.get("LHAPDF_DATA_PATH") + if env: + paths += env.split(os.pathsep) + try: + import lhapdf + paths += list(lhapdf.paths()) + except Exception: + pass + if stored_path: + paths.append(stored_path) + return paths + + +def _locate_pdf_file(stored_file): + """Find a PDF grid/info file: the stored path, else the same / + relative to one of the PDF search paths.""" + stored_file = _gridpack_relative(stored_file) + if os.path.exists(stored_file): + return stored_file + set_dir, name = os.path.split(stored_file) + set_root, set_name = os.path.split(set_dir) + for base in _pdf_search_paths(set_root): + candidate = os.path.join(base, set_name, name) + if os.path.exists(candidate): + return candidate + raise FileNotFoundError( + "PDF file %s not found (set LHAPDF_DATA_PATH to the LHAPDF data directory)" + % stored_file) + + +def load_systematics(run_card, backends=(), me_parameters=None): + """Rebuild the ms.SystematicsCalculator saved with the gridpack + (data/systematics.json) when [systematics] enable is set; None otherwise. + The matrix elements of the mixed-order subprocesses are reloaded into a CPU + context (when a CPU library is available) for the mu_R variations.""" + path = gridpack_path("data", "systematics.json") + try: + enabled = bool(run_card["systematics"]["enable"]) + except Exception: + enabled = False + if not enabled or not os.path.exists(path): + return None + with open(path) as f: + data = json.load(f) + config = ms.SystematicsConfig.from_json(json.dumps(data["config"])) + for spec in config.pdf_members: + spec.grid_file = _locate_pdf_file(spec.grid_file) + spec.info_file = _locate_pdf_file(spec.info_file) + subproc_args = [ + ms.SubprocessSystArgs.from_json(json.dumps(a)) for a in data["subproc_args"] + ] + nominal_pdf = nominal_pdf2 = None + if config.has_pdf: + nominal_pdf = ms.PdfGrid(_locate_pdf_file(data["nominal_grid_file"])) + # the second beam's set, recorded only when it differs from the first + if data.get("nominal_grid_file2"): + nominal_pdf2 = ms.PdfGrid(_locate_pdf_file(data["nominal_grid_file2"])) + # The mu_R variations always need an alpha_s grid, so a gridpack that does + # not name one cannot reweight anything -- drop the systematics rather than + # die here (gridpacks written before the launcher recorded the file for a + # run without parton luminosity). + if not data.get("nominal_info_file"): + print("WARNING systematics: the gridpack records no alpha_s .info file, " + "the scale/PDF weights are dropped") + return None + nominal_alpha_s = ms.AlphaSGrid(_locate_pdf_file(data["nominal_info_file"])) + # PDFs, alpha_s and matrix elements are evaluated on this CPU context + context = ms.Context(device=ms.cpu_device(), thread_count=1) + matrix_elements, flavor_remap = [], [] + need = [i for i, a in enumerate(subproc_args) if a.qcd_power < 0] + # the CPU backend of this run (cpu_mode, resolved when the gridpack was + # saved), else the one the launcher used + backend = data.get("me_backend") + if need and data.get("me_paths"): + cpu = [b for b in backends if not str(b).startswith(("cuda", "hip")) and b != "auto"] + if cpu: + backend = cpu[0] + if need and backend and data.get("me_paths"): + flavor_remap = data.get("flavor_remap", []) + for i, me_path in enumerate(data["me_paths"]): + if i not in need: + matrix_elements.append(None) + continue + lib = _gridpack_relative(me_path.format(device=backend)) + if not os.path.exists(lib): + print("WARNING systematics: %s not found, mu_R variations of the " + "mixed-order subprocesses are dropped" % lib) + matrix_elements, flavor_remap = [], [] + break + api = context.load_matrix_element( + lib, gridpack_path("Cards", "param_card.dat"), me_parameters or {} + ) + matrix_elements.append(ms.MatrixElement( + api, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in], + [ms.MatrixElement.matrix_element_out], + False, + # the frame the events were generated with, not whatever the + # gridpack's run card says now: it has to match the matrix + # element inside the saved channel generators + data.get("me_frame", []), + data.get("incoming_count", 2), + )) + systematics = ms.SystematicsCalculator( + config, subproc_args, nominal_pdf, nominal_alpha_s, + context, matrix_elements, flavor_remap, nominal_pdf2=nominal_pdf2) + for warning in systematics.warnings: + print("WARNING systematics: %s" % warning) + return systematics diff --git a/madgraph/iolibs/template_files/mg7/init_custom_func.py b/madgraph/iolibs/template_files/mg7/init_custom_func.py index 90311d00a2..c8a5064f88 100644 --- a/madgraph/iolibs/template_files/mg7/init_custom_func.py +++ b/madgraph/iolibs/template_files/mg7/init_custom_func.py @@ -1,265 +1,3 @@ -"""Create a skeleton for a custom cut, scale, PDF or bias function. +"""Redirect for the bin/ scripts of outputs made before the move to madgraph/mg7.""" -The function is written for a single phase-space point in C, C++, Fortran, -CUDA or HIP. A generated UMAMI wrapper calls it on batches of events, so that -madspace can load it like a matrix element. The files go into -``custom_/`` in the process directory, and ``make`` there builds -``lib/custom_.so``. The run card option -``[custom_functions] `` takes the path of that library. - -Usage:: - - bin/init_custom_func [-f {cut,scale,pdf,bias}] [-l {c,cpp,fortran,cuda,hip}] - [-y] [--force] [--build | --no-build] - -Options that are not given are asked interactively, unless ``-y`` is passed. -""" - -import argparse -import logging -import os -import shutil -import subprocess -import sys - -logger = logging.getLogger("madevent") - -TEMPLATE_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "custom_functions") -UMAMI_HEADER = os.path.join( - os.path.dirname(os.path.dirname(os.path.realpath(__file__))), "madmatrix", "umami.h" -) - -FUNCTIONS = { - "cut": (["UMAMI_OUT_CUT"], "a cut, true keeps the event and false rejects it"), - "scale": ( - ["UMAMI_OUT_REN_SCALE", "UMAMI_OUT_FACT_SCALE_1", "UMAMI_OUT_FACT_SCALE_2"], - "the renormalization and factorization scales", - ), - "pdf": (["UMAMI_OUT_PDF"], "a parton density x f(x, Q)"), - "bias": (["UMAMI_OUT_BIAS"], "an event bias for the unweighting"), -} -# example parameter of parameters.h: (name, default, run card value) -EXAMPLE_PARAMETERS = { - "cut": ("pt_min", "20.", "30.0"), - "scale": ("ht_fraction", "0.5", "0.25"), - "pdf": ("gluon_norm", "3.", "2.5"), - "bias": ("pt0", "100.", "200.0"), -} -# language: (body template extension, description) -LANGUAGES = { - "c": ("c", "C"), - "cpp": ("cpp", "C++"), - "fortran": ("F90", "Fortran"), - "cuda": ("cu", "CUDA, for runs on NVIDIA GPUs"), - "hip": ("cu", "HIP, for runs on AMD GPUs"), - "python-numpy": ("py", "Python with numpy, for CPU runs, no compilation"), - "python-torch": ("py", "Python with PyTorch, for CPU or GPU runs, no compilation"), -} - -_NONINTERACTIVE = (not sys.stdin.isatty()) or any(a in ("-y", "--yes") for a in sys.argv[1:]) - - -def _ask(question, default, choices): - """Ask through MG5's cmd.ask when available, else input().""" - if _NONINTERACTIVE: - return default - try: - from madgraph.interface.extended_cmd import Cmd - - answer = Cmd().ask(question, default, choices=list(choices), timeout=0) - return default if answer is None else str(answer) - except Exception: - pass - try: - answer = input(f"{question} [{default}]: ").strip() - except EOFError: - return default - return answer or default - - -def ask_choice(title, options, default): - """Numbered menu over ``options`` (name -> description), returns a name.""" - names = list(options) - lines = [title] - width = max(len(name) for name in names) - for i, name in enumerate(names, 1): - lines.append(f" {i}. {name:{width}s} {options[name]}") - print("\n".join(lines)) - choices = names + [str(i) for i in range(1, len(names) + 1)] - while True: - answer = _ask("Enter a number or a name", default, choices).strip().lower() - if answer in names: - return answer - if answer.isdigit() and 1 <= int(answer) <= len(names): - return names[int(answer) - 1] - print(f" Please enter one of {', '.join(names)} or 1-{len(names)}.") - - -def ask_yes_no(question, default=True): - default_s = "y" if default else "n" - while True: - answer = _ask(question, default_s, ["y", "n"]).strip().lower() - if answer in ("y", "yes"): - return True - if answer in ("n", "no"): - return False - print(" Please enter 'y' or 'n'.") - - -def build_parser(): - parser = argparse.ArgumentParser( - prog="bin/init_custom_func", - description=( - "Create a skeleton for a custom cut, scale, PDF or bias function in " - "custom_/. Compiled languages are built into " - "lib/custom_.so, Python functions need no build." - ), - ) - parser.add_argument("-f", "--function", choices=list(FUNCTIONS)) - parser.add_argument("-l", "--language", choices=list(LANGUAGES)) - parser.add_argument( - "-y", "--yes", action="store_true", - help="do not ask, use the defaults (cut, cpp) for options that are not given", - ) - parser.add_argument( - "--force", action="store_true", help="overwrite an existing custom_ directory" - ) - build = parser.add_mutually_exclusive_group() - build.add_argument("--build", dest="build", action="store_true", default=None, - help="run make after creating the files (default)") - build.add_argument("--no-build", dest="build", action="store_false") - run_card = parser.add_mutually_exclusive_group() - run_card.add_argument("--set-run-card", dest="set_run_card", action="store_true", - default=None, help="enable the function in Cards/run_card.toml (default)") - run_card.add_argument("--no-set-run-card", dest="set_run_card", action="store_false") - return parser - - -def is_python(language): - return language.startswith("python-") - - -def function_path(function, language): - """The run card path of the function, relative to the process directory.""" - if is_python(language): - return os.path.join(f"custom_{function}", f"custom_{function}.py") - return os.path.join("lib", f"custom_{function}.so") - - -def write_files(out_dir, function, language): - """Write the function body and, for compiled languages, the wrapper, - umami.h and the Makefile.""" - extension = LANGUAGES[language][0] - body_file = f"custom_{function}.{extension}" - if is_python(language): - framework = language.split("-", 1)[1] - os.makedirs(out_dir, exist_ok=True) - shutil.copy( - os.path.join(TEMPLATE_DIR, f"{function}_{framework}.py"), - os.path.join(out_dir, body_file), - ) - return body_file - values = { - "kind": function, - "function_name": f"custom_{function}", - "output_keys": ", ".join(FUNCTIONS[function][0]), - "is_pdf": int(function == "pdf"), - "is_scale": int(function == "scale"), - "return_type": "bool" if function == "cut" else "double", - "body_file": body_file, - "language": language, - "example": EXAMPLE_PARAMETERS[function][0], - "example_default": EXAMPLE_PARAMETERS[function][1], - "example_value": EXAMPLE_PARAMETERS[function][2], - } - os.makedirs(out_dir, exist_ok=True) - shutil.copy(os.path.join(TEMPLATE_DIR, f"{function}.{extension}"), os.path.join(out_dir, body_file)) - shutil.copy(UMAMI_HEADER, os.path.join(out_dir, "umami.h")) - wrapper = "wrapper_gpu.cu" if language in ("cuda", "hip") else "wrapper_cpu.c" - for template in (wrapper, "Makefile", "parameters.h"): - with open(os.path.join(TEMPLATE_DIR, template)) as f: - text = f.read() % values - with open(os.path.join(out_dir, template), "w") as f: - f.write(text) - return body_file - - -def build(out_dir): - """Run make in ``out_dir``, logging to compile.log. Returns True on success.""" - log_path = os.path.join(out_dir, "compile.log") - with open(log_path, "w") as log: - result = subprocess.run(["make"], cwd=out_dir, stdout=log, stderr=subprocess.STDOUT) - if result.returncode != 0: - with open(log_path) as log: - print(log.read()) - print(f"Building failed, see {log_path}") - return False - return True - - -def set_run_card(me_dir, function, lib_path): - from madgraph.various.banner import RunCardMG7 - - card_path = os.path.join(me_dir, "Cards", "run_card.toml") - run_card = RunCardMG7(card_path) - run_card.set(f"custom_functions.{function}", lib_path) - run_card.write(card_path) - - -def main(me_dir=None, argv=None): - global _NONINTERACTIVE - args = build_parser().parse_args(argv) - if args.yes: - _NONINTERACTIVE = True - me_dir = me_dir or os.getcwd() - - function = args.function or ask_choice( - "Which custom function do you want to create?", - {name: desc for name, (_, desc) in FUNCTIONS.items()}, - "cut", - ) - language = args.language or ask_choice( - "Which language do you want to write it in?", - {name: desc for name, (_, desc) in LANGUAGES.items()}, - "cpp", - ) - - out_dir = os.path.join(me_dir, f"custom_{function}") - if os.path.exists(out_dir): - if not (args.force or (not _NONINTERACTIVE and ask_yes_no( - f"{out_dir} already exists. Overwrite it?", False))): - print(f"{out_dir} already exists, use --force to overwrite it.") - return 1 - shutil.rmtree(out_dir) - - body_file = write_files(out_dir, function, language) - lib_path = function_path(function, language) - print(f"Created {os.path.relpath(out_dir, me_dir)}/. Write your function in " - f"{os.path.relpath(os.path.join(out_dir, body_file), me_dir)}.") - - do_build = args.build - if is_python(language): - if do_build: - print("Python functions need no build.") - do_build = False - elif do_build is None: - do_build = ask_yes_no("Build the default function now?", True) - if do_build and not build(out_dir): - return 1 - - do_set = args.set_run_card - if do_set is None: - do_set = ask_yes_no(f"Set custom_functions.{function} = \"{lib_path}\" in the run card?", True) - if do_set: - set_run_card(me_dir, function, lib_path) - print(f"Enabled it in Cards/run_card.toml: custom_functions.{function} = \"{lib_path}\".") - else: - print(f"To enable it, set custom_functions.{function} = \"{lib_path}\" in Cards/run_card.toml.") - if not do_build and not is_python(language): - print(f"Build it with `make -C {os.path.relpath(out_dir, me_dir)}`. " - "bin/generate_events also rebuilds it before a run.") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) +from madgraph.mg7.init_custom_func import main # noqa: F401 diff --git a/madgraph/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index 6c70fbdbeb..c4a6105e79 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -1,5263 +1,3 @@ -import argparse -import gc -import math -import os -import sys -import time -from datetime import timedelta -from pathlib import Path -import glob -import shutil -import json -import subprocess -import tempfile -import re -import logging -from dataclasses import dataclass, field -from typing import Literal, NamedTuple -import resource +"""Redirect for the bin/ scripts of outputs made before the move to madgraph/mg7.""" -# 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, -) - -_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 madgraph.models.check_param_card import ParamCard, group_auto_widths -from madgraph.models.model_reader import Alphas_Runner -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 -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") - -# From this many subprocess directories to compile, the standard headers are -# precompiled once ('make PCH=1', see madmatrix.mk): each file of a small process -# spends most of its time parsing them. Below, building the precompiled header -# costs about what it saves. -PCH_MIN_SUBPROCESSES = 4 - -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 mlm_pdf_ladder(run_card, mlm_clustering, is_decay=False, leptonic=False): - """Whether an MLM sample's beam densities walk the clustering ladder, the - pdf half of madevent's rewgt: under scale_scheme "madevent", and under - "clustering_mean" with clustering_mean_pdf_ladder. The Python side of - EnergyScale::mlm_pdf_reweighting, which decides the same for the kernel; - the systematics then vary the ladder as madevent's do.""" - beam = run_card["beam"] - return bool( - mlm_clustering and not is_decay and not leptonic - and beam["pdf_reweighting"] - and (beam["scale_scheme"] == "madevent" - or beam["clustering_mean_pdf_ladder"])) - - -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, extra_make_args=(), split_flavors=False): - # appended to the make command that builds the matrix elements - self.extra_make_args = list(extra_make_args) - # one subprocess per flavour of a decay (split_decay_flavors) - self.split_flavors = split_flavors - 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_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")) - misc.configure_gzip(self.run_card.gzip_configuration()) - # 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) - merge_subprocesses = self.run_card["phasespace"]["merge_subprocesses"] - if merge_subprocesses and self.run_card["beam"]["dynamical_scale_choice"] == "mlm": - # A merged subprocess bundles matrix elements whose legs differ - - # q q~ > l l g next to q g > l l q - while the MLM clustering is - # built once per subprocess from a single assignment of flavours to - # legs. It would cluster the gluon-initiated elements as if their - # beam were a quark, and pick the history per diagram from the - # weights of one element only. - logger.warning( - "merge_subprocesses is not supported with dynamical_scale_choice " - "= mlm; the subprocesses are kept separate" - ) - merge_subprocesses = False - if 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 - single = [meta["outgoing"] for meta in self.subprocess_data - if len(clean_pids(meta["outgoing"])) == 1] - if single and self.run_card["beam"]["leptonic"]: - raise ValueError( - f"2 -> 1 processes ({single[0]} ...) need beams with a PDF " - "(leptonic = false in [beam]): with fixed beam energies the " - "one-particle phase space is a delta function in s_hat" - ) - 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() - if getattr(self, "split_flavors", False): - self.split_decay_flavors() - - def split_decay_flavors(self) -> None: - """Integrate every flavour of a decay subprocess on its own. - - A subprocess can hold several flavours with different matrix elements - (z > d d~ g and z > u u~ g in one P1_z_QQxg): its integrand samples - them, and its result is their sum. The widths of compute_decay_widths - need each of them, so each flavour becomes a subprocess of its own, - with the same matrix-element library and flavour index and the - diagrams active for it (see _single_flavor_metas). - """ - if self.merged_subprocess_data is not None: - raise ValueError( - "split_flavors needs merge_subprocesses = false: " - "merged_subprocesses.json indexes the unsplit subprocesses" - ) - self.subprocess_data = [ - single - for meta in self.subprocess_data - for single in _single_flavor_metas(meta) - ] - - 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_mode(self.run_card["run"]["madspace_cpu_mode"]) - - 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 - self.auto_temp_path = temp_dir == "auto" - if self.auto_temp_path: - self.temp_path = tempfile.mkdtemp(prefix="mg7_events_") - elif 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." - ) - parameters = self.run_card.custom_parameters(kind) - apis = load_custom_function( - ms, self.contexts, self.device_types, kind, path, parameters - ) - 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, parameters]) - logger.info("Using the custom %s function %s", kind, path) - - @staticmethod - def build_custom_function(build_dir: str) -> None: - """Run make in a custom_ directory, logging to compile.log.""" - log_path = os.path.join(build_dir, "compile.log") - with open(log_path, "w") as log: - returncode = subprocess.call( - ["make"], cwd=build_dir, stdout=log, stderr=subprocess.STDOUT - ) - if returncode != 0: - raise RuntimeError( - f"Compilation of {build_dir} failed with exit status {returncode}; " - f"see {log_path} for details" - ) - - def custom_function(self, kind: str, particle_count: int, pid_options: list): - """The ms.CustomFunction of [custom_functions] , or None.""" - api = self.custom_functions.get(kind) - if api is None: - return None - function_type = getattr(ms.CustomFunction, kind) - return ms.CustomFunction(api, function_type, particle_count, pid_options) - - def parse_observable(self, name: str, order_observable: str) -> dict: - parts = name.split("-") - sum_momenta = False - 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 apply_auto_ptj_mjj(self, cuts: dict) -> None: - """madevent's auto_ptj_mjj, from Template/LO/SubProcesses/setcuts.f. - - A merging cut already bounds every matrix-element jet from below, so - raising the jet pt and pair-mass cuts to xqcut removes no phase space - that survives the merging anyway while making the integration much more - efficient. The jet dR cuts go the other way and are dropped: - xqcut supersedes them, and a dR cut would carve a hole out of exactly - the region the parton shower is supposed to fill, leaving a gap in the - jet rates at the merging scale. - - With the option off, madevent still drops the dR cuts, and only zeroes - a pt or pair-mass cut that sits above xqcut - it never tightens. - """ - if self.is_decay or self.run_card["beam"]["dynamical_scale_choice"] != "mlm": - return - xqcut = self.run_card["phasespace"]["xqcut"] - if xqcut <= 0: - return - for key in ("jet-delta_r", "jet-lepton-delta_r"): - if cuts.get(key, {}).get("min", 0.0) > 0.0: - cuts[key]["min"] = 0.0 - auto = self.run_card["phasespace"]["auto_ptj_mjj"] - # Both halves of madevent's rule matter, and for different reasons. - # - # ptj reaches the phase space only for jets that are leaf children of - # the root: a jet inside a composite s-channel node - two jets off one - # gluon, say - has its pt_min reset to zero in PhaseSpaceMapping, and - # that node's invariant is then sampled from threshold upward knowing - # nothing about the cut. mjj is what bounds those nodes, because a - # pairwise invariant-mass minimum is carried into the propagator's own - # e_min and so into the s_min the sampler uses. - # - # Leaving mjj out is what makes the highest multiplicity misbehave at a - # tight xqcut: almost every point it proposes is vetoed, and the few - # survivors sit against a boundary it cannot see. - for key in ("jet-pt", "jet-pair_mass"): - if auto: - cuts.setdefault(key, {})["min"] = xqcut - elif cuts.get(key, {}).get("min", 0.0) > xqcut: - cuts[key]["min"] = 0.0 - - def init_cuts(self) -> None: - inf = float("inf") - order_observable = self.run_card["cuts"].get("order_by", "pt") - cuts = { - key: dict(values) - for key, values in self.run_card["cuts"].items() - if key != "order_by" - } - self.apply_auto_ptj_mjj(cuts) - 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 cuts.items() - ] - - # [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. LHAPDF itself is not needed: madspace reads the grid - files directly, so without a usable lhapdf-config (or when 'lhapdf - install' fails) the set is fetched straight from the LHAPDF server.""" - paths = lhapdf_paths() - if paths.find_set(pdf_set): - return paths - if not paths.download_path: - logger.debug("no writable PDF directory: cannot download %s", pdf_set) - return paths - logger.info("PDF set %s not found; downloading into %s", - pdf_set, paths.download_path) - if paths.config: - try: - from madgraph.interface.common_run_interface import CommonRunCmd - os.makedirs(paths.download_path, exist_ok=True) - CommonRunCmd.install_lhapdf_pdfset_static( - paths.config, paths.download_path, pdf_set) - except Exception as err: - logger.debug("lhapdf could not install %s: %s", pdf_set, err) - if not os.path.isdir(os.path.join(paths.download_path, pdf_set)): - try: - misc.download_lhapdf_set(pdf_set, paths.download_path) - 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) - # subprocesses (systematics, MadSpin, reweight) find it here too - os.environ["LHAPDF_DATA_PATH"] = os.pathsep.join(_LHAPDF.data_paths) - 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 lhapdf-config found (not needed for the download)" - if self.lhapdf.download_path: - download = ("the download from %s into %s failed (see the " - "warning above)" % (misc.LHAPDF_SETS_URL, - self.lhapdf.download_path)) - else: - download = "no writable directory to download the set into" - return ( - "LHAPDF set %r, requested by [beam] pdf1/pdf2 in Cards/run_card.toml, " - "was not found.\n" - " searched: %s\n" - " %s\n" - " %s\n" - "Fix this in one of the following ways:\n" - " * check the set name against %s/pdfsets.index and your network\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, download, misc.LHAPDF_SETS_URL, - 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.beam_energies = (float(self.e_cm), 0.0) - self.beam_rapidity = 0.0 - self.leptonic = True - else: - ebeam1 = float(beam_args["ebeam1"]) - ebeam2 = float(beam_args["ebeam2"]) - self.beam_energies = (ebeam1, ebeam2) - # The phase space is generated in the beams' centre-of-mass frame - # (massless beams: s = 4 E1 E2) and boosted by this rapidity into - # the lab frame, where the events are written and the cuts applied. - self.e_cm = 2.0 * math.sqrt(ebeam1 * ebeam2) - self.beam_rapidity = 0.5 * math.log(ebeam1 / ebeam2) - 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, - } - # "mlm" is not one of the closed-form scale choices: it needs the diagram - # topologies of each subprocess, so the EnergyScale is built later, in - # MadgraphSubprocess.init_madspace, from an MLMClustering instance. - if beam_args["dynamical_scale_choice"] == "mlm": - self.mlm_clustering = True - self.scale_kwargs = {} - elif beam_args["dynamical_scale_choice"] in dynamical_scales: - self.mlm_clustering = False - self.scale_kwargs = dict( - dynamical_scale_type=dynamical_scales[ - beam_args["dynamical_scale_choice"] - ], - 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.)), - ) - else: - raise ValueError("Unknown dynamical scale choice") - 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.mlm_clustering and not self.is_decay: - # The merging needs the scales of its own clustering history: a - # custom scale or a factor on it would move the merging scale away - # from the one the shower is matched to. - if self.custom_functions.pop("scale", None) is not None: - logger.warning( - "custom scale ignored under dynamical_scale_choice = mlm") - if float(beam_args.get("scale_factor", 1.)) != 1.: - logger.warning( - "scale_factor ignored under dynamical_scale_choice = mlm") - if self.is_decay: - # One scale is available for a decay -- the decaying mass -- so use - # it, fixed, whatever the card asks for. That includes "mlm": a - # decay has no beams to cluster into, so there is no merging to do. - if self.mlm_clustering: - logger.warning( - "dynamical_scale_choice = mlm has no meaning for a decay; " - "using the decaying mass as a fixed scale instead" - ) - self.mlm_clustering = False - self.scale_kwargs.setdefault( - "dynamical_scale_type", ms.EnergyScale.half_transverse_mass - ) - 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.pdf_grid2 = None - self.lhapdf = None - self.pdf_dir = None - self.pdf_dir2 = 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["pdf1"] - self.pdf_set = pdf_set - self.pdf_set2 = beam_args["pdf2"] - 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)) - # A second grid only when the beams differ (and have a PDF at all); - # otherwise beam 2 shares the first one. - self.pdf_grid2 = None - self.pdf_dir2 = self.pdf_dir - if self.pdf_set2 != pdf_set and not self.leptonic: - self.pdf_dir2 = self.ensure_pdf_set(self.pdf_set2).find_set(self.pdf_set2) - if self.pdf_dir2 is None: - raise RuntimeError(self._no_pdf_message(self.pdf_set2)) - if "pdf" not in self.custom_functions: - self.pdf_grid2 = ms.PdfGrid(os.path.join( - self.pdf_dir2, self.pdf_set2, f"{self.pdf_set2}_0000.dat")) - logger.info("PDF sets: %s (beam 1), %s (beam 2); alpha_s from %s", - pdf_set, self.pdf_set2, 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") - ) - # The grid only covers a band in Q. Asking for a density above its - # highest Q reads past the end of the interpolation and comes back as a - # NaN, so the scale is capped there. A dynamical scale rarely gets near - # it, but the madevent MLM scheme, which multiplies four clustering - # scales together under a fourth root, can. A custom PDF brings no grid - # to read the range from, so it is left uncapped. - grids = [grid for grid in (self.pdf_grid, self.pdf_grid2) - if isinstance(grid, ms.PdfGrid)] - self.max_scale = min(max(grid.q) for grid in grids) if grids else 0.0 - # the bottom of the same range, where the densities freeze - self.min_pdf_scale = max(min(grid.q) for grid in grids) if grids else 0.0 - for context in self.contexts: - if not isinstance(self.pdf_grid, ms.CustomPdf): - self.pdf_grid.initialize_globals(context) - if self.pdf_grid2 is not None: - # the prefix ms.Integrand reads the second beam's grid from - # (Integrand::pdf2_prefix) - self.pdf_grid2.initialize_globals(context, "beam2") - 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"]] - # the emission scale of the MLM alpha_s vertices, which only an MLM - # sample with per-vertex alpha_s has - if (self.mlm_clustering and not self.is_decay - and self.run_card["beam"]["alphas_reweighting"] != "none"): - config.alps = [float(v) for v in syst["alps"]] - else: - config.alps = [1.0] - # the pdf reweighting ladder of an MLM sample, varied as madevent's - # systematics vary it; a varied scale is kept inside the PDF grid, - # where LHAPDF freezes it for madevent - config.pdf_ladder = mlm_pdf_ladder( - self.run_card, self.mlm_clustering, self.is_decay, self.leptonic) - if config.pdf_ladder: - config.ladder_min_scale = float(self.min_pdf_scale) - config.ladder_max_scale = float(self.max_scale) - 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"] - if self.pdf_grid2 is None: - config.pdf_members = self.resolve_pdf_variations() - else: - logger.warning( - "systematics: the beams use different PDF sets (%s, %s), so " - "no PDF variation is computed; the scale variations use both " - "sets", self.pdf_set, self.pdf_set2) - config.pdf_members = [] - 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, - nominal_pdf2=self.pdf_grid2 if config.has_pdf else None) - 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_grid_file2": os.path.join( - self.pdf_dir2, self.pdf_set2, f"{self.pdf_set2}_0000.dat") - if config.has_pdf and self.pdf_grid2 is not None 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 - # a custom scale replaces dynamical_scale_choice, so no variation is redundant - custom_scale = "scale" in getattr(self, "custom_functions", {}) - if not (beam["fixed_ren_scale"] and beam["fixed_fact_scale"] or custom_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 or v.is_emission) - 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))) - emission = systematics_summary.emission_percentages(summary) - if emission is not None: - rows.append(("Emission scale variation:", format_variation(*emission))) - 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) - - # the scale below which the alpha_s of a decay is frozen (write_fixed_alphas_info) - ALPHAS_FREEZE_SCALE = 1.0 - - def write_fixed_alphas_info(self) -> str: - """Write a minimal LHAPDF ``.info`` with the alpha_s of a decay. - - 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, so only alpha_s(M) matters. Like the Fortran decay - matrix elements (Source/alfas_functions.f, ``nloop = 2``) it is the - 2-loop running from the alpha_s(m_Z) of the param card up to M -- not - alpha_s(m_Z) itself, which overshoots alpha_s(125 GeV) by 4.7% (and - h > g g, which goes like alpha_s^2, by 9.7%). - - The grid carries the running over the whole scale range and has a node - exactly at the decaying mass, so the interpolation returns alpha_s(M) - untouched. Below ALPHAS_FREEZE_SCALE (or wherever the running stops - being defined, closer to the Landau pole) alpha_s is frozen at its last - value: a light decaying particle (the muon, ...) has no QCD to speak - of, and its widths must not fail on a coupling they do not use. - """ - # SMINPUTS entry 3 is alpha_s(m_Z), the same value the Fortran - # parameter setup feeds to G. - alpha_s_mz = float(self.param_card.get_value("sminputs", 3)) - runner = Alphas_Runner(alpha_s_mz, 2, cmass=1.42, bmass=4.7) - mass = float(self.decaying_mass) - 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: - # 1 GeV to 100 TeV, and the decaying mass. - nodes = [10 ** (i / 4.0) for i in range(0, 21)] - q_values = sorted( - [q for q in nodes if abs(q / mass - 1) > 1e-9] + [mass] - ) - - def running(q): - try: - value = runner(q) - except (ArithmeticError, ValueError): - return None - return value if math.isfinite(value) and value > 0 else None - - # from the top down, freezing at the first scale where the running is - # not used (below ALPHAS_FREEZE_SCALE) or not defined - values = {} - frozen = None - for q in reversed(q_values): - value = None - if frozen is None and q >= self.ALPHAS_FREEZE_SCALE: - value = running(q) - if value is None: - if frozen is None: - if not values: - raise ValueError( - f"alpha_s({q:g} GeV) cannot be computed from " - f"alpha_s(m_Z) = {alpha_s_mz:g}" - ) - frozen = min(values) - value = values[frozen] - values[q] = value - alpha_s_values = [values[q] for q in q_values] - if frozen is not None and mass < frozen: - logger.info( - "decay mode: alpha_s frozen below %g GeV at alpha_s(%g GeV) = " - "%.6g (decaying mass %g GeV)", frozen, frozen, values[frozen], - mass, - ) - with open(path, "w") as f: - f.write("SetDesc: 2-loop running alpha_s of a decay (no beams)\n") - f.write( - "AlphaS_Qs: [%s]\n" % ", ".join(f"{q:.12g}" for q in q_values) - ) - f.write( - "AlphaS_Vals: [%s]\n" - % ", ".join(f"{a:.12g}" for a in alpha_s_values) - ) - logger.info( - "decay mode: alpha_s(mu = %g GeV) = %.6g (2-loop running from " - "alpha_s(m_Z) = %g, as madevent)", - mass, values[mass], alpha_s_mz, - ) - 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"] - 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: - # by library: the flavours of a split subprocess share one - missing = list({ - meta["me_path"]: meta for meta in self.subprocess_data - if not os.path.isfile(meta["me_path"].format(device=backend)) - }.values()) - 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() - make_args = [f"BACKEND={backend}", f"FPTYPE={fptype}", "USEBUILDDIR=1"] - make_args += self.extra_make_args - if (len(missing) >= PCH_MIN_SUBPROCESSES - and not any(arg.startswith("PCH=") for arg in make_args)): - make_args.append("PCH=1") - self.make_subprocesses(subproc_path, make_args, 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 = [] - 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 - ) - 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, - ) - 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, - min( - (channel.phasespace_mapping.random_dim() for channel in ps.channels), - default=0, - ), - ) - 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 = [] - trained_phasespaces = [] - training_args = [] - self.event_generator = None - for subproc, phasespace in zip(self.subprocesses, self.phasespaces): - if not subproc.madnis_supported: - # no flow for a one-dimensional phase space: keep its VEGAS - # channels, and their event generators, from the survey - madnis_phasespaces.append(phasespace) - continue - 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) - trained_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( - trained_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.""" - 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], list(self.beam_energies) - 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, list(self.beam_energies) - - def _lhapdf_id(self, beam=1): - """Central LHAPDF id of the PDF set of `beam` (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.pdf_set if beam == 1 else self.pdf_set2 - pdf_dir = self.pdf_dir if beam == 1 else self.pdf_dir2 - info = os.path.join(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() - 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=self._lhapdf_id(1), beam2_pdf_id=self._lhapdf_id(2), - 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")) - - include_estimates = self.gridpack_run_mode == "fix_max_weight" - # 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(include_estimates)}" - 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) - # the gridpack refers to its own copies, by paths relative to it - gridpack_custom_paths = [] - for kind, path, parameters in self.custom_function_paths: - rel_path = os.path.normpath(path) - if os.path.isabs(rel_path) or rel_path.startswith(".."): - rel_path = os.path.join("custom_functions", kind, os.path.basename(path)) - target = os.path.join(gridpack_path, rel_path) - if not os.path.exists(target): - os.makedirs(os.path.dirname(target), exist_ok=True) - shutil.copy(path, target) - gridpack_custom_paths.append([kind, rel_path, parameters]) - - 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") - shutil.copy( - os.path.join(os.path.dirname(__file__), "gridpack.py"), gen_events_file - ) - shutil.copy( - os.path.join(os.path.dirname(__file__), "custom_function_loader.py"), - os.path.join(bin_path, "custom_function_loader.py"), - ) - os.chmod(gen_events_file, 0o755) - - data = { - "matrix_elements": matrix_elements, - # loaded before the matrix elements, as in the run - "custom_functions": gridpack_custom_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"]: - archive_path = gridpack_path + ".tar.gz" - use_pigz = self.run_card["run"]["use_pigz"] - pigz = (shutil.which("pigz") if use_pigz is not False else None) - if pigz: - # multi-threaded gzip - with open(archive_path, "wb") as f: - tar = subprocess.Popen( - ["tar", "-cf", "-", "-C", self.run_path, "gridpack"], - stdout=subprocess.PIPE, - ) - gz_command = [pigz, "-c"] - cpu_threads = self.run_card["run"]["cpu_thread_pool_size"] - if cpu_threads > 0: - gz_command += ["-p", str(cpu_threads)] - gz = subprocess.run(gz_command, stdin=tar.stdout, stdout=f) - tar.stdout.close() - if tar.wait() != 0 or gz.returncode != 0: - raise RuntimeError("compressing the gridpack failed") - else: - 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 - # set from the survey by set_madnis_auto_settings - self.madnis_supported = True - self.multi_channel_data = None - - self.unmerged_meta = None - if unmerged_meta is None: - api_path_formats = [self.meta["me_path"]] - else: - api_path_formats = [] - for subproc in self.meta["subprocesses"]: - submeta = unmerged_meta[subproc] - api_path_formats.append(submeta["me_path"]) - 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 - # A flavor standing for both beam orientations draws the orientation - # before the phase-space mapping, so boost, cuts and PDFs all see it. - self.mirror_beams = not self.process.is_decay and self.has_mirrored_flavors() - # 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 - ) - - # The scale range exists to keep the densities inside the PDF grid, so - # a run without beam densities - a lepton collider or a decay - has - # none: a floor there would only throw away events whose scale is - # legitimately low, a decay of a particle lighter than min_scale say. - if self.process.leptonic: - min_scale = max_scale = 0.0 - else: - min_scale = self.process.run_card["beam"]["min_scale"] - max_scale = self.process.max_scale - if self.process.mlm_clustering: - mc_data = self.build_multi_channel_data() - all_pids = (clean_pids(self.meta["incoming"]) - + clean_pids(self.meta["outgoing"])) - self.scale = ms.EnergyScale( - ms.MLMClustering( - topologies=[topo[0] for topo in mc_data.topologies], - permutations=mc_data.permutations, - diagram_indices=mc_data.diagram_indices, - cm_energy=self.process.e_cm, - jet_scale_scheme=getattr( - ms.MLMClustering.JetScaleScheme, - self.process.run_card["beam"]["jet_scale_scheme"], - ), - scale_scheme=getattr( - ms.MLMClustering.ScaleScheme, - self.process.run_card["beam"]["scale_scheme"], - ), - pdg_color_types={ - int(key): value - for key, value in self.meta["pdg_color_types"].items() - }, - xqcut=self.process.run_card["phasespace"]["xqcut"], - bw_cutoff=self.process.run_card["phasespace"]["bw_cutoff"], - jet_radius=self.process.run_card["beam"]["jet_radius"], - hadronic=not self.process.leptonic, - external_pdg_ids=all_pids, - max_jet_flavor=self.process.run_card["beam"]["max_jet_flavor"], - parton_line_scheme=getattr( - ms.MLMClustering.PartonLineScheme, - self.process.run_card["beam"]["parton_line_scheme"], - ), - alphas_scheme=getattr( - ms.MLMClustering.AlphasScheme, - self.process.run_card["beam"]["alphas_reweighting"], - ), - pdf_reweighting=self.process.run_card["beam"][ - "pdf_reweighting" - ], - clustering_measure=getattr( - ms.MLMClustering.ClusteringMeasure, - self.process.run_card["beam"]["clustering_measure"], - ), - clustering_history=getattr( - ms.MLMClustering.ClusteringHistory, - self.process.run_card["beam"]["clustering_history"], - ), - clustering_mean_pdf_ladder=self.process.run_card["beam"][ - "clustering_mean_pdf_ladder" - ], - ), - min_scale=min_scale, - max_scale=max_scale, - ) - else: - self.scale = ms.EnergyScale( - particle_count=self.particle_count, - custom_scale=self.process.custom_function( - "scale", self.particle_count, representative_pids - ), - min_scale=min_scale, - max_scale=max_scale, - **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) - - 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, - beam_rapidity=self.process.beam_rapidity, - mirror_beams=self.mirror_beams, - ) - 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 has_mirrored_flavors(self) -> bool: - """Does any flavor of this subprocess stand for both beam orientations? - Resolved the same way build_integrands resolves flavor_mirror.""" - 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"]] - if flav["mirror"]: - return True - return False - - 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, - beam_rapidity=self.process.beam_rapidity, - mirror_beams=self.mirror_beams, - ) - 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, random_dim: int): - """Choose the MadNIS settings of this subprocess from the relative - spread rsd of the survey. random_dim is the smallest number of random - numbers of its channels: a normalizing flow needs at least two, so a - one-dimensional phase space (a 2 -> 1 process) never uses MadNIS.""" - 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 - self.madnis_supported = random_dim >= 2 - if not self.madnis_supported and self.madnis_settings["enable"]: - logger.warning( - "MadNIS disabled for subprocess %s: its phase space has %d " - "random dimension(s), a normalizing flow needs at least 2", - self.subproc_id, random_dim, - ) - self.madnis_settings["enable"] = False - - 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 = [] - for matrix_element in self.matrix_elements: - if 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, - ) - ) - partial_weights = self.process.systematics_enabled() - madnis_args = self.process.run_card["madnis"] - custom_bias = self.process.custom_function("bias", self.particle_count, flavors) - # With an MLM history picked per diagram, the cross section at the - # scales of that history is the first matrix-element evaluation rescaled - # by alpha_s and PDF ratios when every diagram has the same power of - # alpha_s (-1: mixed powers, evaluate the matrix element again) - mlm_history_qcd_power = int(self.meta.get("qcd_power", -1)) - 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, - None if self.process.leptonic else self.process.pdf_grid2, - mlm_history_qcd_power, - )) - #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 - - -def lhapdf_index_paths() -> list: - """The PDF directories to look an lhaid up in. Without LHAPDF there may be - no pdfsets.index anywhere yet, so fetch one into the download directory.""" - - paths = lhapdf_paths() - data_paths = list(paths.data_paths) - if not any(os.path.exists(os.path.join(p, "pdfsets.index")) for p in data_paths) \ - and paths.download_path \ - and misc.download_lhapdf_index(paths.download_path): - data_paths.append(paths.download_path) - return data_paths - - -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}) - # and so does its error policy: with crash_on_error a failing - # launch has to stop the script with a non-zero exit code, as a - # madevent launch does, instead of returning to the prompt. - if 'crash_on_error' in options: - self.options['crash_on_error'] = options['crash_on_error'] - 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 check_card_consistency(self): - # madevent settings given per beam ("set fixed_fac_scale1 ...") - # that the mg7 card could not follow: said once, at the end, so - # "set fixed_fac_scale1 T" then "set fixed_fac_scale2 T" stays - # quiet. Collected first: lo_set_warnings also forgets them, which - # must happen even if the generic checks below raise. - warnings = [] - if isinstance(getattr(self, "run_card", None), RunCardMG7): - warnings = self.run_card.lo_set_warnings() - try: - super().check_card_consistency() - finally: - for text in warnings: - logger.warning(text) - - 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 - - # legacy madevent cut names ("set ptj 20", "set mmll 200", - # "set etaj -1", ...) -> the [cuts] entry the LO->MG7 - # run_card conversion maps them to (_LO_CUT_MAP), with the - # same "no cut" sentinels (minimum <= 0, maximum < 0). - if rest and nlow in run_card._LO_CUT_MAP: - if rest.lower().startswith("scan"): - logger.warning("ignoring 'set %s %s': the [cuts] entries of " - "the mg7 run_card.toml cannot be scanned", - nlow, rest) - return - try: - cut, bound, val = run_card.set_lo_cut( - nlow, run_card.evaluate(rest, masses)) - except _banner_mod.InvalidCmd: - logger.warning("ignoring 'set %s %s': not a number", nlow, rest) - return - if val is None: - logger.info("set %s (mg7 cut %s.%s) of the run_card.toml: " - "no cut", nlow, cut, bound) - else: - logger.info("set %s (mg7 cut %s.%s) of the run_card.toml to %s", - nlow, cut, bound, val) - self.modified_card.add("run") - return - - # madevent parameters that are not a plain rename: the - # value is translated ("set dynamical_scale_choice 3") or - # several of them make one mg7 entry (fixed_fac_scale1 + - # fixed_fac_scale2 -> beam.fixed_fact_scale). mg7 values - # ("set sde_strategy denominators") come back as None and - # go on to the generic editor. - if rest and nlow in run_card._LO_SET_PARAMS: - try: - done = run_card.set_lo_param(nlow, rest, masses) - except _banner_mod.InvalidCmd as error: - logger.warning("ignoring 'set %s %s': %s", nlow, rest, error) - return - if done is not None: - changes, warnings = done - for key, val in changes.items(): - if val is None: - logger.info("set %s: removed %s from the run_card.toml", - nlow, key) - else: - logger.info("set %s (mg7 %s) of the run_card.toml to %s", - nlow, key, val) - for text in warnings: - logger.warning(text) - if changes: - self.modified_card.add("run") - elif not warnings: - logger.info("set %s %s: nothing to change in the run_card.toml", - nlow, rest) - 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 - - # madevent beam settings (ebeam, lpp, pdlabel, lhaid, per - # beam or for both) onto the per-beam [beam] entries - if rest and nlow in run_card.lo_beam_params: - data_paths = lhapdf_index_paths() if nlow == "lhaid" else () - try: - done = run_card.set_lo_beam_param(nlow, rest, masses, data_paths) - except _banner_mod.InvalidRunCard as error: - logger.warning("ignoring 'set %s %s': %s", name, rest, error) - return - logger.info("set %s %s: %s", nlow, rest, done) - self.modified_card.add("run") - return - - # quantities the card stores per beam: set both beams - if rest and nlow in ("e_cm", "beam.e_cm", "beam.pdf"): - value = rest if nlow == "beam.pdf" else run_card.evaluate(rest, masses) - run_card.set(nlow, value, user=True) - logger.info("set %s of the run_card.toml to %s (both beams)", - nlow, run_card["beam." + nlow.split(".")[-1]]) - 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 madgraph.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:`madgraph.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 madgraph.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 - # "DECAY auto", with the options of compute_widths that may follow - # ("auto@NLO", "auto --nlo --body_decay=3"), as in - # common_run_interface.static_check_param_card: one compute_widths per set of - # options - groups = group_auto_widths(text) - if not groups: - return - pdgs = [pdg for group in groups.values() for pdg in group] - - 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\n" % model - for options, group in groups.items(): - cmds += "compute_widths %s --path=%s %s\n" % ( - " ".join(group), os.path.abspath(param_card_path), " ".join(options)) - 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 compute_decay_widths( - me_dir: str = ".", precision: float = 0.01, nb_core: int | None = None, - extra_make_args=(), -) -> list[dict]: - """Integrate the decays of an mg7 decay directory and return their widths. - - This is the madspace counterpart of madevent's ``survey decay``: only the - survey runs (no unweighting, no events, no systematics), and it stops once - every channel is below ``precision`` relative error, like the - ``--precision_channel`` of ``compute_widths``. - - ``me_dir`` is an output directory of 1 -> n processes of one decaying - particle. It is meant to run in a process of its own (``bin/generate_events - --decay-widths=FILE``, which is how ``compute_widths`` calls it): the - madspace objects write their status file from a destructor, and the - caller removes the temporary directory as soon as the widths are back. - The result holds one entry per decay and flavour assignment:: - - {"incoming": 25, "outgoing": [5, -5], "width": 4.3e-3, "error": 1e-9} - - (signed pdg ids, width in GeV). The flavours of a subprocess (different - matrix elements, e.g. the u u~ g and d d~ g of z > q q~ g) are integrated - separately (MadgraphProcess.split_decay_flavors); the assignments of one - flavour share its matrix element (they are the same function of the - momenta, or they would not be grouped) and split its width evenly. A - decay the decaying particle is too light for has no entry. - """ - cwd = os.getcwd() - os.chdir(me_dir) - process = None - try: - card_path = os.path.join("Cards", "run_card.toml") - card = RunCardMG7(card_path, consistency=False) - card["generation"]["survey_target_precision"] = precision - card["generation"]["survey_max_iters"] = 30 - # one survey result per decay, not per group of mergeable ones - card["phasespace"]["merge_subprocesses"] = False - card["run"]["verbosity"] = "silent" - if nb_core: - card["run"]["cpu_thread_pool_size"] = int(nb_core) - card.write(card_path) - - # one survey result per flavour, not per subprocess - process = MadgraphProcess(extra_make_args=extra_make_args, - split_flavors=True) - if not process.is_decay: - raise ValueError("compute_decay_widths needs a decay (1 -> n) output") - process.survey() - - channel_status = process.event_generator.channel_status() - results = [] - offset = 0 - for subproc, phasespace in zip(process.subprocesses, process.phasespaces): - statuses = channel_status[offset:offset + len(phasespace.channels)] - offset += len(phasespace.channels) - width = sum(status.mean for status in statuses) - error = sum(status.error ** 2 for status in statuses) ** 0.5 - options = _decay_flavor_options(subproc.meta) - for option in options: - results.append({ - "incoming": option[0], - "outgoing": list(option[1:]), - "width": width / len(options), - "error": error / len(options), - }) - return results - finally: - if process is not None: - # StatusFile finishes info.json from its destructor: run it before - # the caller removes the directory (see init_event_dir) - _release_channel_generators(process) - del process - gc.collect() - os.chdir(cwd) - - -def _single_flavor_metas(meta: dict) -> list[dict]: - """The subprocess ``meta`` as one subprocess per flavour (see - MadgraphProcess.split_decay_flavors). - - Each copy keeps the flavour (its matrix-element index and its assignments) - and, in every channel, the diagrams that flavour is active in, renumbered - to its only flavour, 0. A channel left without diagrams is dropped, and a - flavour left without channels (no diagram at all) has no subprocess. A - subprocess of one flavour is returned as it is. - """ - flavors = meta["flavors"] - if len(flavors) == 1: - return [meta] - singles = [] - for position, flavor in enumerate(flavors): - channels = [] - for channel in meta["channels"]: - diagrams = [ - dict(diagram, active_flavors=[0]) - for diagram in channel["diagrams"] - if position in diagram["active_flavors"] - ] - if diagrams: - channels.append(dict(channel, diagrams=diagrams)) - if channels: - singles.append(dict(meta, flavors=[flavor], channels=channels)) - return singles - - -def _decay_flavor_options(meta: dict) -> list[list[int]]: - """The flavour assignments of a subprocess of compute_decay_widths, which - share its width evenly. - - That is right for the assignments of one flavour, which have one matrix - element (d d~ and s s~ of a massless s, the three neutrinos). Different - flavours have different matrix elements (the u u~ and d d~ of a - flavour-grouped subprocess) and the survey of a subprocess only gives - their sum: compute_decay_widths splits them first (split_flavors), and a - subprocess still holding several is an error, not an even split that - would get the branching ratios wrong. - """ - flavors = meta["flavors"] - if len(flavors) != 1: - raise ValueError( - "compute_decay_widths: the subprocess %s has %d flavours with " - "different matrix elements, the survey cannot give the width of " - "each" % (meta.get("path", "?"), len(flavors)) - ) - return [list(option) for option in flavors[0]["options"]] - - -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() - if getattr(process, "auto_temp_path", False): - shutil.rmtree(process.temp_path, ignore_errors=True) - - -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 madgraph.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 madgraph.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") - parser.add_argument( - "--decay-widths", metavar="JSON", default=None, - help="only integrate the 1 -> n decays of this directory and write " - "their widths (GeV) as JSON, see compute_decay_widths") - parser.add_argument( - "--precision", type=float, default=0.01, - help="relative precision aimed at for every channel (--decay-widths)") - parser.add_argument( - "--nb-core", type=int, default=None, - help="number of threads and compilation jobs (--decay-widths)") - parser.add_argument( - "--make-arg", action="append", default=[], - help="extra argument of the matrix element build (--decay-widths)") - args = parser.parse_args() - - if args.decay_widths: - results = compute_decay_widths( - ".", precision=args.precision, nb_core=args.nb_core, - extra_make_args=args.make_arg) - with open(args.decay_widths, "w") as f: - json.dump(results, f) - return - - MG7Cmd().run_cmd("generate_events" if args.ask_edit_cards else "generate_events -f") +from madgraph.mg7.launch import main # noqa: F401 diff --git a/madgraph/iolibs/template_files/mg7/npy_to_lhe.py b/madgraph/iolibs/template_files/mg7/npy_to_lhe.py index df9c62c2f5..8f194255cc 100644 --- a/madgraph/iolibs/template_files/mg7/npy_to_lhe.py +++ b/madgraph/iolibs/template_files/mg7/npy_to_lhe.py @@ -1,252 +1,3 @@ -"""Convert an mg7 numpy event file (``events.npy``) into an LHE file. +"""Redirect for the bin/ scripts of outputs made before the move to madgraph/mg7.""" -The mg7 output writes its events as npy by default. Anything that needs an LHE -file -- a shower, MadSpin, reweighting, an analysis -- normally switches the run -to the LHE format before the generation (``force_lhe_output_if_needed``). This -module covers the case where that request comes *after* the run: it rebuilds -the LHE file from what the npy run left behind. - -* ``output_format = "lhe_npy"``: the file already holds complete LHE events - (pdg ids, status, mothers, colour, helicity); they are only reformatted. -* ``output_format = "compact_npy"``: the file holds the momenta plus the - subprocess/diagram/colour/flavour/helicity indices the generation picked. - Those are completed exactly as the direct LHE output does it, with the - :class:`madspace.LHECompleter` the run saved next to the events - (``lhe_completer.json``). The completion draws random numbers (colour flow - among equivalent ones, resonance assignment), so the result is statistically - but not byte-identical to an ``output_format = "lhe"`` run with the same seed. - -In both cases the LHE header and ```` block (with the systematics -````) are taken verbatim from the ``header.lhe`` file the run writes -next to its npy events; the ``rwgt_`` columns become the per-event -```` block. The ```` reweighting-input block (``[systematics] -write_inputs``) is not reconstructed. - -Usage:: - - bin/npy_to_lhe [RUN_NAME | RUN_DIR | path/to/events.npy] [-o OUTPUT] [--seed N] -""" - -import argparse -import logging -import os -import re -import sys - -logger = logging.getLogger("madevent") - -COMPLETER_FILE = "lhe_completer.json" -HEADER_FILE = "header.lhe" -EVENT_FILE = "events.npy" - -# events converted per LHEFileWriter.write_string call -_CHUNK = 10000 - - -def _madspace(): - from madgraph.iolibs.template_files.mg7.bootstrap import ensure_madspace - ensure_madspace(interactive=False) - import madspace - return madspace - - -def _read_header(run_path): - """The header.lhe text without its closing tag, ready to be followed by - the events.""" - with open(os.path.join(run_path, HEADER_FILE)) as f: - text = f.read() - end = text.rfind("") - return text if end < 0 else text[:end] - - -def _seed(header, seed): - if seed is not None: - return int(seed) - match = re.search(r"\s*(-?\d+)\s*", header) - return int(match.group(1)) if match else 0 - - -# ---------------------------------------------------------------------------- -# conversion -# ---------------------------------------------------------------------------- -def resolve_event_file(path, me_dir=None) -> str: - """``path`` may be an npy file, a run directory or a run name (looked up - under ``/Events``).""" - candidates = [path, os.path.join(path, EVENT_FILE)] - if me_dir is not None: - candidates.append(os.path.join(me_dir, "Events", path, EVENT_FILE)) - for candidate in candidates: - if os.path.isfile(candidate): - return os.path.abspath(candidate) - raise FileNotFoundError("no %s found for '%s'" % (EVENT_FILE, path)) - - -def can_convert(run_path) -> bool: - """True when ``run_path`` holds npy events this module can turn into LHE.""" - return (os.path.isfile(os.path.join(run_path, EVENT_FILE)) - and os.path.isfile(os.path.join(run_path, HEADER_FILE))) - - -def _particle_count(names): - count = 0 - while "part%d_energy" % (count + 1) in names: - count += 1 - return count - - -def _compact_events(ms, events, completer, columns, rng): - names = events.dtype.names - nparticles = _particle_count(names) - momenta = [tuple(events["part%d_%s" % (i + 1, c)] - for c in ("energy", "px", "py", "pz")) - for i in range(nparticles)] - rwgt = [events[c] for c in columns] - for n in range(len(events)): - particles = [] - for energy, px, py, pz in momenta: - if energy[n] == 0.: - break - particles.append(ms.LHEParticle(px=px[n], py=py[n], pz=pz[n], - energy=energy[n])) - event = ms.LHEEvent(weight=events["weight"][n], - scale=events["ren_scale"][n], - alpha_qcd=events["alpha_qcd"][n], - particles=particles) - completer.complete_event_data( - event, int(events["subprocess_index"][n]), - int(events["diagram_index"][n]), int(events["color_index"][n]), - int(events["flavor_index"][n]), int(events["helicity_index"][n]), - rng) - yield event, [float(w[n]) for w in rwgt] - - -_LHE_PARTICLE_FIELDS = ("pdg_id", "status_code", "mother1", "mother2", "color", - "anti_color", "px", "py", "pz", "energy", "mass", - "lifetime", "spin") - - -def _lhe_npy_events(ms, events, columns): - names = events.dtype.names - count = 0 - while "part%d_pdg_id" % (count + 1) in names: - count += 1 - fields = [[events["part%d_%s" % (i + 1, f)] for f in _LHE_PARTICLE_FIELDS] - for i in range(count)] - rwgt = [events[c] for c in columns] - for n in range(len(events)): - particles = [] - for values in fields: - if values[0][n] == 0: - break - particles.append(ms.LHEParticle( - *[v[n].item() for v in values])) - event = ms.LHEEvent(int(events["process_id"][n]), - events["weight"][n], events["scale"][n], - events["alpha_qed"][n], events["alpha_qcd"][n], - particles) - yield event, [float(w[n]) for w in rwgt] - - -def convert(path, output=None, seed=None, compress=True, me_dir=None) -> str: - """Write the LHE file for the npy events at ``path`` (npy file, run - directory or run name) and return its path. By default it is written as - ``events.lhe.gz`` next to the npy file.""" - import numpy as np - from madgraph.various import misc - - ms = _madspace() - npy_path = resolve_event_file(path, me_dir) - run_path = os.path.dirname(npy_path) - # A standalone conversion starts in a fresh process, so it has not seen - # the gzip configuration installed by the event-generation launcher. - card_root = me_dir or os.path.dirname(os.path.dirname(run_path)) - card_path = os.path.join(card_root, "Cards", "run_card.toml") - if os.path.isfile(card_path): - from madgraph.various.banner import RunCardMG7 - misc.configure_gzip(RunCardMG7(card_path).gzip_configuration()) - if not os.path.isfile(os.path.join(run_path, HEADER_FILE)): - raise FileNotFoundError( - "%s is missing: the run predates the npy->LHE conversion support " - "and cannot be converted." % os.path.join(run_path, HEADER_FILE)) - header = _read_header(run_path) - - events = np.load(npy_path, mmap_mode="r") - names = events.dtype.names - columns = [n for n in names if re.fullmatch(r"rwgt_\d+", n)] - ids = [int(c.split("_", 1)[1]) for c in columns] - if "diagram_index" in names: - completer_path = os.path.join(run_path, COMPLETER_FILE) - if not os.path.isfile(completer_path): - raise FileNotFoundError("%s is missing" % completer_path) - completer = ms.LHECompleter.load(completer_path) - rng = ms.MixMaxRandom(_seed(header, seed)) - stream = lambda chunk: _compact_events(ms, chunk, completer, columns, rng) - elif "part1_pdg_id" in names: - stream = lambda chunk: _lhe_npy_events(ms, chunk, columns) - else: - raise ValueError("%s is not an mg7 event file" % npy_path) - - if output is None: - output = os.path.join(run_path, "events.lhe") - if output.endswith(".gz"): - output, compress = output[:-3], True - logger.info("converting %s (%d events) to LHE", npy_path, len(events)) - with open(output, "w") as out: - out.write(header) - for start in range(0, len(events), _CHUNK): - chunk = np.asarray(events[start:start + _CHUNK]) - text = [] - for event, weights in stream(chunk): - if ids: - event.rwgt_ids = ids - event.rwgt = weights - text.append(event.format()) - out.write("".join(text)) - out.write("\n") - if compress: - misc.gzip(output) - output += ".gz" - logger.info("LHE events written to %s", output) - return output - - -def main(argv=None, me_dir=None) -> None: - """Command-line entry point; run names are looked up under - ``/Events`` (default: the working directory).""" - parser = argparse.ArgumentParser( - description="Convert an mg7 events.npy file into an LHE file.") - parser.add_argument("path", nargs="?", default=None, - help="run name, run directory or events.npy file " - "(default: the most recent run with npy events " - "and no LHE file)") - parser.add_argument("-o", "--output", default=None, - help="output file (default: events.lhe.gz next to the " - "npy file)") - parser.add_argument("--seed", type=int, default=None, - help="seed of the LHE completion (default: the run seed)") - parser.add_argument("--no-gzip", action="store_false", dest="compress", - help="do not compress the output") - args = parser.parse_args(argv) - logging.basicConfig(level=logging.INFO, format="%(message)s") - - me_dir = me_dir or os.getcwd() - path = args.path - if path is None: - events_dir = os.path.join(me_dir, "Events") - runs = [os.path.join(events_dir, d) for d in - (os.listdir(events_dir) if os.path.isdir(events_dir) else [])] - runs = [d for d in runs if can_convert(d) and not any( - os.path.exists(os.path.join(d, n)) - for n in ("events.lhe", "events.lhe.gz"))] - if not runs: - sys.exit("no run with unconverted npy events in %s" % events_dir) - path = max(runs, key=lambda d: os.path.getmtime( - os.path.join(d, EVENT_FILE))) - try: - convert(path, args.output, args.seed, args.compress, me_dir) - except (FileNotFoundError, ValueError) as error: - sys.exit(str(error)) - - -if __name__ == "__main__": - main() +from madgraph.mg7.npy_to_lhe import main # noqa: F401 diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 42bcc709af..3151b5b032 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -46,7 +46,7 @@ include_source = %(gridpack.include_source)s include_madspace = %(gridpack.include_madspace)s include_madspace_source = %(gridpack.include_madspace_source)s compress = %(gridpack.compress)s -# end of a run creating a gridpack; options: regular, minimal (no events), fix_max_weight (no events) +# end of a run creating a gridpack; options: regular, minimal (no events), fix_max_weight (no events, adds bin/event_stream.py) run_mode = %(gridpack.run_mode)s # temporary npy files; "" uses the output directory, "auto" a new temporary directory temp_output_dir = %(gridpack.temp_output_dir)s @@ -156,6 +156,27 @@ bias = %(custom_functions.bias)s # events are generated with |w| b and written w # parameters declared by a function, e.g. cut_parameters.pt_min = 30.0 $custom_function_parameters +[surrogate] +# a fast surrogate of the matrix element in the integrand: a .so library, which can +# contain {subprocess} and {device}, or a Python file defining +# custom_surrogate(momenta, pdg_ids, alpha_s); "" disables it +function = %(surrogate.function)s +# unweight or reweight the events with the full matrix element in a second pass +mode = %(surrogate.mode)s +# unweighting with the surrogate (unweight mode); [generation] sets the second pass +first_pass_max_overweight_truncation = %(surrogate.first_pass_max_overweight_truncation)s +first_pass_freeze_max_weight_after = %(surrogate.first_pass_freeze_max_weight_after)s +param_card = %(surrogate.param_card)s # for a .so surrogate, "" uses Cards/param_card.dat + +[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/mg7/__init__.py b/madgraph/mg7/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/madgraph/iolibs/template_files/mg7/bootstrap.py b/madgraph/mg7/bootstrap.py similarity index 70% rename from madgraph/iolibs/template_files/mg7/bootstrap.py rename to madgraph/mg7/bootstrap.py index 416415d5bf..15e62463df 100644 --- a/madgraph/iolibs/template_files/mg7/bootstrap.py +++ b/madgraph/mg7/bootstrap.py @@ -19,9 +19,11 @@ safe to do lazily. """ -import importlib +import importlib.metadata +import importlib.util import json import os +import re import subprocess import sys from pathlib import Path @@ -43,9 +45,39 @@ def madspace_is_installed() -> bool: return (INSTALL_DIR / "madspace").is_dir() +def _pyamplicol_package(): + """The directory of the pyamplicol package that ``import pyamplicol`` would + load, or None. Found, not imported (Symbolica would take its per-user lock).""" + spec = importlib.util.find_spec("pyamplicol") + if spec is None or not spec.origin: + return None # e.g. the source directory as a namespace package + package = Path(spec.origin).parent + if (package.resolve().is_relative_to(PYAMPLICOL_DIR.resolve()) + and not package.resolve().is_relative_to(PYAMPLICOL_INSTALL_DIR.resolve())): + return None + return package + + def pyamplicol_is_installed() -> bool: - """True when a compiled pyamplicol is already available to import.""" - return (PYAMPLICOL_INSTALL_DIR / "pyamplicol").is_dir() + """True when a compiled pyamplicol is already available to import, in + pyamplicol/install or system-wide (``install pyamplicol --system``).""" + return (PYAMPLICOL_INSTALL_DIR / "pyamplicol").is_dir() or _pyamplicol_package() is not None + + +def installed_pyamplicol() -> dict: + """The version and source hash of the pyamplicol that would be imported + (None where unknown), read without importing it.""" + package = _pyamplicol_package() + if package is None: + return {"version": None, "source_hash": None} + source_hash = None + hash_file = package / "_source_hash.py" + if hash_file.is_file(): + match = re.search(r'SOURCE_HASH\s*=\s*"([0-9a-f]*)"', hash_file.read_text()) + source_hash = match.group(1) if match else None + dist = next(importlib.metadata.distributions( + name="pyamplicol", path=[str(package.parent)]), None) + return {"version": dist.version if dist else None, "source_hash": source_hash} def _run_installer(name, source_dir, interactive, jobs) -> None: @@ -113,29 +145,47 @@ def ensure_pyamplicol(interactive=None, jobs=None) -> None: sys.path.insert(0, str(PYAMPLICOL_INSTALL_DIR)) -def check_pyamplicol_source_hash() -> bool: - """Warn if the installed pyamplicol was built from other sources. +def check_pyamplicol_source_hash() -> dict: + """Fail unless the installed pyamplicol was built from the sources in + pyamplicol/, and return :func:`installed_pyamplicol`. - Returns whether the hashes match. Must run after :func:`ensure_pyamplicol`. + Must run after :func:`ensure_pyamplicol`. """ - try: - installed = importlib.import_module("pyamplicol._source_hash").SOURCE_HASH - except ImportError: - installed = None - source = subprocess.run( + installed = installed_pyamplicol() + result = subprocess.run( [sys.executable, str(PYAMPLICOL_DIR / "source_hash.py")], - capture_output=True, text=True, check=True, - ).stdout.strip() - if installed == source: - return True - print() - print( - "\033[1m\033[31mWARNING\033[39m: pyamplicol source and installed package " - "are not compatible (source hash mismatch) — consider recompiling " - "pyamplicol (e.g. `install pyamplicol -y`)\033[0m" + capture_output=True, text=True, ) - print() - return False + if result.returncode != 0: + raise RuntimeError( + "cannot compute the hash of the pyamplicol sources:\n" + + result.stdout + result.stderr + ) + if installed["source_hash"] != result.stdout.strip(): + raise RuntimeError( + "the installed pyamplicol was not built from the sources in " + f"{PYAMPLICOL_DIR} (source hash mismatch): run `install pyamplicol` " + "to rebuild it" + ) + return installed + + +def check_pyamplicol_build(record_path) -> None: + """Exit with an error unless the pyamplicol that would be imported is the + one recorded in record_path by the output. Run by SubProcesses/check-pyamplicol + before the [amplicol] libraries are compiled.""" + with open(record_path) as stream: + recorded = json.load(stream) + installed = installed_pyamplicol() + if installed != {key: recorded.get(key) for key in installed}: + sys.exit( + f"The [amplicol] matrix elements of this output were generated with " + f"pyamplicol {recorded.get('version')} (source hash " + f"{recorded.get('source_hash')}), but the installed one is " + f"{installed['version']} (source hash {installed['source_hash']}): " + "output the process again, or reinstall the pyamplicol it was " + "generated with (`install pyamplicol`)." + ) def check_pyamplicol_license(pythonpath=None) -> None: diff --git a/madgraph/mg7/cli.py b/madgraph/mg7/cli.py new file mode 100644 index 0000000000..a6978818c8 --- /dev/null +++ b/madgraph/mg7/cli.py @@ -0,0 +1,589 @@ +"""The MG7> command line interface and its card editing.""" + +import argparse +import json +import os +import re +import resource + +from madgraph.mg7.common import logger +from madgraph.interface.extended_cmd import Cmd +from madgraph.various.banner import RunCardMG7 +from madgraph.mg7 import npy_to_lhe +from madgraph.mg7.driver import compute_decay_widths, force_lhe_output_if_needed, run_generation +from madgraph.mg7.model_widths import compute_auto_widths, load_process_model +from madgraph.mg7.options import lhapdf_index_paths, load_mg5_options +from madgraph.mg7.postprocess import _off, _setup_logging + +class MG7Cmd(Cmd): + """The mg7 run interface: the command object that drives a launch. + + 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}) + # and so does its error policy: with crash_on_error a failing + # launch has to stop the script with a non-zero exit code, as a + # madevent launch does, instead of returning to the prompt. + if 'crash_on_error' in options: + self.options['crash_on_error'] = options['crash_on_error'] + 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 check_card_consistency(self): + # madevent settings given per beam ("set fixed_fac_scale1 ...") + # that the mg7 card could not follow: said once, at the end, so + # "set fixed_fac_scale1 T" then "set fixed_fac_scale2 T" stays + # quiet. Collected first: lo_set_warnings also forgets them, which + # must happen even if the generic checks below raise. + warnings = [] + if isinstance(getattr(self, "run_card", None), RunCardMG7): + warnings = self.run_card.lo_set_warnings() + try: + super().check_card_consistency() + finally: + for text in warnings: + logger.warning(text) + + 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 + + # legacy madevent cut names ("set ptj 20", "set mmll 200", + # "set etaj -1", ...) -> the [cuts] entry the LO->MG7 + # run_card conversion maps them to (_LO_CUT_MAP), with the + # same "no cut" sentinels (minimum <= 0, maximum < 0). + if rest and nlow in run_card._LO_CUT_MAP: + if rest.lower().startswith("scan"): + logger.warning("ignoring 'set %s %s': the [cuts] entries of " + "the mg7 run_card.toml cannot be scanned", + nlow, rest) + return + try: + cut, bound, val = run_card.set_lo_cut( + nlow, run_card.evaluate(rest, masses)) + except _banner_mod.InvalidCmd: + logger.warning("ignoring 'set %s %s': not a number", nlow, rest) + return + if val is None: + logger.info("set %s (mg7 cut %s.%s) of the run_card.toml: " + "no cut", nlow, cut, bound) + else: + logger.info("set %s (mg7 cut %s.%s) of the run_card.toml to %s", + nlow, cut, bound, val) + self.modified_card.add("run") + return + + # madevent parameters that are not a plain rename: the + # value is translated ("set dynamical_scale_choice 3") or + # several of them make one mg7 entry (fixed_fac_scale1 + + # fixed_fac_scale2 -> beam.fixed_fact_scale). mg7 values + # ("set sde_strategy denominators") come back as None and + # go on to the generic editor. + if rest and nlow in run_card._LO_SET_PARAMS: + try: + done = run_card.set_lo_param(nlow, rest, masses) + except _banner_mod.InvalidCmd as error: + logger.warning("ignoring 'set %s %s': %s", nlow, rest, error) + return + if done is not None: + changes, warnings = done + for key, val in changes.items(): + if val is None: + logger.info("set %s: removed %s from the run_card.toml", + nlow, key) + else: + logger.info("set %s (mg7 %s) of the run_card.toml to %s", + nlow, key, val) + for text in warnings: + logger.warning(text) + if changes: + self.modified_card.add("run") + elif not warnings: + logger.info("set %s %s: nothing to change in the run_card.toml", + nlow, rest) + 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 + + # madevent beam settings (ebeam, lpp, pdlabel, lhaid, per + # beam or for both) onto the per-beam [beam] entries + if rest and nlow in run_card.lo_beam_params: + data_paths = lhapdf_index_paths() if nlow == "lhaid" else () + try: + done = run_card.set_lo_beam_param(nlow, rest, masses, data_paths) + except _banner_mod.InvalidRunCard as error: + logger.warning("ignoring 'set %s %s': %s", name, rest, error) + return + logger.info("set %s %s: %s", nlow, rest, done) + self.modified_card.add("run") + return + + # quantities the card stores per beam: set both beams + if rest and nlow in ("e_cm", "beam.e_cm", "beam.pdf"): + value = rest if nlow == "beam.pdf" else run_card.evaluate(rest, masses) + run_card.set(nlow, value, user=True) + logger.info("set %s of the run_card.toml to %s (both beams)", + nlow, run_card["beam." + nlow.split(".")[-1]]) + 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") + parser.add_argument( + "--decay-widths", metavar="JSON", default=None, + help="only integrate the 1 -> n decays of this directory and write " + "their widths (GeV) as JSON, see compute_decay_widths") + parser.add_argument( + "--precision", type=float, default=0.01, + help="relative precision aimed at for every channel (--decay-widths)") + parser.add_argument( + "--nb-core", type=int, default=None, + help="number of threads and compilation jobs (--decay-widths)") + parser.add_argument( + "--make-arg", action="append", default=[], + help="extra argument of the matrix element build (--decay-widths)") + args = parser.parse_args() + + if args.decay_widths: + results = compute_decay_widths( + ".", precision=args.precision, nb_core=args.nb_core, + extra_make_args=args.make_arg) + with open(args.decay_widths, "w") as f: + json.dump(results, f) + return + + MG7Cmd().run_cmd("generate_events" if args.ask_edit_cards else "generate_events -f") diff --git a/madgraph/mg7/common.py b/madgraph/mg7/common.py new file mode 100644 index 0000000000..163dc5c3d2 --- /dev/null +++ b/madgraph/mg7/common.py @@ -0,0 +1,205 @@ +"""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.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 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 mlm_pdf_ladder(run_card, mlm_clustering, is_decay=False, leptonic=False): + """Whether an MLM sample's beam densities walk the clustering ladder, the + pdf half of madevent's rewgt: under scale_scheme "madevent", and under + "clustering_mean" with clustering_mean_pdf_ladder. The Python side of + EnergyScale::mlm_pdf_reweighting, which decides the same for the kernel; + the systematics then vary the ladder as madevent's do.""" + beam = run_card["beam"] + return bool( + mlm_clustering and not is_decay and not leptonic + and beam["pdf_reweighting"] + and (beam["scale_scheme"] == "madevent" + or beam["clustering_mean_pdf_ladder"])) + + +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/mg7/driver.py b/madgraph/mg7/driver.py new file mode 100644 index 0000000000..d23add81fc --- /dev/null +++ b/madgraph/mg7/driver.py @@ -0,0 +1,286 @@ +"""Runs, scans and the entry point of the mg7 launcher.""" + +import gc +import logging +import os +import re +import shutil + +from madgraph.mg7.common import logger +from madgraph.various.banner import RunCardMG7 +from madgraph.mg7.model_widths import compute_auto_widths +from madgraph.mg7.postprocess import run_lhe_postprocessing, run_selected_tools, selected_tools +from madgraph.mg7.madgraph_process import MadgraphProcess + +def compute_decay_widths( + me_dir: str = ".", precision: float = 0.01, nb_core: int | None = None, + extra_make_args=(), +) -> list[dict]: + """Integrate the decays of an mg7 decay directory and return their widths. + + This is the madspace counterpart of madevent's ``survey decay``: only the + survey runs (no unweighting, no events, no systematics), and it stops once + every channel is below ``precision`` relative error, like the + ``--precision_channel`` of ``compute_widths``. + + ``me_dir`` is an output directory of 1 -> n processes of one decaying + particle. It is meant to run in a process of its own (``bin/generate_events + --decay-widths=FILE``, which is how ``compute_widths`` calls it): the + madspace objects write their status file from a destructor, and the + caller removes the temporary directory as soon as the widths are back. + The result holds one entry per decay and flavour assignment:: + + {"incoming": 25, "outgoing": [5, -5], "width": 4.3e-3, "error": 1e-9} + + (signed pdg ids, width in GeV). The flavours of a subprocess (different + matrix elements, e.g. the u u~ g and d d~ g of z > q q~ g) are integrated + separately (MadgraphProcess.split_decay_flavors); the assignments of one + flavour share its matrix element (they are the same function of the + momenta, or they would not be grouped) and split its width evenly. A + decay the decaying particle is too light for has no entry. + """ + cwd = os.getcwd() + os.chdir(me_dir) + process = None + try: + card_path = os.path.join("Cards", "run_card.toml") + card = RunCardMG7(card_path, consistency=False) + card["generation"]["survey_target_precision"] = precision + card["generation"]["survey_max_iters"] = 30 + # one survey result per decay, not per group of mergeable ones + card["phasespace"]["merge_subprocesses"] = False + card["run"]["verbosity"] = "silent" + if nb_core: + card["run"]["cpu_thread_pool_size"] = int(nb_core) + card.write(card_path) + + # one survey result per flavour, not per subprocess + process = MadgraphProcess(extra_make_args=extra_make_args, + split_flavors=True) + if not process.is_decay: + raise ValueError("compute_decay_widths needs a decay (1 -> n) output") + process.survey() + + channel_status = process.event_generator.channel_status() + results = [] + offset = 0 + for subproc, phasespace in zip(process.subprocesses, process.phasespaces): + statuses = channel_status[offset:offset + len(phasespace.channels)] + offset += len(phasespace.channels) + width = sum(status.mean for status in statuses) + error = sum(status.error ** 2 for status in statuses) ** 0.5 + options = _decay_flavor_options(subproc.meta) + for option in options: + results.append({ + "incoming": option[0], + "outgoing": list(option[1:]), + "width": width / len(options), + "error": error / len(options), + }) + return results + finally: + if process is not None: + # StatusFile finishes info.json from its destructor: run it before + # the caller removes the directory (see init_event_dir) + _release_channel_generators(process) + del process + gc.collect() + os.chdir(cwd) + + +def _decay_flavor_options(meta: dict) -> list[list[int]]: + """The flavour assignments of a subprocess of compute_decay_widths, which + share its width evenly. + + That is right for the assignments of one flavour, which have one matrix + element (d d~ and s s~ of a massless s, the three neutrinos). Different + flavours have different matrix elements (the u u~ and d d~ of a + flavour-grouped subprocess) and the survey of a subprocess only gives + their sum: compute_decay_widths splits them first (split_flavors), and a + subprocess still holding several is an error, not an even split that + would get the branching ratios wrong. + """ + flavors = meta["flavors"] + if len(flavors) != 1: + raise ValueError( + "compute_decay_widths: the subprocess %s has %d flavours with " + "different matrix elements, the survey cannot give the width of " + "each" % (meta.get("path", "?"), len(flavors)) + ) + return [list(option) for option in flavors[0]["options"]] + + +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() + if getattr(process, "auto_temp_path", False): + shutil.rmtree(process.temp_path, ignore_errors=True) + + +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 madgraph.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 madgraph.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/hwu_output.py b/madgraph/mg7/hwu_output.py similarity index 100% rename from madgraph/iolibs/template_files/mg7/hwu_output.py rename to madgraph/mg7/hwu_output.py diff --git a/madgraph/mg7/init_custom_func.py b/madgraph/mg7/init_custom_func.py new file mode 100644 index 0000000000..3133851e6d --- /dev/null +++ b/madgraph/mg7/init_custom_func.py @@ -0,0 +1,266 @@ +"""Create a skeleton for a custom cut, scale, PDF or bias function. + +The function is written for a single phase-space point in C, C++, Fortran, +CUDA or HIP. A generated UMAMI wrapper calls it on batches of events, so that +madspace can load it like a matrix element. The files go into +``custom_/`` in the process directory, and ``make`` there builds +``lib/custom_.so``. The run card option +``[custom_functions] `` takes the path of that library. + +Usage:: + + bin/init_custom_func [-f {cut,scale,pdf,bias}] [-l {c,cpp,fortran,cuda,hip}] + [-y] [--force] [--build | --no-build] + +Options that are not given are asked interactively, unless ``-y`` is passed. +""" + +import argparse +import logging +import os +import shutil +import subprocess +import sys + +logger = logging.getLogger("madevent") + +_TEMPLATE_FILES = os.path.join( + os.path.dirname(os.path.dirname(os.path.realpath(__file__))), "iolibs", "template_files" +) +TEMPLATE_DIR = os.path.join(_TEMPLATE_FILES, "mg7", "custom_functions") +UMAMI_HEADER = os.path.join(_TEMPLATE_FILES, "madmatrix", "umami.h") + +FUNCTIONS = { + "cut": (["UMAMI_OUT_CUT"], "a cut, true keeps the event and false rejects it"), + "scale": ( + ["UMAMI_OUT_REN_SCALE", "UMAMI_OUT_FACT_SCALE_1", "UMAMI_OUT_FACT_SCALE_2"], + "the renormalization and factorization scales", + ), + "pdf": (["UMAMI_OUT_PDF"], "a parton density x f(x, Q)"), + "bias": (["UMAMI_OUT_BIAS"], "an event bias for the unweighting"), +} +# example parameter of parameters.h: (name, default, run card value) +EXAMPLE_PARAMETERS = { + "cut": ("pt_min", "20.", "30.0"), + "scale": ("ht_fraction", "0.5", "0.25"), + "pdf": ("gluon_norm", "3.", "2.5"), + "bias": ("pt0", "100.", "200.0"), +} +# language: (body template extension, description) +LANGUAGES = { + "c": ("c", "C"), + "cpp": ("cpp", "C++"), + "fortran": ("F90", "Fortran"), + "cuda": ("cu", "CUDA, for runs on NVIDIA GPUs"), + "hip": ("cu", "HIP, for runs on AMD GPUs"), + "python-numpy": ("py", "Python with numpy, for CPU runs, no compilation"), + "python-torch": ("py", "Python with PyTorch, for CPU or GPU runs, no compilation"), +} + +_NONINTERACTIVE = (not sys.stdin.isatty()) or any(a in ("-y", "--yes") for a in sys.argv[1:]) + + +def _ask(question, default, choices): + """Ask through MG5's cmd.ask when available, else input().""" + if _NONINTERACTIVE: + return default + try: + from madgraph.interface.extended_cmd import Cmd + + answer = Cmd().ask(question, default, choices=list(choices), timeout=0) + return default if answer is None else str(answer) + except Exception: + pass + try: + answer = input(f"{question} [{default}]: ").strip() + except EOFError: + return default + return answer or default + + +def ask_choice(title, options, default): + """Numbered menu over ``options`` (name -> description), returns a name.""" + names = list(options) + lines = [title] + width = max(len(name) for name in names) + for i, name in enumerate(names, 1): + lines.append(f" {i}. {name:{width}s} {options[name]}") + print("\n".join(lines)) + choices = names + [str(i) for i in range(1, len(names) + 1)] + while True: + answer = _ask("Enter a number or a name", default, choices).strip().lower() + if answer in names: + return answer + if answer.isdigit() and 1 <= int(answer) <= len(names): + return names[int(answer) - 1] + print(f" Please enter one of {', '.join(names)} or 1-{len(names)}.") + + +def ask_yes_no(question, default=True): + default_s = "y" if default else "n" + while True: + answer = _ask(question, default_s, ["y", "n"]).strip().lower() + if answer in ("y", "yes"): + return True + if answer in ("n", "no"): + return False + print(" Please enter 'y' or 'n'.") + + +def build_parser(): + parser = argparse.ArgumentParser( + prog="bin/init_custom_func", + description=( + "Create a skeleton for a custom cut, scale, PDF or bias function in " + "custom_/. Compiled languages are built into " + "lib/custom_.so, Python functions need no build." + ), + ) + parser.add_argument("-f", "--function", choices=list(FUNCTIONS)) + parser.add_argument("-l", "--language", choices=list(LANGUAGES)) + parser.add_argument( + "-y", "--yes", action="store_true", + help="do not ask, use the defaults (cut, cpp) for options that are not given", + ) + parser.add_argument( + "--force", action="store_true", help="overwrite an existing custom_ directory" + ) + build = parser.add_mutually_exclusive_group() + build.add_argument("--build", dest="build", action="store_true", default=None, + help="run make after creating the files (default)") + build.add_argument("--no-build", dest="build", action="store_false") + run_card = parser.add_mutually_exclusive_group() + run_card.add_argument("--set-run-card", dest="set_run_card", action="store_true", + default=None, help="enable the function in Cards/run_card.toml (default)") + run_card.add_argument("--no-set-run-card", dest="set_run_card", action="store_false") + return parser + + +def is_python(language): + return language.startswith("python-") + + +def function_path(function, language): + """The run card path of the function, relative to the process directory.""" + if is_python(language): + return os.path.join(f"custom_{function}", f"custom_{function}.py") + return os.path.join("lib", f"custom_{function}.so") + + +def write_files(out_dir, function, language): + """Write the function body and, for compiled languages, the wrapper, + umami.h and the Makefile.""" + extension = LANGUAGES[language][0] + body_file = f"custom_{function}.{extension}" + if is_python(language): + framework = language.split("-", 1)[1] + os.makedirs(out_dir, exist_ok=True) + shutil.copy( + os.path.join(TEMPLATE_DIR, f"{function}_{framework}.py"), + os.path.join(out_dir, body_file), + ) + return body_file + values = { + "kind": function, + "function_name": f"custom_{function}", + "output_keys": ", ".join(FUNCTIONS[function][0]), + "is_pdf": int(function == "pdf"), + "is_scale": int(function == "scale"), + "return_type": "bool" if function == "cut" else "double", + "body_file": body_file, + "language": language, + "example": EXAMPLE_PARAMETERS[function][0], + "example_default": EXAMPLE_PARAMETERS[function][1], + "example_value": EXAMPLE_PARAMETERS[function][2], + } + os.makedirs(out_dir, exist_ok=True) + shutil.copy(os.path.join(TEMPLATE_DIR, f"{function}.{extension}"), os.path.join(out_dir, body_file)) + shutil.copy(UMAMI_HEADER, os.path.join(out_dir, "umami.h")) + wrapper = "wrapper_gpu.cu" if language in ("cuda", "hip") else "wrapper_cpu.c" + for template in (wrapper, "Makefile", "parameters.h"): + with open(os.path.join(TEMPLATE_DIR, template)) as f: + text = f.read() % values + with open(os.path.join(out_dir, template), "w") as f: + f.write(text) + return body_file + + +def build(out_dir): + """Run make in ``out_dir``, logging to compile.log. Returns True on success.""" + log_path = os.path.join(out_dir, "compile.log") + with open(log_path, "w") as log: + result = subprocess.run(["make"], cwd=out_dir, stdout=log, stderr=subprocess.STDOUT) + if result.returncode != 0: + with open(log_path) as log: + print(log.read()) + print(f"Building failed, see {log_path}") + return False + return True + + +def set_run_card(me_dir, function, lib_path): + from madgraph.various.banner import RunCardMG7 + + card_path = os.path.join(me_dir, "Cards", "run_card.toml") + run_card = RunCardMG7(card_path) + run_card.set(f"custom_functions.{function}", lib_path) + run_card.write(card_path) + + +def main(me_dir=None, argv=None): + global _NONINTERACTIVE + args = build_parser().parse_args(argv) + if args.yes: + _NONINTERACTIVE = True + me_dir = me_dir or os.getcwd() + + function = args.function or ask_choice( + "Which custom function do you want to create?", + {name: desc for name, (_, desc) in FUNCTIONS.items()}, + "cut", + ) + language = args.language or ask_choice( + "Which language do you want to write it in?", + {name: desc for name, (_, desc) in LANGUAGES.items()}, + "cpp", + ) + + out_dir = os.path.join(me_dir, f"custom_{function}") + if os.path.exists(out_dir): + if not (args.force or (not _NONINTERACTIVE and ask_yes_no( + f"{out_dir} already exists. Overwrite it?", False))): + print(f"{out_dir} already exists, use --force to overwrite it.") + return 1 + shutil.rmtree(out_dir) + + body_file = write_files(out_dir, function, language) + lib_path = function_path(function, language) + print(f"Created {os.path.relpath(out_dir, me_dir)}/. Write your function in " + f"{os.path.relpath(os.path.join(out_dir, body_file), me_dir)}.") + + do_build = args.build + if is_python(language): + if do_build: + print("Python functions need no build.") + do_build = False + elif do_build is None: + do_build = ask_yes_no("Build the default function now?", True) + if do_build and not build(out_dir): + return 1 + + do_set = args.set_run_card + if do_set is None: + do_set = ask_yes_no(f"Set custom_functions.{function} = \"{lib_path}\" in the run card?", True) + if do_set: + set_run_card(me_dir, function, lib_path) + print(f"Enabled it in Cards/run_card.toml: custom_functions.{function} = \"{lib_path}\".") + else: + print(f"To enable it, set custom_functions.{function} = \"{lib_path}\" in Cards/run_card.toml.") + if not do_build and not is_python(language): + print(f"Build it with `make -C {os.path.relpath(out_dir, me_dir)}`. " + "bin/generate_events also rebuilds it before a run.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/madgraph/mg7/launch.py b/madgraph/mg7/launch.py new file mode 100644 index 0000000000..e3e69dfb59 --- /dev/null +++ b/madgraph/mg7/launch.py @@ -0,0 +1,78 @@ +"""Entry point of the mg7 launcher, re-exporting its modules.""" + +from madgraph.mg7.common import * # noqa: F401,F403 +from madgraph.mg7.common import ms, logger # noqa: F401 +from madgraph.mg7.observables import ( # noqa: F401 + CutItem, + HistItem, + _MERGED_PID_REPRESENTATIVE, + clean_pids, + _SIGNED_OBSERVABLES, + _flavor_options, + observable_pids, + build_observable, + decay_products, + pid_is_qcd, +) +from madgraph.mg7.subprocess_madmatrix import ( # noqa: F401 + Channel, + PhaseSpace, + MultiChannelData, + build_topologies, + build_multi_channel_data, + MadgraphSubprocess, +) +from madgraph.mg7.madgraph_process import ( # noqa: F401 + MadgraphProcess, + PCH_MIN_SUBPROCESSES, + _single_flavor_metas, +) +from madgraph.mg7.options import ( # noqa: F401 + _ROOTED_OPTIONS, + load_mg5_options, + lhapdf_paths, + lhapdf_index_paths, +) +from madgraph.mg7.model_widths import ( # noqa: F401 + _model_hash_warned, + _model_cache, + read_stored_model, + _proc_characteristic, + load_process_model, + compute_auto_widths, +) +from madgraph.mg7.postprocess import ( # noqa: F401 + _find_event_file, + _find_or_convert_event_file, + _report_failure, + _off, + _has_effective_handler, + _TOOL_LOGGING_READY, + _setup_logging, + selected_tools, + run_selected_tools, + _add_time_of_flight, + _lhapdf_config_path, + _run_systematics, + run_lhe_postprocessing, +) +from madgraph.mg7.driver import ( # noqa: F401 + compute_decay_widths, + _decay_flavor_options, + _release_channel_generators, + run_single, + detect_run_scan, + detect_param_scan, + run_scan, + run_generation, + _lhe_needed_by, + force_lhe_output_if_needed, +) +from madgraph.mg7.cli import ( # noqa: F401 + MG7Cmd, + build_selector_cmd, + ask_edit_cards, + prune_unselected_tool_cards, + _raise_open_file_limit, + main, +) diff --git a/madgraph/mg7/madgraph_process.py b/madgraph/mg7/madgraph_process.py new file mode 100644 index 0000000000..3b9107b3fe --- /dev/null +++ b/madgraph/mg7/madgraph_process.py @@ -0,0 +1,2306 @@ +"""The run of one process directory: cards, integration, events and output.""" + +import glob +import json +import math +import os +import shutil +import subprocess +import tempfile +import time + +from madgraph.mg7.common import backend_of, fptype_of, interference_helicities_summed, lhe_helicities, lhe_weight_info, logger, me_parameters, mlm_pdf_ladder, ms, resolve_auto_backend, resolve_seed, resolve_verbosity +from madgraph.various import misc +from madgraph.various.banner import RunCardMG7 +from madgraph.models.check_param_card import ParamCard +from madgraph.models.model_reader import Alphas_Runner +from madgraph.mg7 import hwu_output +from madgraph.mg7 import npy_to_lhe +from madgraph.mg7 import options +from madgraph.mg7 import plots +from madgraph.mg7 import systematics_summary +from madgraph.mg7.bootstrap import INSTALL_DIR as _INSTALL_DIR +from madgraph.mg7.bootstrap import MADSPACE_DIR as _MADSPACE_DIR +from madgraph.iolibs.template_files.mg7.custom_function_loader import load_custom_function +from madgraph.mg7.observables import CutItem, HistItem, build_observable, clean_pids, observable_pids +from madgraph.mg7.options import lhapdf_paths +from madgraph.mg7.subprocess_amplicol import AmplicolSubprocess, amplicol_subproc_args, is_amplicol +from madgraph.mg7.subprocess_madmatrix import MadgraphSubprocess, PhaseSpace, build_multi_channel_data + +# From this many subprocess directories to compile, the standard headers are +# precompiled once ('make PCH=1', see madmatrix.mk): each file of a small process +# spends most of its time parsing them. Below, building the precompiled header +# costs about what it saves. +PCH_MIN_SUBPROCESSES = 4 + + +class MadgraphProcess: + def __init__(self, extra_make_args=(), split_flavors=False): + # appended to the make command that builds the matrix elements + self.extra_make_args = list(extra_make_args) + # one subprocess per flavour of a decay (split_decay_flavors) + self.split_flavors = split_flavors + 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")) + misc.configure_gzip(self.run_card.gzip_configuration()) + # 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) + merge_subprocesses = self.run_card["phasespace"]["merge_subprocesses"] + if merge_subprocesses and self.run_card["beam"]["dynamical_scale_choice"] == "mlm": + # A merged subprocess bundles matrix elements whose legs differ - + # q q~ > l l g next to q g > l l q - while the MLM clustering is + # built once per subprocess from a single assignment of flavours to + # legs. It would cluster the gluon-initiated elements as if their + # beam were a quark, and pick the history per diagram from the + # weights of one element only. + logger.warning( + "merge_subprocesses is not supported with dynamical_scale_choice " + "= mlm; the subprocesses are kept separate" + ) + merge_subprocesses = False + if 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 + single = [meta["outgoing"] for meta in self.subprocess_data + if len(clean_pids(meta["outgoing"])) == 1] + if single and self.run_card["beam"]["leptonic"]: + raise ValueError( + f"2 -> 1 processes ({single[0]} ...) need beams with a PDF " + "(leptonic = false in [beam]): with fixed beam energies the " + "one-particle phase space is a delta function in s_hat" + ) + 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() + if getattr(self, "split_flavors", False): + self.split_decay_flavors() + + def split_decay_flavors(self) -> None: + """Integrate every flavour of a decay subprocess on its own. + + A subprocess can hold several flavours with different matrix elements + (z > d d~ g and z > u u~ g in one P1_z_QQxg): its integrand samples + them, and its result is their sum. The widths of compute_decay_widths + need each of them, so each flavour becomes a subprocess of its own, + with the same matrix-element library and flavour index and the + diagrams active for it (see _single_flavor_metas). + """ + if self.merged_subprocess_data is not None: + raise ValueError( + "split_flavors needs merge_subprocesses = false: " + "merged_subprocesses.json indexes the unsplit subprocesses" + ) + self.subprocess_data = [ + single + for meta in self.subprocess_data + for single in _single_flavor_metas(meta) + ] + + 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_mode(self.run_card["run"]["madspace_cpu_mode"]) + + 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 + self.auto_temp_path = temp_dir == "auto" + if self.auto_temp_path: + self.temp_path = tempfile.mkdtemp(prefix="mg7_events_") + elif 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 = [] + # AmpliCol libraries, shared by all channels of their process + self.amplicol_apis = {} + # [path format, init path, parameters] of every load_matrix_element, in order + self.loaded_matrix_elements = [] + self.device_types = [] + self.devices = [] + self.pool_sizes = [] + 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)) + + def load_matrix_element(self, path_format: str, init_path: str, parameters: dict): + """Load a matrix-element library into every context, recorded for the + gridpack to load the same ones at the same indices. The init path is + the param card, or the artifact of an AmpliCol library.""" + for context, backend in zip(self.contexts, self.backends): + api = context.load_matrix_element( + path_format.format(device=backend), os.path.abspath(init_path), + parameters + ) + self.loaded_matrix_elements.append([path_format, init_path, parameters]) + return api + + CUSTOM_FUNCTION_KINDS = ("cut", "scale", "pdf", "bias") + + def init_custom_functions(self) -> None: + """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." + ) + parameters = self.run_card.custom_parameters(kind) + apis = load_custom_function( + ms, self.contexts, self.device_types, kind, path, parameters + ) + 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, parameters]) + 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 any(is_amplicol(meta) for meta in self.subprocess_data): + raise RuntimeError( + "[surrogate] is not supported with [amplicol] subprocesses: " + "set [surrogate] function = \"\" or use the madmatrix backend" + ) + # a .so path with {subprocess} or {device} is checked when it is loaded + if "{" not in path and not os.path.isfile(path): + raise RuntimeError(f"[surrogate] function = {path!r}: file not found") + if self.run_card["phasespace"]["sde_strategy"] != "denominators": + raise RuntimeError( + "[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 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 not (channel_indices and steps): + continue + group = ms.WeightingGroup(channel_indices, steps) + # [generation] sets the last unweighting, the first_pass_ ones the earlier + first_pass = subproc.first_pass_settings() + if first_pass is not None: + group.first_pass_max_overweight_truncation = ( + first_pass["first_pass_max_overweight_truncation"] + ) + group.first_pass_freeze_max_weight_after = ( + first_pass["first_pass_freeze_max_weight_after"] + ) + groups.append(group) + return groups + + @staticmethod + 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 apply_auto_ptj_mjj(self, cuts: dict) -> None: + """madevent's auto_ptj_mjj, from Template/LO/SubProcesses/setcuts.f. + + A merging cut already bounds every matrix-element jet from below, so + raising the jet pt and pair-mass cuts to xqcut removes no phase space + that survives the merging anyway while making the integration much more + efficient. The jet dR cuts go the other way and are dropped: + xqcut supersedes them, and a dR cut would carve a hole out of exactly + the region the parton shower is supposed to fill, leaving a gap in the + jet rates at the merging scale. + + With the option off, madevent still drops the dR cuts, and only zeroes + a pt or pair-mass cut that sits above xqcut - it never tightens. + """ + if self.is_decay or self.run_card["beam"]["dynamical_scale_choice"] != "mlm": + return + xqcut = self.run_card["phasespace"]["xqcut"] + if xqcut <= 0: + return + for key in ("jet-delta_r", "jet-lepton-delta_r"): + if cuts.get(key, {}).get("min", 0.0) > 0.0: + cuts[key]["min"] = 0.0 + auto = self.run_card["phasespace"]["auto_ptj_mjj"] + # Both halves of madevent's rule matter, and for different reasons. + # + # ptj reaches the phase space only for jets that are leaf children of + # the root: a jet inside a composite s-channel node - two jets off one + # gluon, say - has its pt_min reset to zero in PhaseSpaceMapping, and + # that node's invariant is then sampled from threshold upward knowing + # nothing about the cut. mjj is what bounds those nodes, because a + # pairwise invariant-mass minimum is carried into the propagator's own + # e_min and so into the s_min the sampler uses. + # + # Leaving mjj out is what makes the highest multiplicity misbehave at a + # tight xqcut: almost every point it proposes is vetoed, and the few + # survivors sit against a boundary it cannot see. + for key in ("jet-pt", "jet-pair_mass"): + if auto: + cuts.setdefault(key, {})["min"] = xqcut + elif cuts.get(key, {}).get("min", 0.0) > xqcut: + cuts[key]["min"] = 0.0 + + def init_cuts(self) -> None: + inf = float("inf") + order_observable = self.run_card["cuts"].get("order_by", "pt") + cuts = { + key: dict(values) + for key, values in self.run_card["cuts"].items() + if key != "order_by" + } + self.apply_auto_ptj_mjj(cuts) + 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 cuts.items() + ] + + # [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. LHAPDF itself is not needed: madspace reads the grid + files directly, so without a usable lhapdf-config (or when 'lhapdf + install' fails) the set is fetched straight from the LHAPDF server.""" + paths = lhapdf_paths() + if paths.find_set(pdf_set): + return paths + if not paths.download_path: + logger.debug("no writable PDF directory: cannot download %s", pdf_set) + return paths + logger.info("PDF set %s not found; downloading into %s", + pdf_set, paths.download_path) + if paths.config: + try: + from madgraph.interface.common_run_interface import CommonRunCmd + os.makedirs(paths.download_path, exist_ok=True) + CommonRunCmd.install_lhapdf_pdfset_static( + paths.config, paths.download_path, pdf_set) + except Exception as err: + logger.debug("lhapdf could not install %s: %s", pdf_set, err) + if not os.path.isdir(os.path.join(paths.download_path, pdf_set)): + try: + misc.download_lhapdf_set(pdf_set, paths.download_path) + 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) + # subprocesses (systematics, MadSpin, reweight) find it here too + os.environ["LHAPDF_DATA_PATH"] = os.pathsep.join(options._LHAPDF.data_paths) + 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 lhapdf-config found (not needed for the download)" + if self.lhapdf.download_path: + download = ("the download from %s into %s failed (see the " + "warning above)" % (misc.LHAPDF_SETS_URL, + self.lhapdf.download_path)) + else: + download = "no writable directory to download the set into" + return ( + "LHAPDF set %r, requested by [beam] pdf1/pdf2 in Cards/run_card.toml, " + "was not found.\n" + " searched: %s\n" + " %s\n" + " %s\n" + "Fix this in one of the following ways:\n" + " * check the set name against %s/pdfsets.index and your network\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, download, misc.LHAPDF_SETS_URL, + 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.beam_energies = (float(self.e_cm), 0.0) + self.beam_rapidity = 0.0 + self.leptonic = True + else: + ebeam1 = float(beam_args["ebeam1"]) + ebeam2 = float(beam_args["ebeam2"]) + self.beam_energies = (ebeam1, ebeam2) + # The phase space is generated in the beams' centre-of-mass frame + # (massless beams: s = 4 E1 E2) and boosted by this rapidity into + # the lab frame, where the events are written and the cuts applied. + self.e_cm = 2.0 * math.sqrt(ebeam1 * ebeam2) + self.beam_rapidity = 0.5 * math.log(ebeam1 / ebeam2) + 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, + } + # "mlm" is not one of the closed-form scale choices: it needs the diagram + # topologies of each subprocess, so the EnergyScale is built later, in + # MadgraphSubprocess.init_madspace, from an MLMClustering instance. + if beam_args["dynamical_scale_choice"] == "mlm": + self.mlm_clustering = True + self.scale_kwargs = {} + elif beam_args["dynamical_scale_choice"] in dynamical_scales: + self.mlm_clustering = False + self.scale_kwargs = dict( + dynamical_scale_type=dynamical_scales[ + beam_args["dynamical_scale_choice"] + ], + 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.)), + ) + else: + raise ValueError("Unknown dynamical scale choice") + 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.mlm_clustering and not self.is_decay: + # The merging needs the scales of its own clustering history: a + # custom scale or a factor on it would move the merging scale away + # from the one the shower is matched to. + if self.custom_functions.pop("scale", None) is not None: + logger.warning( + "custom scale ignored under dynamical_scale_choice = mlm") + if float(beam_args.get("scale_factor", 1.)) != 1.: + logger.warning( + "scale_factor ignored under dynamical_scale_choice = mlm") + if self.is_decay: + # One scale is available for a decay -- the decaying mass -- so use + # it, fixed, whatever the card asks for. That includes "mlm": a + # decay has no beams to cluster into, so there is no merging to do. + if self.mlm_clustering: + logger.warning( + "dynamical_scale_choice = mlm has no meaning for a decay; " + "using the decaying mass as a fixed scale instead" + ) + self.mlm_clustering = False + self.scale_kwargs.setdefault( + "dynamical_scale_type", ms.EnergyScale.half_transverse_mass + ) + 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.pdf_grid2 = None + self.lhapdf = None + self.pdf_dir = None + self.pdf_dir2 = 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["pdf1"] + self.pdf_set = pdf_set + self.pdf_set2 = beam_args["pdf2"] + 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)) + # A second grid only when the beams differ (and have a PDF at all); + # otherwise beam 2 shares the first one. + self.pdf_grid2 = None + self.pdf_dir2 = self.pdf_dir + if self.pdf_set2 != pdf_set and not self.leptonic: + self.pdf_dir2 = self.ensure_pdf_set(self.pdf_set2).find_set(self.pdf_set2) + if self.pdf_dir2 is None: + raise RuntimeError(self._no_pdf_message(self.pdf_set2)) + if "pdf" not in self.custom_functions: + self.pdf_grid2 = ms.PdfGrid(os.path.join( + self.pdf_dir2, self.pdf_set2, f"{self.pdf_set2}_0000.dat")) + logger.info("PDF sets: %s (beam 1), %s (beam 2); alpha_s from %s", + pdf_set, self.pdf_set2, 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") + ) + # The grid only covers a band in Q. Asking for a density above its + # highest Q reads past the end of the interpolation and comes back as a + # NaN, so the scale is capped there. A dynamical scale rarely gets near + # it, but the madevent MLM scheme, which multiplies four clustering + # scales together under a fourth root, can. A custom PDF brings no grid + # to read the range from, so it is left uncapped. + grids = [grid for grid in (self.pdf_grid, self.pdf_grid2) + if isinstance(grid, ms.PdfGrid)] + self.max_scale = min(max(grid.q) for grid in grids) if grids else 0.0 + # the bottom of the same range, where the densities freeze + self.min_pdf_scale = max(min(grid.q) for grid in grids) if grids else 0.0 + for context in self.contexts: + if not isinstance(self.pdf_grid, ms.CustomPdf): + self.pdf_grid.initialize_globals(context) + if self.pdf_grid2 is not None: + # the prefix ms.Integrand reads the second beam's grid from + # (Integrand::pdf2_prefix) + self.pdf_grid2.initialize_globals(context, "beam2") + 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"]] + # the emission scale of the MLM alpha_s vertices, which only an MLM + # sample with per-vertex alpha_s has + if (self.mlm_clustering and not self.is_decay + and self.run_card["beam"]["alphas_reweighting"] != "none"): + config.alps = [float(v) for v in syst["alps"]] + else: + config.alps = [1.0] + # the pdf reweighting ladder of an MLM sample, varied as madevent's + # systematics vary it; a varied scale is kept inside the PDF grid, + # where LHAPDF freezes it for madevent + config.pdf_ladder = mlm_pdf_ladder( + self.run_card, self.mlm_clustering, self.is_decay, self.leptonic) + if config.pdf_ladder: + config.ladder_min_scale = float(self.min_pdf_scale) + config.ladder_max_scale = float(self.max_scale) + 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"] + if self.pdf_grid2 is None: + config.pdf_members = self.resolve_pdf_variations() + else: + logger.warning( + "systematics: the beams use different PDF sets (%s, %s), so " + "no PDF variation is computed; the scale variations use both " + "sets", self.pdf_set, self.pdf_set2) + config.pdf_members = [] + 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, + nominal_pdf2=self.pdf_grid2 if config.has_pdf else None) + 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_grid_file2": os.path.join( + self.pdf_dir2, self.pdf_set2, f"{self.pdf_set2}_0000.dat") + if config.has_pdf and self.pdf_grid2 is not None 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 + # a custom scale replaces dynamical_scale_choice, so no variation is redundant + custom_scale = "scale" in getattr(self, "custom_functions", {}) + if not (beam["fixed_ren_scale"] and beam["fixed_fact_scale"] or custom_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 or v.is_emission) + 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))) + emission = systematics_summary.emission_percentages(summary) + if emission is not None: + rows.append(("Emission scale variation:", format_variation(*emission))) + 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) + + # the scale below which the alpha_s of a decay is frozen (write_fixed_alphas_info) + ALPHAS_FREEZE_SCALE = 1.0 + + def write_fixed_alphas_info(self) -> str: + """Write a minimal LHAPDF ``.info`` with the alpha_s of a decay. + + 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, so only alpha_s(M) matters. Like the Fortran decay + matrix elements (Source/alfas_functions.f, ``nloop = 2``) it is the + 2-loop running from the alpha_s(m_Z) of the param card up to M -- not + alpha_s(m_Z) itself, which overshoots alpha_s(125 GeV) by 4.7% (and + h > g g, which goes like alpha_s^2, by 9.7%). + + The grid carries the running over the whole scale range and has a node + exactly at the decaying mass, so the interpolation returns alpha_s(M) + untouched. Below ALPHAS_FREEZE_SCALE (or wherever the running stops + being defined, closer to the Landau pole) alpha_s is frozen at its last + value: a light decaying particle (the muon, ...) has no QCD to speak + of, and its widths must not fail on a coupling they do not use. + """ + # SMINPUTS entry 3 is alpha_s(m_Z), the same value the Fortran + # parameter setup feeds to G. + alpha_s_mz = float(self.param_card.get_value("sminputs", 3)) + runner = Alphas_Runner(alpha_s_mz, 2, cmass=1.42, bmass=4.7) + mass = float(self.decaying_mass) + 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: + # 1 GeV to 100 TeV, and the decaying mass. + nodes = [10 ** (i / 4.0) for i in range(0, 21)] + q_values = sorted( + [q for q in nodes if abs(q / mass - 1) > 1e-9] + [mass] + ) + + def running(q): + try: + value = runner(q) + except (ArithmeticError, ValueError): + return None + return value if math.isfinite(value) and value > 0 else None + + # from the top down, freezing at the first scale where the running is + # not used (below ALPHAS_FREEZE_SCALE) or not defined + values = {} + frozen = None + for q in reversed(q_values): + value = None + if frozen is None and q >= self.ALPHAS_FREEZE_SCALE: + value = running(q) + if value is None: + if frozen is None: + if not values: + raise ValueError( + f"alpha_s({q:g} GeV) cannot be computed from " + f"alpha_s(m_Z) = {alpha_s_mz:g}" + ) + frozen = min(values) + value = values[frozen] + values[q] = value + alpha_s_values = [values[q] for q in q_values] + if frozen is not None and mass < frozen: + logger.info( + "decay mode: alpha_s frozen below %g GeV at alpha_s(%g GeV) = " + "%.6g (decaying mass %g GeV)", frozen, frozen, values[frozen], + mass, + ) + with open(path, "w") as f: + f.write("SetDesc: 2-loop running alpha_s of a decay (no beams)\n") + f.write( + "AlphaS_Qs: [%s]\n" % ", ".join(f"{q:.12g}" for q in q_values) + ) + f.write( + "AlphaS_Vals: [%s]\n" + % ", ".join(f"{a:.12g}" for a in alpha_s_values) + ) + logger.info( + "decay mode: alpha_s(mu = %g GeV) = %.6g (2-loop running from " + "alpha_s(m_Z) = %g, as madevent)", + mass, values[mass], alpha_s_mz, + ) + 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"] + 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): + 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( + subproc_class(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] + 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 = madmatrix_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: + # by library: the flavours of a split subprocess share one + missing = list({ + meta["me_path"]: meta for meta in madmatrix_data + if not os.path.isfile(meta["me_path"].format(device=backend)) + }.values()) + 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() + make_args = [f"BACKEND={backend}", f"FPTYPE={fptype}", "USEBUILDDIR=1"] + make_args += self.extra_make_args + if (len(missing) >= PCH_MIN_SUBPROCESSES + and not any(arg.startswith("PCH=") for arg in make_args)): + make_args.append("PCH=1") + self.make_subprocesses(subproc_path, make_args, nb_core, log_path) + logger.info( + f"Compilation of SubProcesses done in {time.time() - start_time:.1f} s" + ) + 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 check in sorted({ + os.path.join(os.path.dirname(os.path.normpath(path)), "check-pyamplicol") + for path in paths + }): + # older outputs have no check script + if os.path.isfile(check): + self.check_pyamplicol(check) + for path in paths: + logger.info(f"Compiling the AmpliCol matrix elements in {path}, see {log_path}") + self.make_subprocesses(path, [], 1, log_path) + + @staticmethod + def check_pyamplicol(script: str) -> None: + """Run SubProcesses/check-pyamplicol, which fails if the installed + pyamplicol is not the one the output was made with.""" + result = subprocess.run([script], capture_output=True, text=True) + if result.returncode != 0: + message = (result.stderr.strip() or result.stdout.strip() + or f"{script} failed with exit status {result.returncode}") + raise RuntimeError(message) + + @staticmethod + def make_subprocesses( + subproc_path: str, args: list[str], nb_core: int, log_path: str + ) -> 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) + + groups = self.weighting_groups(subprocess_channels) + # the groups of the last generator built, the one a gridpack stores + self.weighting_groups_json = ms.weighting_groups_to_json(groups) + event_generator = ms.EventGenerator( + contexts=self.contexts, + channels=channel_generators, + status_file=self.status_file, + config=self.event_generator_config, + seed=self.run_seed, + groups=groups, + ) + 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, + min( + (channel.phasespace_mapping.random_dim() for channel in ps.channels), + default=0, + ), + ) + 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 = [] + trained_phasespaces = [] + training_args = [] + self.event_generator = None + for subproc, phasespace in zip(self.subprocesses, self.phasespaces): + if not subproc.madnis_supported: + # no flow for a one-dimensional phase space: keep its VEGAS + # channels, and their event generators, from the survey + madnis_phasespaces.append(phasespace) + continue + 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) + trained_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( + trained_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 + + @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 + 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.""" + 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], list(self.beam_energies) + 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, list(self.beam_energies) + + def _lhapdf_id(self, beam=1): + """Central LHAPDF id of the PDF set of `beam` (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.pdf_set if beam == 1 else self.pdf_set2 + pdf_dir = self.pdf_dir if beam == 1 else self.pdf_dir2 + info = os.path.join(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() + 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=self._lhapdf_id(1), beam2_pdf_id=self._lhapdf_id(2), + 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 = [ + 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 = [ + 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, + 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")) + + include_estimates = self.gridpack_run_mode == "fix_max_weight" + # 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(include_estimates)}" + for channel in self.event_generator.channels() + ] + with open(os.path.join(data_path, "channels.json"), "w") as f: + f.write("{" + ",".join(channel_entries) + "}") + with open(os.path.join(data_path, "groups.json"), "w") as f: + f.write(self.weighting_groups_json) + + lib_path = os.path.join(gridpack_path, "lib") + if self.run_card["gridpack"]["include_source"]: + 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) + # the gridpack refers to its own copies, by paths relative to it + gridpack_custom_paths = [] + for kind, path, parameters in self.custom_function_paths: + rel_path = os.path.normpath(path) + if os.path.isabs(rel_path) or rel_path.startswith(".."): + rel_path = os.path.join("custom_functions", kind, os.path.basename(path)) + target = os.path.join(gridpack_path, rel_path) + if not os.path.exists(target): + os.makedirs(os.path.dirname(target), exist_ok=True) + shutil.copy(path, target) + gridpack_custom_paths.append([kind, rel_path, parameters]) + + 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"), + ) + + # the AmpliCol artifacts are read at run time; the ratio library finds + # them next to its own, otherwise empty, reweight directory + for artifact_dir in sorted({ + os.path.dirname(meta["artifact"]) + for meta in self.subprocess_data if is_amplicol(meta) + }): + for name in ("lc", "full"): + shutil.copytree( + os.path.join(artifact_dir, name), + os.path.join(gridpack_path, artifact_dir, name), + ) + os.makedirs(os.path.join(gridpack_path, artifact_dir, "reweight")) + + cards_path = os.path.join(gridpack_path, "Cards") + os.mkdir(cards_path) + 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.join(os.path.dirname(os.path.dirname(__file__)), + "iolibs", "template_files", "mg7") + shutil.copy(os.path.join(template_dir, "gridpack.py"), gen_events_file) + shutil.copy(os.path.join(template_dir, "custom_function_loader.py"), bin_path) + os.chmod(gen_events_file, 0o755) + 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": self.loaded_matrix_elements, + # loaded before the matrix elements, as in the run + "custom_functions": gridpack_custom_paths, + "me_parameters": me_parameters(self.run_card), + "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"]: + archive_path = gridpack_path + ".tar.gz" + use_pigz = self.run_card["run"]["use_pigz"] + pigz = (shutil.which("pigz") if use_pigz is not False else None) + if pigz: + # multi-threaded gzip + with open(archive_path, "wb") as f: + tar = subprocess.Popen( + ["tar", "-cf", "-", "-C", self.run_path, "gridpack"], + stdout=subprocess.PIPE, + ) + gz_command = [pigz, "-c"] + cpu_threads = self.run_card["run"]["cpu_thread_pool_size"] + if cpu_threads > 0: + gz_command += ["-p", str(cpu_threads)] + gz = subprocess.run(gz_command, stdin=tar.stdout, stdout=f) + tar.stdout.close() + if tar.wait() != 0 or gz.returncode != 0: + raise RuntimeError("compressing the gridpack failed") + else: + 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) + + +def _single_flavor_metas(meta: dict) -> list[dict]: + """The subprocess ``meta`` as one subprocess per flavour (see + MadgraphProcess.split_decay_flavors). + + Each copy keeps the flavour (its matrix-element index and its assignments) + and, in every channel, the diagrams that flavour is active in, renumbered + to its only flavour, 0. A channel left without diagrams is dropped, and a + flavour left without channels (no diagram at all) has no subprocess. A + subprocess of one flavour is returned as it is. + """ + flavors = meta["flavors"] + if len(flavors) == 1: + return [meta] + singles = [] + for position, flavor in enumerate(flavors): + channels = [] + for channel in meta["channels"]: + diagrams = [ + dict(diagram, active_flavors=[0]) + for diagram in channel["diagrams"] + if position in diagram["active_flavors"] + ] + if diagrams: + channels.append(dict(channel, diagrams=diagrams)) + if channels: + singles.append(dict(meta, flavors=[flavor], channels=channels)) + return singles diff --git a/madgraph/mg7/model_widths.py b/madgraph/mg7/model_widths.py new file mode 100644 index 0000000000..bd37ee210d --- /dev/null +++ b/madgraph/mg7/model_widths.py @@ -0,0 +1,160 @@ +"""The model of a process directory and its auto widths. Does not import madspace.""" + +import logging +import os +import re +import subprocess + +from madgraph.models.check_param_card import group_auto_widths +from madgraph.various import misc +from madgraph.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:`madgraph.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 madgraph.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 + # "DECAY auto", with the options of compute_widths that may follow + # ("auto@NLO", "auto --nlo --body_decay=3"), as in + # common_run_interface.static_check_param_card: one compute_widths per set of + # options + groups = group_auto_widths(text) + if not groups: + return + pdgs = [pdg for group in groups.values() for pdg in group] + + 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\n" % model + for options, group in groups.items(): + cmds += "compute_widths %s --path=%s %s\n" % ( + " ".join(group), os.path.abspath(param_card_path), " ".join(options)) + 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/mg7/npy_to_lhe.py b/madgraph/mg7/npy_to_lhe.py new file mode 100644 index 0000000000..d7d810097d --- /dev/null +++ b/madgraph/mg7/npy_to_lhe.py @@ -0,0 +1,252 @@ +"""Convert an mg7 numpy event file (``events.npy``) into an LHE file. + +The mg7 output writes its events as npy by default. Anything that needs an LHE +file -- a shower, MadSpin, reweighting, an analysis -- normally switches the run +to the LHE format before the generation (``force_lhe_output_if_needed``). This +module covers the case where that request comes *after* the run: it rebuilds +the LHE file from what the npy run left behind. + +* ``output_format = "lhe_npy"``: the file already holds complete LHE events + (pdg ids, status, mothers, colour, helicity); they are only reformatted. +* ``output_format = "compact_npy"``: the file holds the momenta plus the + subprocess/diagram/colour/flavour/helicity indices the generation picked. + Those are completed exactly as the direct LHE output does it, with the + :class:`madspace.LHECompleter` the run saved next to the events + (``lhe_completer.json``). The completion draws random numbers (colour flow + among equivalent ones, resonance assignment), so the result is statistically + but not byte-identical to an ``output_format = "lhe"`` run with the same seed. + +In both cases the LHE header and ```` block (with the systematics +````) are taken verbatim from the ``header.lhe`` file the run writes +next to its npy events; the ``rwgt_`` columns become the per-event +```` block. The ```` reweighting-input block (``[systematics] +write_inputs``) is not reconstructed. + +Usage:: + + bin/npy_to_lhe [RUN_NAME | RUN_DIR | path/to/events.npy] [-o OUTPUT] [--seed N] +""" + +import argparse +import logging +import os +import re +import sys + +logger = logging.getLogger("madevent") + +COMPLETER_FILE = "lhe_completer.json" +HEADER_FILE = "header.lhe" +EVENT_FILE = "events.npy" + +# events converted per LHEFileWriter.write_string call +_CHUNK = 10000 + + +def _madspace(): + from madgraph.mg7.bootstrap import ensure_madspace + ensure_madspace(interactive=False) + import madspace + return madspace + + +def _read_header(run_path): + """The header.lhe text without its closing tag, ready to be followed by + the events.""" + with open(os.path.join(run_path, HEADER_FILE)) as f: + text = f.read() + end = text.rfind("") + return text if end < 0 else text[:end] + + +def _seed(header, seed): + if seed is not None: + return int(seed) + match = re.search(r"\s*(-?\d+)\s*", header) + return int(match.group(1)) if match else 0 + + +# ---------------------------------------------------------------------------- +# conversion +# ---------------------------------------------------------------------------- +def resolve_event_file(path, me_dir=None) -> str: + """``path`` may be an npy file, a run directory or a run name (looked up + under ``/Events``).""" + candidates = [path, os.path.join(path, EVENT_FILE)] + if me_dir is not None: + candidates.append(os.path.join(me_dir, "Events", path, EVENT_FILE)) + for candidate in candidates: + if os.path.isfile(candidate): + return os.path.abspath(candidate) + raise FileNotFoundError("no %s found for '%s'" % (EVENT_FILE, path)) + + +def can_convert(run_path) -> bool: + """True when ``run_path`` holds npy events this module can turn into LHE.""" + return (os.path.isfile(os.path.join(run_path, EVENT_FILE)) + and os.path.isfile(os.path.join(run_path, HEADER_FILE))) + + +def _particle_count(names): + count = 0 + while "part%d_energy" % (count + 1) in names: + count += 1 + return count + + +def _compact_events(ms, events, completer, columns, rng): + names = events.dtype.names + nparticles = _particle_count(names) + momenta = [tuple(events["part%d_%s" % (i + 1, c)] + for c in ("energy", "px", "py", "pz")) + for i in range(nparticles)] + rwgt = [events[c] for c in columns] + for n in range(len(events)): + particles = [] + for energy, px, py, pz in momenta: + if energy[n] == 0.: + break + particles.append(ms.LHEParticle(px=px[n], py=py[n], pz=pz[n], + energy=energy[n])) + event = ms.LHEEvent(weight=events["weight"][n], + scale=events["ren_scale"][n], + alpha_qcd=events["alpha_qcd"][n], + particles=particles) + completer.complete_event_data( + event, int(events["subprocess_index"][n]), + int(events["diagram_index"][n]), int(events["color_index"][n]), + int(events["flavor_index"][n]), int(events["helicity_index"][n]), + rng) + yield event, [float(w[n]) for w in rwgt] + + +_LHE_PARTICLE_FIELDS = ("pdg_id", "status_code", "mother1", "mother2", "color", + "anti_color", "px", "py", "pz", "energy", "mass", + "lifetime", "spin") + + +def _lhe_npy_events(ms, events, columns): + names = events.dtype.names + count = 0 + while "part%d_pdg_id" % (count + 1) in names: + count += 1 + fields = [[events["part%d_%s" % (i + 1, f)] for f in _LHE_PARTICLE_FIELDS] + for i in range(count)] + rwgt = [events[c] for c in columns] + for n in range(len(events)): + particles = [] + for values in fields: + if values[0][n] == 0: + break + particles.append(ms.LHEParticle( + *[v[n].item() for v in values])) + event = ms.LHEEvent(int(events["process_id"][n]), + events["weight"][n], events["scale"][n], + events["alpha_qed"][n], events["alpha_qcd"][n], + particles) + yield event, [float(w[n]) for w in rwgt] + + +def convert(path, output=None, seed=None, compress=True, me_dir=None) -> str: + """Write the LHE file for the npy events at ``path`` (npy file, run + directory or run name) and return its path. By default it is written as + ``events.lhe.gz`` next to the npy file.""" + import numpy as np + from madgraph.various import misc + + ms = _madspace() + npy_path = resolve_event_file(path, me_dir) + run_path = os.path.dirname(npy_path) + # A standalone conversion starts in a fresh process, so it has not seen + # the gzip configuration installed by the event-generation launcher. + card_root = me_dir or os.path.dirname(os.path.dirname(run_path)) + card_path = os.path.join(card_root, "Cards", "run_card.toml") + if os.path.isfile(card_path): + from madgraph.various.banner import RunCardMG7 + misc.configure_gzip(RunCardMG7(card_path).gzip_configuration()) + if not os.path.isfile(os.path.join(run_path, HEADER_FILE)): + raise FileNotFoundError( + "%s is missing: the run predates the npy->LHE conversion support " + "and cannot be converted." % os.path.join(run_path, HEADER_FILE)) + header = _read_header(run_path) + + events = np.load(npy_path, mmap_mode="r") + names = events.dtype.names + columns = [n for n in names if re.fullmatch(r"rwgt_\d+", n)] + ids = [int(c.split("_", 1)[1]) for c in columns] + if "diagram_index" in names: + completer_path = os.path.join(run_path, COMPLETER_FILE) + if not os.path.isfile(completer_path): + raise FileNotFoundError("%s is missing" % completer_path) + completer = ms.LHECompleter.load(completer_path) + rng = ms.MixMaxRandom(_seed(header, seed)) + stream = lambda chunk: _compact_events(ms, chunk, completer, columns, rng) + elif "part1_pdg_id" in names: + stream = lambda chunk: _lhe_npy_events(ms, chunk, columns) + else: + raise ValueError("%s is not an mg7 event file" % npy_path) + + if output is None: + output = os.path.join(run_path, "events.lhe") + if output.endswith(".gz"): + output, compress = output[:-3], True + logger.info("converting %s (%d events) to LHE", npy_path, len(events)) + with open(output, "w") as out: + out.write(header) + for start in range(0, len(events), _CHUNK): + chunk = np.asarray(events[start:start + _CHUNK]) + text = [] + for event, weights in stream(chunk): + if ids: + event.rwgt_ids = ids + event.rwgt = weights + text.append(event.format()) + out.write("".join(text)) + out.write("\n") + if compress: + misc.gzip(output) + output += ".gz" + logger.info("LHE events written to %s", output) + return output + + +def main(argv=None, me_dir=None) -> None: + """Command-line entry point; run names are looked up under + ``/Events`` (default: the working directory).""" + parser = argparse.ArgumentParser( + description="Convert an mg7 events.npy file into an LHE file.") + parser.add_argument("path", nargs="?", default=None, + help="run name, run directory or events.npy file " + "(default: the most recent run with npy events " + "and no LHE file)") + parser.add_argument("-o", "--output", default=None, + help="output file (default: events.lhe.gz next to the " + "npy file)") + parser.add_argument("--seed", type=int, default=None, + help="seed of the LHE completion (default: the run seed)") + parser.add_argument("--no-gzip", action="store_false", dest="compress", + help="do not compress the output") + args = parser.parse_args(argv) + logging.basicConfig(level=logging.INFO, format="%(message)s") + + me_dir = me_dir or os.getcwd() + path = args.path + if path is None: + events_dir = os.path.join(me_dir, "Events") + runs = [os.path.join(events_dir, d) for d in + (os.listdir(events_dir) if os.path.isdir(events_dir) else [])] + runs = [d for d in runs if can_convert(d) and not any( + os.path.exists(os.path.join(d, n)) + for n in ("events.lhe", "events.lhe.gz"))] + if not runs: + sys.exit("no run with unconverted npy events in %s" % events_dir) + path = max(runs, key=lambda d: os.path.getmtime( + os.path.join(d, EVENT_FILE))) + try: + convert(path, args.output, args.seed, args.compress, me_dir) + except (FileNotFoundError, ValueError) as error: + sys.exit(str(error)) + + +if __name__ == "__main__": + main() diff --git a/madgraph/mg7/observables.py b/madgraph/mg7/observables.py new file mode 100644 index 0000000000..5c55cc57be --- /dev/null +++ b/madgraph/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.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/mg7/options.py b/madgraph/mg7/options.py new file mode 100644 index 0000000000..a983efbf0a --- /dev/null +++ b/madgraph/mg7/options.py @@ -0,0 +1,99 @@ +"""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 + + +def lhapdf_index_paths() -> list: + """The PDF directories to look an lhaid up in. Without LHAPDF there may be + no pdfsets.index anywhere yet, so fetch one into the download directory.""" + + paths = lhapdf_paths() + data_paths = list(paths.data_paths) + if not any(os.path.exists(os.path.join(p, "pdfsets.index")) for p in data_paths) \ + and paths.download_path \ + and misc.download_lhapdf_index(paths.download_path): + data_paths.append(paths.download_path) + return data_paths diff --git a/madgraph/iolibs/template_files/mg7/plots.py b/madgraph/mg7/plots.py similarity index 100% rename from madgraph/iolibs/template_files/mg7/plots.py rename to madgraph/mg7/plots.py diff --git a/madgraph/mg7/postprocess.py b/madgraph/mg7/postprocess.py new file mode 100644 index 0000000000..55c891bf2a --- /dev/null +++ b/madgraph/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.mg7 import npy_to_lhe +from madgraph.mg7.options import lhapdf_paths, load_mg5_options + +def _find_event_file(run_path): + """Locate the LHE event file produced by generate_events (if any).""" + 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.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 madgraph.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/mg7/run_interface.py similarity index 99% rename from madgraph/iolibs/template_files/mg7/run_interface.py rename to madgraph/mg7/run_interface.py index 24f6c8adef..6e541e800b 100644 --- a/madgraph/iolibs/template_files/mg7/run_interface.py +++ b/madgraph/mg7/run_interface.py @@ -261,7 +261,7 @@ def get_model(self): tools this class drives (MadSpin, reweight, the auto widths and the 'update dependent' pass of their card questions) work here too. """ - from madgraph.iolibs.template_files.mg7.launch import load_process_model + from madgraph.mg7.model_widths import load_process_model return load_process_model(self.me_dir) def do_set(self, line, log=True): diff --git a/madgraph/mg7/subprocess_amplicol.py b/madgraph/mg7/subprocess_amplicol.py new file mode 100644 index 0000000000..45339c6514 --- /dev/null +++ b/madgraph/mg7/subprocess_amplicol.py @@ -0,0 +1,444 @@ +"""Subprocesses with AmpliCol matrix elements ('[amplicol]' processes). + +Each subprocess is one UMAMI channel of an AmpliCol provider: colour orderings +integrated at leading colour over a common set of colour-ordered phase-space +channels. They differ in where the colour singlets attach to the colour lines +and in the s-channel resonances they map out, and are weighted by the densities +of their mappings. Resonances that cannot go on shell with the param card +masses and widths are dropped at launch. A second pass re- or +unweights the events to full colour with the ratio of the full-colour to the +leading-colour matrix element summed over all colour orderings, at the +helicity of the event. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from madgraph.mg7.common import lhe_helicities, ms +from madgraph.mg7.observables import clean_pids +from madgraph.mg7.subprocess_madmatrix import ( + Channel, + MadgraphSubprocess, + PhaseSpace, +) + +if TYPE_CHECKING: + from madgraph.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 channel_topology( + incoming_masses: list[float], + outgoing_masses: list[float], + order: list[int], + resonances: list[tuple[int, list[int]]], + process: MadgraphProcess, +) -> tuple[ms.Topology, list[int]]: + """The topology of a colour-ordered phase-space channel, a t-channel chain + through the final-state legs and the resonances (pdg, legs) that decay to + them, and its colour order over the children of that chain.""" + n_in = len(incoming_masses) + legs = [frozenset(r_legs) for _, r_legs in resonances] + top = [i for i, r in enumerate(legs) if not any(r < other for other in legs)] + children = [] + for slot in order: + if slot < n_in: + continue + child = next((("r", i) for i in top if slot in legs[i]), ("o", slot)) + if child not in children: + children.append(child) + elif children[-1] != child: + raise RuntimeError(f"the resonance legs of {order} are not contiguous") + chain_count = len(children) - 1 + + def ref(child): + kind, index = child + return f"o{index - n_in}" if kind == "o" else f"p{chain_count + index}" + + refs = [ref(child) for child in children] + if len(children) == 1: + vertices = [["i0", "i1", refs[0]]] + else: + vertices = [["i0", "p0", refs[0]]] + vertices += [[f"p{i - 1}", f"p{i}", refs[i]] for i in range(1, chain_count)] + vertices.append(["i1", f"p{chain_count - 1}", refs[-1]]) + position = {slot: i for i, slot in enumerate(order)} + for index, r_legs in enumerate(legs): + inner = [ + i for i, other in enumerate(legs) + if other < r_legs and not any(other < mid < r_legs for mid in legs) + ] + covered = set().union(*(legs[i] for i in inner)) + decay = [(min(position[s] for s in legs[i]), ref(("r", i))) for i in inner] + decay += [(position[s], ref(("o", s))) for s in r_legs - covered] + vertices.append([f"p{chain_count + index}", *(r for _, r in sorted(decay))]) + propagators = [ms.Propagator() for _ in range(chain_count)] + for pdg, _ in resonances: + propagators.append(ms.Propagator( + mass=process.get_mass(abs(pdg)), + width=process.get_width(abs(pdg)), + pdg_id=pdg, + )) + topology = ms.Topology( + ms.Diagram(incoming_masses, outgoing_masses, propagators, vertices) + ) + color_order = [] + for slot in order: + node = slot if slot < n_in else n_in + children.index( + next((("r", i) for i in top if slot in legs[i]), ("o", slot)) + ) + if node not in color_order: + color_order.append(node) + return topology, color_order + + +def channel_topologies(meta: dict, process: MadgraphProcess) -> list[tuple]: + """channel_topology of every phase-space channel of a subprocess.""" + incoming_masses = [process.get_mass(pid) for pid in clean_pids(meta["incoming"])] + outgoing_masses = [process.get_mass(pid) for pid in clean_pids(meta["outgoing"])] + return [ + channel_topology( + incoming_masses, + outgoing_masses, + channel["order"], + [(r["pdg"], r["legs"]) for r in channel["resonances"]], + process, + ) + for channel in meta["phasespace_channels"] + ] + + +def active_resonances(meta: dict, get_mass, get_width) -> dict: + """The subprocess meta without the s-channel resonances that cannot go on + shell: a resonance is active if its width is positive and its mass above + the masses of its decay legs. Phase-space channels that then agree are + merged, and the flavours renumbered to them.""" + leg_masses = [ + get_mass(pid) for pid in clean_pids(meta["incoming"] + meta["outgoing"]) + ] + + def active(resonance): + pdg = abs(resonance["pdg"]) + return get_width(pdg) > 0 and get_mass(pdg) > sum( + leg_masses[leg] for leg in resonance["legs"] + ) + + for resonance in meta.get("colored_resonances", []): + if active(resonance): + raise RuntimeError( + f"the coloured s-channel resonance {resonance['pdg']} can go on " + "shell, which [amplicol] phase-space mappings do not follow yet; " + "generate the process with on-shell decays or with the madmatrix " + "backend" + ) + channels, lookup, remap, kept = [], {}, [], [] + for channel in meta["phasespace_channels"]: + keep = [i for i, r in enumerate(channel["resonances"]) if active(r)] + resonances = [channel["resonances"][i] for i in keep] + key = ( + tuple(channel["order"]), + tuple((abs(r["pdg"]), tuple(r["legs"])) for r in resonances), + ) + if key not in lookup: + lookup[key] = len(channels) + channels.append({**channel, "resonances": resonances}) + remap.append(lookup[key]) + kept.append(keep) + flavors = [] + for flavor in meta["flavors"]: + indices, pdgs = [], [] + for index, resonance_pdgs in zip(flavor["channels"], flavor["resonance_pdgs"]): + if remap[index] not in indices: + indices.append(remap[index]) + pdgs.append([resonance_pdgs[i] for i in kept[index]]) + flavors.append({**flavor, "channels": indices, "resonance_pdgs": pdgs}) + return {**meta, "phasespace_channels": channels, "flavors": flavors} + + +def lhe_diagrams(meta: dict) -> tuple[list[tuple[int, tuple]], list[list[float]]]: + """The diagrams that the LHE writer takes the resonances of an event from, + one per phase-space channel and resonance charges, and per flavour a mask + of those it can be written with. None if there are no resonances.""" + if not any(channel["resonances"] for channel in meta["phasespace_channels"]): + return [], [] + diagrams = [] + for flavor in meta["flavors"]: + for channel, pdgs in zip(flavor["channels"], flavor["resonance_pdgs"]): + if (channel, tuple(pdgs)) not in diagrams: + diagrams.append((channel, tuple(pdgs))) + masks = [ + [ + float(diagram in zip(flavor["channels"], map(tuple, flavor["resonance_pdgs"]))) + for diagram in diagrams + ] + for flavor in meta["flavors"] + ] + return diagrams, masks + + +def amplicol_subproc_args(meta: dict, process: MadgraphProcess) -> ms.SubprocArgs: + """LHE completion of an AmpliCol subprocess: its diagrams are the + phase-space channels, with the resonances they map out.""" + meta = active_resonances(meta, process.get_mass, process.get_width) + diagrams, _ = lhe_diagrams(meta) + if diagrams: + topologies = channel_topologies(meta, process) + particle_count = len(meta["incoming"]) + len(meta["outgoing"]) + colors = sorted({flavor["color_index"] for flavor in meta["flavors"]}) + by_channel = [ + [(d, pdgs) for d, (c, pdgs) in enumerate(diagrams) if c == channel] + for channel in range(len(topologies)) + ] + # the chain propagators come first: one less than the root's children + chain_counts = [len(order) - len(meta["incoming"]) - 1 for _, order in topologies] + layout = dict( + topologies=[topology for topology, _ in topologies], + permutations=[ + [list(range(particle_count))] * len(channel) for channel in by_channel + ], + diagram_indices=[[d for d, _ in channel] for channel in by_channel], + diagram_color_indices=[[colors] * len(channel) for channel in by_channel], + diagram_propagator_pdgs=[ + [[0] * count + list(pdgs) for _, pdgs in channel] + for channel, count in zip(by_channel, chain_counts) + ], + ) + else: + layout = dict( + topologies=[chain_topology( + [process.get_mass(pid) for pid in clean_pids(meta["incoming"])], + [process.get_mass(pid) for pid in clean_pids(meta["outgoing"])], + )], + permutations=[[]], + diagram_indices=[[]], + diagram_color_indices=[[]], + ) + return ms.SubprocArgs( + **layout, + 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"]], + pdg_id_weights=[flavor.get("weights", []) for flavor in meta["flavors"]], + ) + + +class AmplicolSubprocess(MadgraphSubprocess): + def __init__(self, process, meta, subproc_id, unmerged_meta=None): + if process.mlm_clustering: + # the clustering needs Feynman-diagram topologies, which AmpliCol lacks + raise RuntimeError( + "dynamical_scale_choice = mlm is not supported for [amplicol] " + f"subprocesses ({meta['incoming']} > {meta['outgoing']}): " + "use a closed-form scale choice" + ) + super().__init__(process, meta, subproc_id, unmerged_meta) + self.meta = active_resonances(self.meta, process.get_mass, process.get_width) + self.check_flavor_blind_cuts() + + def check_flavor_blind_cuts(self): + """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, 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") + gridpack_args = self.process.run_card["gridpack"] + if gridpack_args["save_gridpack"] and gridpack_args["include_source"]: + raise RuntimeError( + "gridpacks of [amplicol] processes cannot include the sources" + ) + if self.meta["qcd_power"] < 0 and self.process.systematics_enabled(): + raise RuntimeError( + "the scale and PDF systematics of [amplicol] processes that mix " + "coupling orders are not supported yet: disable [systematics]" + ) + self.supports_alpha_s = self.meta.get("supports_alpha_s", True) + if not self.supports_alpha_s and not self.process.run_card["beam"]["fixed_ren_scale"]: + raise RuntimeError( + "the AmpliCol matrix elements of this model cannot take the running " + "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"] + 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): + # loaded once for all channels; each load reads the whole pyamplicol artifact + apis = self.process.amplicol_apis + if path not in apis: + apis[path] = self.process.load_matrix_element(path, artifact, parameters) + return apis[path] + + def build_multi_channel_data(self): + return None + + def build_multichannel_phasespace(self) -> PhaseSpace: + """One channel per phase-space channel of the UMAMI channel, weighted + by the densities of their mappings.""" + flavors = self.meta["flavors"] + particle_count = len(self.incoming_masses) + len(self.outgoing_masses) + channels = [] + for index, (topology, color_order) in enumerate( + channel_topologies(self.meta, self.process) + ): + mapping = ms.PhaseSpaceMapping( + topology, + self.process.e_cm, + leptonic=self.process.leptonic, + invariant_power=self.process.run_card["phasespace"]["invariant_power"], + t_channel_mode=ms.PhaseSpaceMapping.color_ordered, + cuts=self.cuts, + permutations=[list(range(particle_count))], + color_ordered=ms.ColorOrderedMapping.Options(color_order), + beam_rapidity=self.process.beam_rapidity, + mirror_beams=self.mirror_beams, + ) + prefix = f"subproc{self.subproc_id}.amplicol{index}" + discrete_sym, discrete_flavor = self.build_discrete( + 1, len(flavors), prefix + ) + channels.append(Channel( + phasespace_mapping=mapping, + adaptive_mapping=self.build_vegas(mapping, prefix), + discrete_sym=discrete_sym, + discrete_flavor=discrete_flavor, + channel_weight_indices=[index], + name=f"C{index}", + active_flavors=[[ + f for f, flavor in enumerate(flavors) if index in flavor["channels"] + ]], + )) + diagrams, diagram_masks = lhe_diagrams(self.meta) + weights = None + if len(channels) > 1 or diagrams: + weights = ms.MappingChannelWeights( + [channel.phasespace_mapping for channel in channels], + [ + [float(index in flavor["channels"]) for index in range(len(channels))] + for flavor in flavors + ], + [channel for channel, _ in diagrams], + diagram_masks, + ) + return PhaseSpace( + mode="multichannel", + channels=channels, + symfact=[None] * len(channels), + mapping_chan_weights=weights, + ) + + 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: + # without an alpha_s input, the model value of the fixed scale is used + inputs = [ + ms.MatrixElement.momenta_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, + ] + 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"], + )] + + def first_pass_settings(self) -> dict | None: + amplicol = self.process.run_card["amplicol"] + if self.ratio_matrix_element is None or amplicol["mode"] != "unweight": + return None + return amplicol diff --git a/madgraph/mg7/subprocess_madmatrix.py b/madgraph/mg7/subprocess_madmatrix.py new file mode 100644 index 0000000000..c035e110c2 --- /dev/null +++ b/madgraph/mg7/subprocess_madmatrix.py @@ -0,0 +1,1258 @@ +"""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.mg7.common import logger, me_parameters, ms +from madgraph.iolibs.template_files.mg7.custom_function_loader import load_surrogate +from madgraph.mg7.observables import build_observable, clean_pids, decay_products, observable_pids, pid_is_qcd + +if TYPE_CHECKING: + from madgraph.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 + mapping_chan_weights: ms.MappingChannelWeights | 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 + # set from the survey by set_madnis_auto_settings + self.madnis_supported = True + 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 + + 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 + # A flavor standing for both beam orientations draws the orientation + # before the phase-space mapping, so boost, cuts and PDFs all see it. + self.mirror_beams = not self.process.is_decay and self.has_mirrored_flavors() + # 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 + ) + + # The scale range exists to keep the densities inside the PDF grid, so + # a run without beam densities - a lepton collider or a decay - has + # none: a floor there would only throw away events whose scale is + # legitimately low, a decay of a particle lighter than min_scale say. + if self.process.leptonic: + min_scale = max_scale = 0.0 + else: + min_scale = self.process.run_card["beam"]["min_scale"] + max_scale = self.process.max_scale + if self.process.mlm_clustering: + self.scale = self.build_mlm_scale(min_scale, max_scale) + else: + self.scale = ms.EnergyScale( + particle_count=self.particle_count, + custom_scale=self.process.custom_function( + "scale", self.particle_count, representative_pids + ), + min_scale=min_scale, + max_scale=max_scale, + **self.process.scale_kwargs, + ) + + # the libraries were all built up front by MadgraphProcess.compile_matrix_elements + self.load_matrix_elements(api_path_formats, member_metas) + + def build_mlm_scale(self, min_scale: float, max_scale: float): + """The scales of the MLM clustering history of each event.""" + mc_data = self.build_multi_channel_data() + beam = self.process.run_card["beam"] + all_pids = (clean_pids(self.meta["incoming"]) + + clean_pids(self.meta["outgoing"])) + return ms.EnergyScale( + ms.MLMClustering( + topologies=[topo[0] for topo in mc_data.topologies], + permutations=mc_data.permutations, + diagram_indices=mc_data.diagram_indices, + cm_energy=self.process.e_cm, + jet_scale_scheme=getattr( + ms.MLMClustering.JetScaleScheme, beam["jet_scale_scheme"]), + scale_scheme=getattr( + ms.MLMClustering.ScaleScheme, beam["scale_scheme"]), + pdg_color_types={ + int(key): value + for key, value in self.meta["pdg_color_types"].items() + }, + xqcut=self.process.run_card["phasespace"]["xqcut"], + bw_cutoff=self.process.run_card["phasespace"]["bw_cutoff"], + jet_radius=beam["jet_radius"], + hadronic=not self.process.leptonic, + external_pdg_ids=all_pids, + max_jet_flavor=beam["max_jet_flavor"], + parton_line_scheme=getattr( + ms.MLMClustering.PartonLineScheme, beam["parton_line_scheme"]), + alphas_scheme=getattr( + ms.MLMClustering.AlphasScheme, beam["alphas_reweighting"]), + pdf_reweighting=beam["pdf_reweighting"], + clustering_measure=getattr( + ms.MLMClustering.ClusteringMeasure, beam["clustering_measure"]), + clustering_history=getattr( + ms.MLMClustering.ClusteringHistory, beam["clustering_history"]), + clustering_mean_pdf_ladder=beam["clustering_mean_pdf_ladder"], + ), + min_scale=min_scale, + max_scale=max_scale, + ) + + def load_matrix_elements(self, api_path_formats, member_metas): + """The matrix element and [surrogate] of every member subprocess.""" + if self.process.run_card["run"]["dummy_matrix_element"]: + self.matrix_elements = [None] * len(api_path_formats) + else: + self.matrix_elements = [ + self.process.load_matrix_element( + path_format, self.process.param_card_path, + me_parameters(self.process.run_card) + ) + for path_format in api_path_formats + ] + self.surrogates = [] + if self.process.surrogate is not None: + for path_format, member_meta in zip(api_path_formats, member_metas): + 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"], + me_parameters(self.process.run_card), + ) + if len({api.index() for api in apis}) != 1: + 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.""" + 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 first_pass_settings(self) -> dict | None: + """The run card section with the first_pass_ unweighting settings, if + the events are unweighted again after the first unweighting.""" + if not self.surrogates or self.process.surrogate["mode"] != "unweight": + return None + return self.process.surrogate + + def flavor_pids(self) -> list[list[int]]: + """PDG ids of the external particles for every flavor of the + integrand, in the order of its flavor index.""" + 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, + beam_rapidity=self.process.beam_rapidity, + mirror_beams=self.mirror_beams, + ) + 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 has_mirrored_flavors(self) -> bool: + """Does any flavor of this subprocess stand for both beam orientations? + Resolved the same way build_integrands resolves flavor_mirror.""" + 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"]] + if flav["mirror"]: + return True + return False + + 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, + beam_rapidity=self.process.beam_rapidity, + mirror_beams=self.mirror_beams, + ) + 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, random_dim: int): + """Choose the MadNIS settings of this subprocess from the relative + spread rsd of the survey. random_dim is the smallest number of random + numbers of its channels: a normalizing flow needs at least two, so a + one-dimensional phase space (a 2 -> 1 process) never uses MadNIS.""" + 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 + self.madnis_supported = random_dim >= 2 + if not self.madnis_supported and self.madnis_settings["enable"]: + logger.warning( + "MadNIS disabled for subprocess %s: its phase space has %d " + "random dimension(s), a normalizing flow needs at least 2", + self.subproc_id, random_dim, + ) + self.madnis_settings["enable"] = False + + 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, + mapping_chan_weights=phasespace.mapping_chan_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 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: + # 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, + ) + 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( + 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=self.has_weighting_steps(), + ) + ) + partial_weights = self.process.systematics_enabled() + madnis_args = self.process.run_card["madnis"] + custom_bias = self.process.custom_function("bias", self.particle_count, flavors) + # With an MLM history picked per diagram, the cross section at the + # scales of that history is the first matrix-element evaluation rescaled + # by alpha_s and PDF ratios when every diagram has the same power of + # alpha_s (-1: mixed powers, evaluate the matrix element again) + mlm_history_qcd_power = int(self.meta.get("qcd_power", -1)) + 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, + None if self.process.leptonic else self.process.pdf_grid2, + phasespace.mapping_chan_weights, + mlm_history_qcd_power, + )) + #print(integrands[1].function()) + #for i in integrands: print(i.function()) + return integrands diff --git a/madgraph/iolibs/template_files/mg7/systematics_summary.py b/madgraph/mg7/systematics_summary.py similarity index 100% rename from madgraph/iolibs/template_files/mg7/systematics_summary.py rename to madgraph/mg7/systematics_summary.py diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index eaa49f8703..192df15a2c 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -7164,6 +7164,31 @@ def default_setup(self): self.add_toml_param('custom_functions', 'bias', "", comment="custom event bias") + # ---------------------------- [surrogate] --------------------- + # a fast surrogate of the matrix element for the integrand, followed by + # a second pass with the full matrix element; "" disables it + self.add_toml_param('surrogate', 'function', "", + comment="surrogate matrix element, a .so library or a .py file") + self.add_toml_param('surrogate', 'mode', "unweight", allowed=["unweight", "reweight"], + comment="unweight or reweight the events with the full matrix element") + self.add_toml_param('surrogate', 'first_pass_max_overweight_truncation', 0.01, + comment="max_overweight_truncation of the unweighting with the surrogate") + self.add_toml_param('surrogate', 'first_pass_freeze_max_weight_after', 10000, + comment="freeze_max_weight_after of the unweighting with the surrogate") + self.add_toml_param('surrogate', 'param_card', "", + comment="param card for a .so surrogate; \"\" uses Cards/param_card.dat") + # ---------------------------- [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, + gridpack=True, comment="max_overweight_truncation of the leading-colour unweighting") + self.add_toml_param('amplicol', 'first_pass_freeze_max_weight_after', 10000, + gridpack=True, comment="freeze_max_weight_after of the leading-colour unweighting") + # ----------------------------- [madnis] ----------------------- self.add_toml_param('madnis', 'enable', False, allowed=["auto", True, False], auto=True, comment="train MadNIS networks to improve the sampling; auto decides by process") @@ -8315,9 +8340,11 @@ def write_gridpack_card(self, output_file): # multiplicities, madevent's default_mlm_xqcut = 30.0 - def create_default_for_process(self, proc_characteristic, history, proc_def): + def create_default_for_process(self, proc_characteristic, history, proc_def, + amplicol=False): """Adjust default values depending on the process, then read the - site/user default file (``input/default_run_card_mg7.toml``).""" + site/user default file (``input/default_run_card_mg7.toml``). + ``amplicol``: the output has [amplicol] processes, which MLM does not support.""" # collect the initial-state (beam) particle ids beam_id = set() @@ -8365,7 +8392,7 @@ def create_default_for_process(self, proc_characteristic, history, proc_def): # drjl = 0). The shower then does the jet matching (is_mlm_merged), and # auto_ptj_mjj ties the jet cuts to xqcut. if beam_id & hadronic and not self['beam']['leptonic'] and \ - mlm_matching_needed(proc_def): + not amplicol and mlm_matching_needed(proc_def): self['beam']['dynamical_scale_choice'] = 'mlm' self['phasespace']['xqcut'] = self.default_mlm_xqcut cuts = self.dynamic_sections['cuts'] @@ -8497,7 +8524,7 @@ def _parameter_value(model, name): # MadGraph gives a leg whose flavours were merged one of these codes; the # mg7 runtime resolves them to the same representatives (clean_pids / - # _MERGED_PID_REPRESENTATIVE in template_files/mg7/launch.py), so the + # _MERGED_PID_REPRESENTATIVE in madgraph/mg7/launch.py), so the # observables have to be written in terms of those. merged_pid_representative = {81: 1, 82: 11, 83: 12} diff --git a/madmatrix/output.py b/madmatrix/output.py index 27ef9a6c35..5b6c13e5fc 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 diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index c7989a423f..8ac9d60c91 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -210,6 +210,8 @@ add_library( src/driver/context.cpp src/driver/event_generator.cpp src/driver/channel_generator.cpp + src/driver/event_utils.cpp + src/driver/event_stream.cpp src/driver/generator_data.cpp src/driver/status_file.cpp src/driver/thread_pool.cpp @@ -223,6 +225,7 @@ add_library( src/driver/event_histograms.cpp src/driver/madnis_training.cpp src/driver/random.cpp + src/driver/weighting_pipeline.cpp src/compgraphs/type.cpp src/compgraphs/function.cpp src/compgraphs/instruction.cpp @@ -260,6 +263,7 @@ add_library( src/phasespace/mlm_clustering.cpp src/phasespace/cross_section.cpp src/phasespace/scale.cpp + src/phasespace/weighting_step.cpp include/madspace/util.hpp include/madspace/driver.hpp include/madspace/driver/adam_optimizer.hpp @@ -279,6 +283,7 @@ add_library( include/madspace/driver/logger.hpp include/madspace/driver/madnis_training.hpp include/madspace/driver/random.hpp + include/madspace/driver/weighting_pipeline.hpp include/madspace/compgraphs.hpp include/madspace/compgraphs/type.hpp include/madspace/compgraphs/function.hpp @@ -318,6 +323,7 @@ add_library( include/madspace/phasespace/mlm_clustering.hpp include/madspace/phasespace/cross_section.hpp include/madspace/phasespace/scale.hpp + include/madspace/phasespace/weighting_step.hpp include/madspace/mixmax/mixmax.hpp include/madspace/mixmax/mixmax_skip_N17.c ) diff --git a/madspace/generate_pyi.py b/madspace/generate_pyi.py index b83c686e40..79c664b41d 100644 --- a/madspace/generate_pyi.py +++ b/madspace/generate_pyi.py @@ -57,6 +57,7 @@ def main(): "--enum-class-locations=Activation:MLP.Activation", "--enum-class-locations=CutMode:Cuts.CutMode", "--enum-class-locations=LRSchedule:AdamOptimizer.LRSchedule", + "--enum-class-locations=Mode:WeightingStep.Mode", "--enum-class-locations=JetScaleScheme:MLMClustering.JetScaleScheme", "--enum-class-locations=ScaleScheme:MLMClustering.ScaleScheme", "--enum-class-locations=PartonLineScheme:MLMClustering.PartonLineScheme", diff --git a/madspace/include/madspace/driver.hpp b/madspace/include/madspace/driver.hpp index 948427276c..74142e3ff4 100644 --- a/madspace/include/madspace/driver.hpp +++ b/madspace/include/madspace/driver.hpp @@ -6,6 +6,8 @@ #include "driver/discrete_optimizer.hpp" #include "driver/event_generator.hpp" #include "driver/event_histograms.hpp" +#include "driver/event_stream.hpp" +#include "driver/event_utils.hpp" #include "driver/format.hpp" #include "driver/io.hpp" #include "driver/lhe_output.hpp" @@ -17,3 +19,4 @@ #include "driver/tensor.hpp" #include "driver/thread_pool.hpp" #include "driver/vegas_optimizer.hpp" +#include "driver/weighting_pipeline.hpp" diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index b8476a1e32..6a7fdf6d4e 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -20,6 +20,35 @@ namespace madspace { +/// Events of one job that passed all re- and unweighting steps, see +/// @ref WeightingPipeline. +struct StepEvents { + /// Index of the channel. + std::size_t channel_index; + /// Id of the generation job of the events. + std::size_t job_id; + /// The events on the CPU, as returned by the unweighter. + TensorVec events; + /// The weights of the first step. + std::vector weights; + /// The maximum weight of the first step they were unweighted with. + double max_weight; + /// Per step, the ratios of the matrix elements. + std::vector> ratios; + /// Per step, the maximum weight they were unweighted with. + std::vector max_weights; + /// The biases of the events, empty without. + std::vector biases; + /// Number of points generated for these events. + std::size_t generated_count; +}; + +/// Sums of the final weights over the generated points. +struct StepIntegral { + std::size_t count = 0; + double sum = 0., sum2 = 0., abs_sum = 0., abs_sum2 = 0.; +}; + /** * Event generation and integration for one integration channel. * @@ -53,8 +82,10 @@ class ChannelEventGenerator { /** * @param contexts One context per device to run on. * @param integrand The channel's compiled integrand. - * @param event_file Path of the weighted event file. - * @param weight_file Path of the max-weight tracking file. + * @param event_file Path of the weighted event file; empty for no + * file (events are only kept in memory). + * @param weight_file Path of the max-weight tracking file; empty + * for no file. * @param config Generator configuration. * @param subprocess_index Index of the owning subprocess. * @param name Human-readable channel name. @@ -87,12 +118,17 @@ class ChannelEventGenerator { std::size_t count_requested_opt() const { return _count_requested_opt; } /// The filled observable histograms, if any were requested. const std::vector& histograms() const { return _histograms; } - /// The weighted event file. - EventFile& event_file() { return _event_file; } - /// The max-weight tracking file. - EventFile& weight_file() { return _weight_file; } + /// The weighted event file. Throws if the channel was created without one. + EventFile& event_file(); + /// The max-weight tracking file. Throws if the channel was created without + /// one. + EventFile& weight_file(); + /// Number of particles per event. + std::size_t particle_count() const { return _particle_count; } /// Current maximum-weight estimate, used for unweighting. double max_weight() const { return _max_weight; } + /// Whether the maximum weight was fixed with @ref set_fixed_max_weight. + bool max_weight_fixed() const { return _max_weight_fixed; } /// Use `max_weight` for unweighting from now on, instead of estimating it /// from the generated events. Must be positive. void set_fixed_max_weight(double max_weight); @@ -105,7 +141,7 @@ class ChannelEventGenerator { void set_integral_prior(double mean, double error); /// Whether this run has too few samples for an integral estimate of its own, so /// that @ref EventGenerator uses the priors instead. - bool uses_integral_prior() const; + bool uses_integral_prior(bool final_step = true) const; /// The prior of the signed integral as (mean, standard error), or of the /// absolute one if no signed prior was set. std::pair integral_prior() const; @@ -115,8 +151,9 @@ class ChannelEventGenerator { } /// The integral of the absolute weights: the mean of this run, combined /// with the prior (if any) weighted by the inverse variances. Used to share - /// the events between channels. - double abs_integral_estimate() const; + /// the events between channels. With `final_step`, the integral after the + /// weighting steps, see @ref result_integral, otherwise of the first step. + double abs_integral_estimate(bool final_step = true) const; /// Current generation batch size. std::size_t batch_size() const { return _batch_size; } /// Whether a VEGAS grid or channel-weight optimization pass is still @@ -137,6 +174,134 @@ class ChannelEventGenerator { int particle_layout_extra_flags() const { return _particle_layout_extra_flags; } /// Layout of @ref event_file. const DataLayout& event_file_layout() const { return _event_file_layout; } + /// Layout of @ref weight_file. + const DataLayout& weight_layout() const { + return _step_settings.empty() ? weight_file_layout : step_weight_file_layout; + } + /// Records per event in @ref weight_file, none without steps. + std::size_t weight_record_count() const { + return _step_settings.empty() ? 0 : _step_settings.size() + 1; + } + /// Types of the unweighted events, the integrand outputs before `random`. + const NamedVector& event_types() const { return _event_types; } + /// Whether the matrix element of the integrand returns the color, helicity + /// and diagram index. + const std::array& provided_indices() const { return _provided_indices; } + + /// Integral of the final weights: after the last re- or unweighting step, + /// or the weights of the integrand without steps. + struct IntegralEstimate { + double mean; + double variance; + double mean_abs; + double variance_abs; + std::size_t count; + }; + /// See @ref IntegralEstimate. With steps, but no events through all of + /// them yet, the integral of the first step. + IntegralEstimate result_integral() const; + /// The integral of the first step, see @ref IntegralEstimate. + IntegralEstimate first_step_integral() const; + + /// Set the re- and unweighting steps after the unweighting of this channel, + /// see @ref WeightingStep. Their events are generated by the + /// @ref EventGenerator, before any events of this channel. + void set_weighting_steps(const std::vector& steps); + /// The re- and unweighting steps. + const std::vector& weighting_steps() const { return _step_settings; } + /// Replace @ref GeneratorConfig::max_overweight_truncation and + /// @ref GeneratorConfig::freeze_max_weight_after of the first unweighting, + /// where set. + void set_first_pass_settings( + std::optional max_overweight_truncation, + std::optional freeze_max_weight_after + ); + /// The truncation of the first unweighting, see + /// @ref GeneratorConfig::max_overweight_truncation. + double max_overweight_truncation() const { + return _config.max_overweight_truncation; + } + /// The freeze count of the first unweighting, see + /// @ref GeneratorConfig::freeze_max_weight_after. + std::size_t freeze_max_weight_after() const { + return _config.freeze_max_weight_after; + } + /// The freeze count set with @ref set_first_pass_settings, if any. + std::optional first_pass_freeze_max_weight_after() const { + return _first_pass_freeze_max_weight_after; + } + /// Maximum weight of step `step`, counted from 1. + double step_max_weight(std::size_t step) const { + return _step_states.at(step - 1).max_weight; + } + /// Whether the maximum weights of all unweighting steps are fixed. + bool step_max_weights_fixed() const { + for (auto [state, settings] : zip(_step_states, _step_settings)) { + if (settings.unweight && !state.max_weight_fixed) { + return false; + } + } + return true; + } + /// Fix the maximum weight of step `step`, like @ref set_fixed_max_weight. + void set_fixed_step_max_weight(std::size_t step, double max_weight); + /// Update the maximum weight of step `step` from the values it unweights, + /// with the truncation budget for `count_target` events. + void update_step_max_weight( + std::size_t step, const std::vector& values, std::size_t count_target + ); + /// Count `count` events given the weight 0 by a step, see + /// @ref GeneratorStatus::count_invalid_step_values. + void add_invalid_step_values(std::size_t count) { + _status.count_invalid_step_values += count; + } + /// Apply the truncation budget of the current target count to the maxima + /// of the steps, which are updated with the target count of each job. Like + /// @ref set_target_count for the first maximum. + void apply_step_truncation_budgets(); + /// Count `count` events accepted by step `step`, which passed the earlier + /// unweightings with the maxima `max_weight` (first step) and + /// `max_weights` (the steps before). + void add_step_accepted( + std::size_t step, + double count, + double max_weight, + const std::vector& max_weights + ); + /// Estimated number of events accepted by the first unweighting and each + /// step so far; without steps, just @ref GeneratorStatus::count_unweighted. + std::vector step_accepted_counts() const { + if (_step_settings.empty()) { + return {_status.count_unweighted}; + } + std::vector counts{_first_step_count}; + for (auto& state : _step_states) { + counts.push_back(state.count_accepted); + } + return counts; + } + /// Count `count` events accepted by the first unweighting with + /// `job_max_weight`, with steps. + void add_first_step_events(std::size_t count, double job_max_weight); + /// Incremented whenever the events are cleared, so that events still in the + /// steps can be discarded. + std::size_t event_epoch() const { return _event_epoch; } + /// Write events that passed all steps and add them to the integral. + void write_step_events(const StepEvents& step_events); + /// The contribution of `step_events` to the integral of the final weights. + StepIntegral step_integral(const StepEvents& step_events) const; + /// Add to the integral of the final weights. + void add_step_integral(const StepIntegral& integral); + /// Pack `step_events` into an event buffer and a weight buffer (layouts + /// @ref event_file_layout and @ref weight_layout). + void fill_step_buffers( + const StepEvents& step_events, + EventBuffer& event_buffer, + EventBuffer& weight_buffer + ) const; + /// Signed weight of event `index` of a buffer read from @ref weight_file, + /// in units of the maximum weights, 0 for a rejected event. + double unweighted_weight(EventBuffer& weight_buffer, std::size_t index) const; /// Run a final unweighting pass over @ref weight_file into @ref /// event_file, using `rand_gen`. @@ -170,11 +335,20 @@ class ChannelEventGenerator { std::size_t next_vegas_batch_size(); /// Discard the accumulated event and weight files. void clear_events(); - /// Update @ref max_weight from a batch of `weights`. - void update_max_weight(Tensor weights); + /// Update @ref max_weight from a batch of `weights`, which are unweighted + /// later if `unweight` is true. + void update_max_weight(Tensor weights, bool unweight); /// Append `unweighted_events` to @ref event_file, recording /// `job_max_weight`. void write_events(const TensorVec& unweighted_events, double job_max_weight); + /// Pack `unweighted_events` into `event_buffer` (layout @ref + /// event_file_layout) and their weights into `weight_buffer` (layout + /// `weight_file_layout`), resizing both. + void fill_event_buffers( + const TensorVec& unweighted_events, + EventBuffer& event_buffer, + EventBuffer& weight_buffer + ) const; /// Serialize this generator's state to `file_name`; see @ref load. void save(const std::string& file_name) const; /// The JSON text that @ref save writes; see @ref load_json. With @@ -185,7 +359,20 @@ class ChannelEventGenerator { std::string to_json(bool include_estimates = false) const; private: + struct StepState { + double max_weight = 0.; + bool max_weight_fixed = false; + std::vector large_weights; + double count_accepted = 0.; + // like _freeze_count, the expected accepted events of this step alone + double freeze_count = 0.; + }; + void apply_truncation_budget(); + void apply_step_truncation_budget(std::size_t step, std::size_t count_target); + void update_status_integrals(); + void unweight_step_file(MixMaxRandom& rand_gen); + StepChainResult step_chain(EventBuffer& weight_buffer, std::size_t index) const; ChannelEventGenerator( const std::vector& contexts, std::size_t particle_count, @@ -213,6 +400,7 @@ class ChannelEventGenerator { RuntimePtr vegas_histogram = nullptr; RuntimePtr discrete_histogram = nullptr; RuntimePtr observable_histograms = nullptr; + RuntimePtr step_random = nullptr; }; struct FieldIndices { @@ -226,6 +414,7 @@ class ChannelEventGenerator { std::array ladder; int subprocess_index; int bias; + int matrix_element; int random, rest; }; @@ -236,8 +425,9 @@ class ChannelEventGenerator { int _event_layout_extra_flags; int _particle_layout_extra_flags; DataLayout _event_file_layout; - EventFile _event_file; - EventFile _weight_file; + std::optional _event_file; + std::optional _weight_file; + std::string _weight_file_path; std::optional _vegas_optimizer; std::optional _discrete_optimizer; std::size_t _batch_size; @@ -250,11 +440,6 @@ class ChannelEventGenerator { RunningIntegral _cross_section; RunningIntegral _abs_cross_section; double _max_weight = 0.; - // Expected number of accepted events, sum of min(|w| / max_weight, 1) over the - // committed generate stages, rescaled like count_unweighted when the max - // weight rises. Decides when the max weight is frozen: unlike count_unweighted - // it only depends on the generate stages, which commit in a fixed order. - double _expected_unweighted = 0.; bool _max_weight_fixed = false; // mean and standard error std::optional> _abs_integral_prior; @@ -277,6 +462,18 @@ class ChannelEventGenerator { std::vector _histograms; std::unordered_set _used_globals; FieldIndices _field_indices; + NamedVector _event_types; + std::array _provided_indices{true, true, true}; + std::vector _step_settings; + std::vector _step_states; + std::vector _loaded_step_max_weights; + double _first_step_count = 0.; + // expected number of unweighted events, counted in the commit order of the + // jobs to freeze the maximum weight independently of the thread timing + double _freeze_count = 0.; + StepIntegral _final_integral; + std::size_t _event_epoch = 0; + std::optional _first_pass_freeze_max_weight_after; friend void to_json(nlohmann::json& j, const ChannelEventGenerator& channel); }; diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 53f4e7f10e..b0079843ec 100644 --- a/madspace/include/madspace/driver/event_generator.hpp +++ b/madspace/include/madspace/driver/event_generator.hpp @@ -20,6 +20,7 @@ #include "madspace/driver/status_file.hpp" #include "madspace/driver/systematics.hpp" #include "madspace/driver/vegas_optimizer.hpp" +#include "madspace/driver/weighting_pipeline.hpp" #include "madspace/phasespace.hpp" namespace madspace { @@ -54,13 +55,18 @@ class EventGenerator { * @param status_file Optional status file progress is periodically * written to. * @param config Generator configuration. + * @param groups Groups of channels with re- and unweighting steps + * after their first unweighting, see @ref WeightingStep. + * The integrals and event counts are then those after + * the last step. */ EventGenerator( const std::vector& contexts, const std::vector>& channels, std::uint64_t seed, std::shared_ptr status_file = nullptr, - const GeneratorConfig& config = default_config + const GeneratorConfig& config = default_config, + const std::vector& groups = {} ); EventGenerator(EventGenerator&&) = default; EventGenerator& operator=(EventGenerator&&) = default; @@ -153,6 +159,8 @@ class EventGenerator { std::vector _channel_optimizing; std::vector _channel_integral_fractions; std::vector _context_job_counts; + // invalid step values per channel that were already warned about + std::vector _warned_invalid_step_values; // True while a channel has a steady-state batch dispatched but not yet fully // committed; keeps next_batch_event_count() from double-counting in-flight work. std::vector _channel_batch_pending; @@ -184,6 +192,7 @@ class EventGenerator { // dispatch, drained with priority by start_jobs(). std::vector> _context_unweight_queue; ResultQueue _result_queue; + std::unique_ptr _pipeline; // Base seed for reproducible event generation. std::uint64_t _seed; @@ -229,25 +238,28 @@ class EventGenerator { void reset_jobs(); void update_integral(); void update_integral_status(); - void update_integral_fractions(); + // With `final_step`, the integrals after the weighting steps are used. They + // only change in a fixed order once the steps are drained. + void update_integral_fractions(bool final_step = true); void update_counts(); void reset_start_time(); void add_timing_data(const std::string& key); void unweight_all(); - std::tuple, std::size_t, double> init_combine(); + std::tuple, std::size_t, std::vector> + init_combine(); void read_and_combine( std::vector& channel_data, EventBuffer& buffer, - double norm_factor, - MixMaxRandom& rand_gen - ); - void fill_lhe_event( - LHECompleter& lhe_completer, - LHEEvent& lhe_event, - EventBuffer& buffer, - std::size_t event_index, + const std::vector& norm_factors, MixMaxRandom& rand_gen ); + // run the re- and unweighting steps of all collected events + void drain_pipeline(); + // unweighting efficiency of each pass, the first one before and after cuts + std::tuple> pass_efficiencies() const; + // start the jobs of the steps whose batches are full, and the partial ones + // once the generation jobs leave threads idle + std::size_t start_pipeline_jobs(); // Layout of the combined events as read from the channel files DataLayout combined_layout() const; // Layout of the written events: the reweighting inputs are dropped unless @@ -293,6 +305,7 @@ class EventGenerator { std::size_t iter ); + void warn_invalid_step_values(); void print_gen_init(); void print_gen_update(bool done); void print_gen_update_pretty(bool done); diff --git a/madspace/include/madspace/driver/event_stream.hpp b/madspace/include/madspace/driver/event_stream.hpp new file mode 100644 index 0000000000..535154a2f4 --- /dev/null +++ b/madspace/include/madspace/driver/event_stream.hpp @@ -0,0 +1,185 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "madspace/driver/channel_generator.hpp" +#include "madspace/driver/context.hpp" +#include "madspace/driver/generator_data.hpp" +#include "madspace/driver/io.hpp" +#include "madspace/driver/lhe_output.hpp" +#include "madspace/driver/random.hpp" +#include "madspace/driver/thread_pool.hpp" +#include "madspace/driver/weighting_pipeline.hpp" + +namespace madspace { + +/** + * Endless stream of unweighted events, generated in memory. + * + * Every channel needs a fixed maximum weight (see @ref + * ChannelEventGenerator::set_fixed_max_weight), so that an event is final as + * soon as it is unweighted; channels without one are skipped. With re- and + * unweighting steps (see @ref WeightingStep), the maximum weights of the steps + * have to be fixed as well, and an event is final after the last step. A coordinator + * thread keeps generation jobs running on the context's thread pool and + * buffers each channel's unweighted events, in launch order, until they are + * consumed. @ref next_events picks the channel of every event at random, + * proportional to its current cross-section estimate, and completes the + * events with an @ref LHECompleter. + * + * The integral estimates only include the jobs whose events were requested + * so far, so the event sequence is a pure function of the seed: it depends + * neither on thread timing nor on how the requests are split into batches. + * + * Event weights are in units of the maximum weight: +-1, except for + * overweight events. After reweighting steps, they also contain the ratios of + * the reweighting steps. + */ +class EventStream { +public: + /// Install `func`, called periodically while @ref next_events waits for + /// events, to check for a requested abort. + static void set_abort_check_function(std::function func) { + _abort_check_function = func; + } + + /** + * Start the coordinator thread and the first generation jobs. + * + * @param context Context to generate on. + * @param channels The channels, each created with `context` as its + * only context. + * @param seed Run seed every random stream is derived from. + * @param lhe_completer Completes the events into LHE events. + * @param config Generator configuration; the batch sizes and cut + * settings are used. + * @param groups Groups of channels with re- and unweighting steps, + * see @ref EventGenerator. + */ + EventStream( + ContextPtr context, + const std::vector>& channels, + std::uint64_t seed, + const LHECompleter& lhe_completer, + const GeneratorConfig& config, + const std::vector& groups = {} + ); + ~EventStream(); + EventStream(const EventStream&) = delete; + EventStream& operator=(const EventStream&) = delete; + + /// The next `count` events of the stream. Blocks until they are available. + std::vector next_events(std::size_t count); + /// Stop the coordinator thread and discard the running jobs. Called by + /// the destructor. + void close(); + /// Combined status over the channels with fixed maximum weights, the others + /// return no events. `count_unweighted` is the number of events returned + /// so far. + GeneratorStatus status() const; + /// Per-channel status; `count_unweighted` is the number of events of the + /// channel returned so far. + std::vector channel_status() const; + /// Current probability of each channel to be picked for the next event. + std::vector channel_probabilities() const; + /// Number of events returned so far. + std::size_t event_count() const; + /// Largest particle count of the returned LHE events. + std::size_t max_particle_count() const { + return _lhe_completer.max_particle_count(); + } + +private: + // A job whose events are buffered, waiting to be consumed. + struct ClosedJob { + GeneratorBatchJob job; + EventBuffer events; + EventBuffer weights; + bool opened = false; + std::size_t next_event = 0; + // with steps: the integral after the steps and the sum of the + // absolute event weights + StepIntegral step_integral; + double abs_weight_sum = 0.; + }; + struct ChannelState { + bool active = false; + // job ids in launch order, waiting to be closed + std::deque launch_order; + std::map finished; + std::deque closed; + std::size_t in_flight = 0; + std::size_t buffered_events = 0; + std::size_t unopened_jobs = 0; + // statistics of the completed and of the opened jobs + std::size_t completed_jobs = 0; + std::size_t completed_events = 0; + std::size_t opened_jobs = 0; + double opened_abs_sum = 0.; + double opened_capped_sum = 0.; + std::size_t returned_events = 0; + }; + // thrown inside the coordinator to stop it + struct StopRequest {}; + + inline static std::function _abort_check_function = [] {}; + + ContextPtr _context; + std::vector> _channels; + std::uint64_t _seed; + LHECompleter _lhe_completer; + GeneratorConfig _config; + DataLayout _combined_layout; + MixMaxRandom _select_rng; + MixMaxRandom _lhe_rng; + + // guards everything below, shared by the consumer and the coordinator + mutable std::mutex _mutex; + std::condition_variable _coordinator_cv; + std::condition_variable _consumer_cv; + std::vector _channel_states; + std::vector _cum_probabilities; + bool _probabilities_dirty = true; + bool _dispatch_requested = false; + bool _stopping = false; + std::exception_ptr _error; + std::size_t _lookahead; + std::size_t _event_count = 0; + bool _started = false; + + // only used by the consumer, which next_events() serializes + std::mutex _consumer_mutex; + // only used by the coordinator thread + std::unordered_map _running_jobs; + std::size_t _results_pending = 0; + std::size_t _job_id = 0; + ResultQueue _result_queue; + std::unique_ptr _pipeline; + // jobs whose events are in the steps + std::unordered_map _step_jobs; + std::thread _coordinator; + + void run_coordinator(); + void dispatch_jobs(); + void handle_result(std::size_t job_id); + void + close_job(std::size_t job_id, std::size_t channel_index, ClosedJob&& closed_job); + void close_step_events(StepEvents& step_events); + double expected_job_events(std::size_t channel_index) const; + void update_probabilities(); + std::size_t select_channel(); + ClosedJob& next_job(std::size_t channel_index, std::unique_lock& lock); + void open_job(std::size_t channel_index, ClosedJob& closed_job); + void wait_for_events(std::size_t channel_index, std::unique_lock& lock); +}; + +} // namespace madspace diff --git a/madspace/include/madspace/driver/event_utils.hpp b/madspace/include/madspace/driver/event_utils.hpp new file mode 100644 index 0000000000..380666c4cd --- /dev/null +++ b/madspace/include/madspace/driver/event_utils.hpp @@ -0,0 +1,86 @@ +#pragma once + +#include +#include +#include +#include + +#include "madspace/driver/channel_generator.hpp" +#include "madspace/driver/generator_data.hpp" +#include "madspace/driver/io.hpp" +#include "madspace/driver/lhe_output.hpp" +#include "madspace/driver/random.hpp" + +namespace madspace { + +/// Weight of an unweighted event in units of its channel's maximum weight: +/// +-1, or larger in magnitude for an overweight event. +inline double unweighted_event_weight(double weight, double max_weight) { + return std::copysign(std::max(1., std::abs(weight) / max_weight), weight); +} + +/// Layout of combined events: weight, subprocess index and event data, plus +/// the extra fields (`extra_event_flags`, `extra_particle_flags`) the channels +/// carry, except for the bias, which is divided out of the weight. +DataLayout combined_event_layout(int extra_event_flags, int extra_particle_flags); + +/// The MLM alpha_s vertex scale and pdf ladder particle flags of any of +/// `channels`: a merged run mixes 2 -> 1 samples without vertices with the jet +/// samples, and channel 0 may be either. +int mlm_particle_flags( + const std::vector>& channels +); + +/** + * Copy event `in_index` of a channel's event buffer to event `out_index` of a + * combined buffer (layout @ref combined_event_layout), padding the momenta of + * missing particles with zeros. + * + * @param weight Weight to store for the event, divided by the bias + * if the channel records one. + * @param subprocess_index Stored if the channel records no per-event + * subprocess index. + * @param extra_flags The channel's extra `EventRecord` layout flags. + * @param in_particle_flags The channel's extra `ParticleRecord` layout flags. + * @param out_particle_flags The extra `ParticleRecord` layout flags of the + * combined buffer. + */ +void copy_channel_event( + EventBuffer& in_buffer, + std::size_t in_index, + EventBuffer& out_buffer, + std::size_t out_index, + double weight, + int subprocess_index, + int extra_flags, + int in_particle_flags, + int out_particle_flags +); + +/// Build `lhe_event` from event `event_index` of a combined buffer and +/// complete it with `lhe_completer`. `extra_event_flags` and +/// `extra_particle_flags` are the extra layout flags of the buffer's channels. +void fill_lhe_event( + LHECompleter& lhe_completer, + LHEEvent& lhe_event, + EventBuffer& buffer, + std::size_t event_index, + MixMaxRandom& rand_gen, + int extra_event_flags, + int extra_particle_flags +); + +/// Write the particles of `lhe_event` to event `event_index` of a buffer with +/// LHE particle layout, padding the remaining slots with empty particles. +void write_lhe_particles( + const LHEEvent& lhe_event, EventBuffer& buffer, std::size_t event_index +); + +/// Sum the integrals and event counts of `channels` into `status`. Throws if +/// an integral is not finite. +void sum_channel_status( + const std::vector>& channels, + GeneratorStatus& status +); + +} // namespace madspace diff --git a/madspace/include/madspace/driver/generator_data.hpp b/madspace/include/madspace/driver/generator_data.hpp index c089958006..02e4c52e82 100644 --- a/madspace/include/madspace/driver/generator_data.hpp +++ b/madspace/include/madspace/driver/generator_data.hpp @@ -112,6 +112,57 @@ struct GeneratorConfig { double batch_overshoot_sigma = 1.0; }; +/// Settings of a re- or unweighting step after the first unweighting, as far +/// as the channels need them; see @ref WeightingStep. +struct StepSettings { + /// Unweighting step if true, reweighting step otherwise. + bool unweight = true; + /// See @ref GeneratorConfig::max_overweight_truncation. + double max_overweight_truncation = 0.01; + /// See @ref GeneratorConfig::freeze_max_weight_after. + std::size_t freeze_max_weight_after = 10000; + /// Whether the step returns the ratio instead of the matrix element. + bool ratio = false; +}; + +/// Result of @ref evaluate_step_chain. +struct StepChainResult { + /// Probability that the event is accepted by all steps. + double acceptance; + /// Signed weight of the event after all steps, in units of the product of + /// the maximum weights of the unweighting steps. + double weight; +}; + +/** + * Acceptance probability and weight of an event after the first unweighting and + * the re- and unweighting steps that follow. + * + * The first step accepts the event with probability + * @f$\min(1, |w_0| / m_0)@f$ and carries the overweight factor + * @f$o_0 = \max(1, |w_0| / m_0)@f$. Each reweighting step multiplies the weight + * by its ratio @f$r_k@f$. Each unweighting step unweights + * @f$u_k = o \prod r@f$, the overweight factor of the last unweighting step + * times the ratios since, with probability @f$\min(1, u_k / m_k)@f$, and carries + * @f$o_k = \max(1, u_k / m_k)@f$. Raising any maximum weight only lowers the + * acceptance, so the events can be unweighted again with larger maxima. + * + * @param weight0 The stored weight of the first step, @f$\max(|w_0|, m)@f$ + * with the sign of @f$w_0@f$, for a maximum @f$m \le m_0@f$. + * @param max_weight0 The maximum weight @f$m_0@f$ of the first step. + * @param ratios The ratios @f$r_k@f$ of the steps. + * @param max_weights The maximum weights @f$m_k@f$ of the steps, ignored for + * reweighting steps. + * @param steps The steps. + */ +StepChainResult evaluate_step_chain( + double weight0, + double max_weight0, + const double* ratios, + const double* max_weights, + const std::vector& steps +); + /** * Number of events to schedule in the next generation batch. * @@ -200,6 +251,9 @@ struct GeneratorStatus { bool optimized; /// Whether generation of this subprocess is complete. bool done; + /// Events given the weight 0 by a weighting step, because their matrix + /// element ratio was not finite or the earlier matrix element was 0. + std::size_t count_invalid_step_values = 0; }; /// One binned observable, accumulated during event generation; see @ref @@ -257,6 +311,8 @@ struct GeneratorBatchJob { std::size_t context_index; std::size_t job_id; double max_weight; + // target count of the channel when max_weight was taken + std::size_t count_target = 0; // Top-level seed plus job identity, used to derive this job's DerivedSeed(s) // independently for generate vs unweight (see start_job()/submit_unweight_job()). std::optional rng_seed; @@ -271,6 +327,9 @@ struct GeneratorBatchJob { // Nominal batch size start_job() requested; the job generates more raw points // when it repeats batches to get enough of them past the cuts. std::size_t requested_event_count = 0; + // With re- and unweighting steps: unweighted_events stay on the device and + // these are their weights, matrix elements and biases on the CPU + TensorVec step_host; }; void to_json(nlohmann::json& j, const GeneratorStatus& status); diff --git a/madspace/include/madspace/driver/io.hpp b/madspace/include/madspace/driver/io.hpp index 4aca6da092..21d2601e49 100644 --- a/madspace/include/madspace/driver/io.hpp +++ b/madspace/include/madspace/driver/io.hpp @@ -180,6 +180,10 @@ class ParticleRecord { UnalignedRef ladder_info() { return &_data[_offsets[2] + 24]; } UnalignedRef ladder_lead_q() { return &_data[_offsets[2] + 32]; } + // weight files with re- and unweighting steps store one record per step + UnalignedRef step_ratio() { return &_data[0]; } + UnalignedRef step_max_weight() { return &_data[8]; } + void from_lhe_particle(const LHEParticle& particle, bool has_clustering = false) { lhe_pdg_id() = particle.pdg_id; lhe_status_code() = particle.status_code; @@ -391,6 +395,13 @@ inline const DataLayout weight_file_layout = DataLayout( ParticleRecord::layout(ParticleRecord::f_none) ); +/// Weight file of a channel with re- and unweighting steps: the weight of the +/// first step and, as "particles", the ratio and the maximum weight of each step. +inline const DataLayout step_weight_file_layout = DataLayout( + EventRecord::layout(EventRecord::f_weight), + {{"ratio", FieldLayout::f64, 0}, {"max_weight", FieldLayout::f64, 0}} +); + class EventBuffer { public: EventBuffer( diff --git a/madspace/include/madspace/driver/lhe_output.hpp b/madspace/include/madspace/driver/lhe_output.hpp index b224372195..4e65d1cfb0 100644 --- a/madspace/include/madspace/driver/lhe_output.hpp +++ b/madspace/include/madspace/driver/lhe_output.hpp @@ -189,7 +189,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; @@ -209,6 +210,10 @@ class LHECompleter { /// then by `Topology::Decay::flat_propagator_index`. Falls back to /// `Topology::Decay::pdg_id` if empty. nested_vector3 diagram_propagator_pdgs; + /// Relative weight of each PDG-id option of each flavor: an option is + /// written with probability proportional to its weight. Uniform if empty, + /// or for a flavor whose weights are empty. + nested_vector2 pdg_id_weights; }; /// @param subproc_args Bookkeeping for every subprocess, indexed like @@ -224,7 +229,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. @@ -269,6 +275,9 @@ class LHECompleter { std::vector _helicities; std::vector> _pdg_id_and_count; std::vector _pdg_ids; + // flavor slot -> offset of its cumulative option probabilities, if not uniform + std::unordered_map _pdg_id_weight_offsets; + std::vector _pdg_id_cumulative; std::unordered_map> _propagator_index_and_count; std::vector _propagators; diff --git a/madspace/include/madspace/driver/madnis_training.hpp b/madspace/include/madspace/driver/madnis_training.hpp index 8d274c0ab6..5ba75131ec 100644 --- a/madspace/include/madspace/driver/madnis_training.hpp +++ b/madspace/include/madspace/driver/madnis_training.hpp @@ -352,6 +352,7 @@ class MultiMadnisTraining { std::chrono::time_point _last_print_time; PrettyBox _pretty_box_upper; PrettyBox _pretty_box_lower; + nested_vector2 _pretty_subproc_rows; std::shared_ptr _status_file; nlohmann::json _trainings_status; }; diff --git a/madspace/include/madspace/driver/random.hpp b/madspace/include/madspace/driver/random.hpp index 821d054a25..e617ecd9f4 100644 --- a/madspace/include/madspace/driver/random.hpp +++ b/madspace/include/madspace/driver/random.hpp @@ -32,6 +32,9 @@ struct DerivedSeed { madnis_unweight, madnis_sample_buffer, global_init, + // per-event random numbers of the re- and unweighting steps; at most + // 16 types fit into the seed + weighting_step_random, }; std::array seed_parts; diff --git a/madspace/include/madspace/driver/weighting_pipeline.hpp b/madspace/include/madspace/driver/weighting_pipeline.hpp new file mode 100644 index 0000000000..da476a08f3 --- /dev/null +++ b/madspace/include/madspace/driver/weighting_pipeline.hpp @@ -0,0 +1,193 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "madspace/driver/channel_generator.hpp" +#include "madspace/driver/context.hpp" +#include "madspace/driver/generator_data.hpp" +#include "madspace/driver/thread_pool.hpp" +#include "madspace/phasespace/weighting_step.hpp" + +namespace madspace { + +/// A group of channels with the same re- and unweighting steps after their +/// first unweighting; see @ref WeightingStep. +struct WeightingGroup { + /// Indices of the channels in the @ref EventGenerator. + std::vector channel_indices; + /// The steps, in order. + std::vector steps; + /// If set, replaces @ref GeneratorConfig::max_overweight_truncation for + /// the first unweighting of the channels of the group. + std::optional first_pass_max_overweight_truncation; + /// If set, replaces @ref GeneratorConfig::freeze_max_weight_after for the + /// first unweighting of the channels of the group. + std::optional first_pass_freeze_max_weight_after; +}; + +/// JSON text of `groups`, for example to store them in a gridpack. The matrix +/// elements are stored by their index in the context. +std::string weighting_groups_to_json(const std::vector& groups); +/// The groups stored with @ref weighting_groups_to_json. The matrix elements +/// must be loaded into the contexts in the same order as before. +std::vector weighting_groups_from_json(const std::string& text); + +/** + * Runs the re- and unweighting steps of the channel groups of an + * @ref EventGenerator. + * + * The events that pass the first unweighting stay on the device. They are + * collected per group and context until a device batch is full, then the + * matrix element of the next step is evaluated on them. The decisions of the + * steps are made on the scheduling thread, in the order of the jobs of each + * channel, and the random numbers they use are drawn per event when the + * events are first unweighted. The events of the channels sharing a batch + * therefore don't influence each other, and the results don't depend on how + * the events are batched. + */ +class WeightingPipeline { +public: + /// Job ids of the pipeline have this bit set. + static constexpr std::size_t job_id_flag = std::size_t(1) << 62; + + /// Receives the events after the last step, in the job order of each + /// channel. + using Sink = std::function; + + /// Without `sink`, the events are written to the channel files with + /// @ref ChannelEventGenerator::write_step_events. + WeightingPipeline( + const std::vector& contexts, + const std::vector>& channels, + const std::vector& groups, + const GeneratorConfig& config, + Sink sink = {} + ); + /// Whether channel `channel_index` has steps. + bool contains(std::size_t channel_index) const { + return _channel_groups.at(channel_index) != -1; + } + /// Whether `job_id` belongs to the pipeline. + static bool owns_job(std::size_t job_id) { return job_id & job_id_flag; } + /// Take the unweighted events of `job`, in the job order of its channel. + void add_events(GeneratorBatchJob& job); + /// Start the jobs of full batches. With `flush`, all collected events are + /// split over the free threads of each context, where `busy_counts` are + /// the other jobs running on it. Returns the number of started jobs. + std::size_t start_jobs( + ResultQueue& result_queue, + bool flush, + const std::vector& busy_counts = {} + ); + /// Process the result of the finished job `job_id`. + void commit_job(std::size_t job_id); + /// Whether events are still collected or in jobs. + bool busy() const { return _chunk_count > 0; } + /// Drop all events, after a failed job. + void reset(); + +private: + // the unweighted events of one job of a channel + struct Chunk { + std::size_t channel_index; + std::size_t job_id; + std::size_t sequence; + std::size_t epoch; + std::size_t context_index; + std::size_t generated_count; + std::size_t count_target; + // on the device: rows offset to offset + size() of source, before the + // last decision, which keeps the rows in keep + std::shared_ptr source; + std::size_t offset = 0; + std::optional> keep; + TensorVec events; + // per event on the CPU + std::vector weights, matrix_elements, biases; + std::vector overweights, ratio_products; + std::vector> ratios; + double max_weight; + std::vector max_weights; + // results of the last job + std::vector step_matrix_elements, step_random; + TensorVec cpu_events; + + std::size_t size() const { return weights.size(); } + void clear(); + void compact(const std::vector& indices); + }; + using ChunkPtr = std::shared_ptr; + + struct ContextRuntimes { + RuntimePtr gather; + // by the number of concatenated parts, built when first needed + std::map concats; + std::vector steps; + }; + struct Group { + NamedVector types; + std::vector settings; + std::vector runtimes; + // [context][stage - 1], stage step_count + 1 copies the events to the CPU + std::vector>> buffers; + std::vector> buffer_sizes; + std::vector multiple_matrix_elements; + int momenta, alpha_s, flavor, me_index, color, helicity, diagram, random; + }; + struct Job { + std::size_t group; + std::size_t stage; + std::size_t context_index; + std::vector chunks; + }; + + void advance(ChunkPtr chunk, std::size_t stage); + void set_ready(ChunkPtr chunk, std::size_t stage); + void commit_stage(ChunkPtr chunk, std::size_t stage); + void decide(Chunk& chunk, std::size_t step); + void start_job( + ResultQueue& result_queue, + std::size_t group_index, + std::size_t context_index, + std::size_t stage, + std::vector chunks + ); + void run_job(Job& job); + void gather_events(Chunk& chunk, ContextRuntimes& runtimes); + // the rows of `chunks` from their shared `source`, concatenated + TensorVec gather_rows( + const TensorVec& source, + const std::vector& chunks, + std::size_t context_index, + ContextRuntimes& runtimes + ); + RuntimePtr& concat_runtime( + std::size_t group_index, std::size_t context_index, std::size_t part_count + ); + std::size_t batch_size(std::size_t context_index) const; + + std::vector _contexts; + std::vector> _channels; + GeneratorConfig _config; + std::vector _groups; + std::vector _channel_groups; + std::vector _channel_sequences; + // per channel and stage: the next sequence to commit, and the finished + // chunks waiting for their turn + std::vector> _next_sequences; + std::vector>> _ready; + std::unordered_map _jobs; + std::vector _context_job_counts; + std::mutex _concat_mutex; + std::size_t _job_counter = 0; + std::size_t _chunk_count = 0; + Sink _sink; +}; + +} // namespace madspace diff --git a/madspace/include/madspace/phasespace.hpp b/madspace/include/madspace/phasespace.hpp index dafca32363..1e27f9bf42 100644 --- a/madspace/include/madspace/phasespace.hpp +++ b/madspace/include/madspace/phasespace.hpp @@ -31,3 +31,4 @@ #include "phasespace/two_particle.hpp" #include "phasespace/unweighter.hpp" #include "phasespace/vegas.hpp" +#include "phasespace/weighting_step.hpp" diff --git a/madspace/include/madspace/phasespace/channel_weights.hpp b/madspace/include/madspace/phasespace/channel_weights.hpp index e058bad6d5..b699f948b5 100644 --- a/madspace/include/madspace/phasespace/channel_weights.hpp +++ b/madspace/include/madspace/phasespace/channel_weights.hpp @@ -1,6 +1,7 @@ #pragma once #include "madspace/phasespace/base.hpp" +#include "madspace/phasespace/phasespace.hpp" #include "madspace/phasespace/topology.hpp" #include "madspace/util.hpp" @@ -113,4 +114,82 @@ class SubchannelWeights : public FunctionGenerator { std::vector _subchannel_indices; }; +/** + * Local multi-channel weights from the sampling densities of the channel + * mappings. + * + * The weight of channel @f$c@f$ for an event of flavour @f$f@f$ is + * @f$\alpha_c(x) = m_{fc} g_c(x) / \sum_k m_{fk} g_k(x)@f$ [1, 2], where + * @f$g_k@f$ is the density with which the @ref PhaseSpaceMapping of channel + * @f$k@f$ samples the event, read off its inverse with uniformly sampled + * discrete choices, and @f$m_{fk}@f$ is 1 if channel @f$k@f$ integrates + * flavour @f$f@f$, else 0. The adaptive samplers do not enter, so the weights + * of one channel do not depend on the state of the others. A channel whose + * mapping cannot produce the event gets the weight 0. If no channel of the + * flavour has a valid density, the weights are uniform over these channels. + * + * Optionally, it also samples a diagram for the event record: diagram + * @f$d@f$ belongs to channel @f$c_d@f$ and is drawn with a probability + * proportional to @f$\alpha_{c_d}@f$ among the diagrams of the flavour. + * + * `batch` is the leading batch dimension. `c` indexes the channels. + * + * **Arguments** + * - `momenta` – `float`, shape `(batch, n, 4)` – the lab-frame momenta. + * - `x1`, `x2` – `float`, shape `(batch,)` – the momentum fractions returned + * by the mapping that generated the event. + * - `flavor` – `int`, shape `(batch,)` – the flavour index. + * - `mirror_index` – `int`, shape `(batch,)` – the orientation of the beams. + * Present only if the mappings are built with `mirror_beams`. + * + * **Returns** + * - `channel_weights` – `float`, shape `(batch, c)` – the normalized + * per-channel weights. + * - `diagram_index` – `int`, shape `(batch,)` – the sampled diagram. Present + * only if diagrams are given. + * + * **References** + * - [1] R. Kleiss, R. Pittau, "Weight optimization in multichannel Monte + * Carlo", https://arxiv.org/abs/hep-ph/9405257 + * - [2] R. Frederix, T. Vitos, "Leading-colour-based unweighted event + * generation for multi-parton tree-level processes", + * https://arxiv.org/abs/2409.12128 + */ +class MappingChannelWeights : public FunctionGenerator { +public: + /** + * @param mappings The mapping of each channel, without + * permutations. They take the same momenta. + * @param flavor_masks Per flavour, 1 for each channel that + * integrates it, else 0. + * @param diagram_channels The channel of each diagram. + * @param flavor_diagram_masks Per flavour, 1 for each diagram it can be + * written with, else 0. + */ + MappingChannelWeights( + const std::vector& mappings, + const nested_vector2& flavor_masks, + const std::vector& diagram_channels = {}, + const nested_vector2& flavor_diagram_masks = {} + ); + + /// Number of channels. + std::size_t channel_count() const { return _mappings.size(); } + /// Whether the beams' orientation is an argument. + bool mirror_beams() const { return _mirror_beams; } + /// Whether a diagram is sampled. + bool has_diagrams() const { return !_diagram_channels.empty(); } + +private: + NamedVector build_function_impl( + FunctionBuilder& fb, const NamedVector& args + ) const override; + + std::vector _mappings; + nested_vector2 _flavor_masks; + std::vector _diagram_channels; + nested_vector2 _flavor_diagram_masks; + bool _mirror_beams; +}; + } // namespace madspace diff --git a/madspace/include/madspace/phasespace/cross_section.hpp b/madspace/include/madspace/phasespace/cross_section.hpp index 6143b95f30..498e646d9e 100644 --- a/madspace/include/madspace/phasespace/cross_section.hpp +++ b/madspace/include/madspace/phasespace/cross_section.hpp @@ -35,6 +35,8 @@ namespace madspace { * **Returns** * - the return values of @p matrix_element, scaled to the differential cross * section. + * - `raw_matrix_element` – `float`, shape `(batch,)` – the squared matrix + * element without PDFs and flux (with @p return_matrix_element). * * **References** * - [1] T. Heimel, O. Mattelaer, R. Winterhalder, "MadSpace", @@ -62,6 +64,8 @@ class DifferentialCrossSection : public FunctionGenerator { * rather than reconstructed. * @param decay If true, describe the decay of a single * incoming particle instead of a collision. + * @param return_matrix_element If true, also return the squared matrix + * element itself as `raw_matrix_element`. */ DifferentialCrossSection( const MatrixElement& matrix_element, @@ -73,7 +77,8 @@ class DifferentialCrossSection : public FunctionGenerator { const std::variant& pdf1 = std::monostate{}, const std::variant& pdf2 = std::monostate{}, bool input_momentum_fraction = true, - bool decay = false + bool decay = false, + bool return_matrix_element = false ); /// The allowed initial-parton flavour combinations. @@ -86,6 +91,8 @@ class DifferentialCrossSection : public FunctionGenerator { const std::optional& running_coupling() const { return _running_coupling; } + /// Whether the squared matrix element is returned as `raw_matrix_element`. + bool returns_matrix_element() const { return _return_matrix_element; } private: NamedVector build_function_impl( @@ -101,6 +108,7 @@ class DifferentialCrossSection : public FunctionGenerator { std::variant _energy_scale; bool _input_momentum_fraction; bool _decay; + bool _return_matrix_element; }; } // namespace madspace diff --git a/madspace/include/madspace/phasespace/integrand.hpp b/madspace/include/madspace/phasespace/integrand.hpp index 0bab9ef189..a2d2a76b1f 100644 --- a/madspace/include/madspace/phasespace/integrand.hpp +++ b/madspace/include/madspace/phasespace/integrand.hpp @@ -41,14 +41,25 @@ namespace madspace { * - `weight` – `float`, shape `(batch,)` – the event weight. * - `momenta` – `float`, shape `(batch, particle_count, 4)` – the momenta. * - `color_index`, `helicity_index`, `diagram_index`, `flavor_index` – `int`, - * shape `(batch,)` – the sampled discrete quantities. + * shape `(batch,)` – the sampled discrete quantities. Indices the matrix + * element does not return are -1. * - `ren_scale`, `alpha_qcd`, `x1`, `x2`, `fact_scale1`, `fact_scale2` – * `float`, shape `(batch,)` – scales and momentum fractions. * - `random` – `float`, shape `(batch, mapping.random_dim())` – the underlying * random numbers. + * - `matrix_element` – `float`, shape `(batch,)` – the squared matrix element, + * present only if the differential cross sections return it, followed by + * `matrix_element_flavor` – `int` – the flavor index passed to it, and with + * several differential cross sections `matrix_element_index` – `int` – the + * index of the one evaluated. * - `bias` – `float`, shape `(batch,)` – the event bias, present only with a * custom bias. `weight` then includes the bias factor. * + * Only `matrix_element_out` is required from the matrix element. + * `diagram_amp2_out` is needed for diagram-based channel weights, i.e. with + * several channels and neither @p prop_chan_weights nor + * @p mapping_chan_weights. + * * **References** * - [1] T. Heimel, O. Mattelaer, R. Winterhalder, "MadSpace", * https://arxiv.org/abs/2602.06895 @@ -73,7 +84,7 @@ class Integrand : public FunctionGenerator { MatrixElement::random_helicity_in, MatrixElement::random_diagram_in, }; - /// The @ref MatrixElement outputs an `Integrand` requests. + /// The @ref MatrixElement outputs an `Integrand` can use. inline static const std::vector matrix_element_outputs = { MatrixElement::matrix_element_out, @@ -84,7 +95,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 @@ -146,6 +159,12 @@ class Integrand : public FunctionGenerator { * the prefix @ref pdf2_prefix. * Not allowed with a @ref CustomPdf, * which receives the beam index. + * @param mapping_chan_weights Optional + * @ref MappingChannelWeights, in + * place of @p prop_chan_weights. + * Its sampled diagram is the + * diagram index of the event when + * the matrix element returns none. * @param mlm_history_qcd_power Power of alpha_s in |M|^2 when it * is the same for every diagram, * -1 otherwise. With an MLM history @@ -187,6 +206,7 @@ class Integrand : public FunctionGenerator { std::size_t compressed_channel_weight_count = 50, const std::optional& custom_bias = std::nullopt, const std::optional& pdf_grid2 = std::nullopt, + const std::optional& mapping_chan_weights = std::nullopt, int mlm_history_qcd_power = -1 ); /// Globals prefix under which the second beam's PDF grid is registered @@ -236,6 +256,10 @@ class Integrand : public FunctionGenerator { const std::optional& prop_chan_weights() const { return _prop_chan_weights; } + /// The mapping-density channel weights, if any. + const std::optional& mapping_chan_weights() const { + return _mapping_chan_weights; + } /// The channel-weight network, if any. const std::optional& chan_weight_net() const { return _chan_weight_net; @@ -321,6 +345,7 @@ class Integrand : public FunctionGenerator { // Empty when the clustering does not depend on the flavours. std::vector _mlm_leg_flavors; std::optional _prop_chan_weights; + std::optional _mapping_chan_weights; std::optional _subchan_weights; std::optional _chan_weight_net; nested_vector2 _first_chan_weight_remap; diff --git a/madspace/include/madspace/phasespace/matrix_element.hpp b/madspace/include/madspace/phasespace/matrix_element.hpp index dfccc64cd6..c7964775b0 100644 --- a/madspace/include/madspace/phasespace/matrix_element.hpp +++ b/madspace/include/madspace/phasespace/matrix_element.hpp @@ -69,7 +69,10 @@ class MatrixElement : public FunctionGenerator { fact_scale1_out, ///< Custom factorization scale, first incoming. fact_scale2_out, ///< Custom factorization scale, second incoming. pdf_out, ///< Custom parton density x f(x, Q). - bias_out ///< Custom event bias. + bias_out, ///< Custom event bias. + /// Ratio of the matrix element to a reference matrix element, for a + /// @ref WeightingStep. + matrix_element_ratio_out }; /** @@ -97,6 +100,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, @@ -106,7 +111,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 @@ -115,6 +121,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, @@ -122,7 +131,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(), @@ -132,7 +142,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; } @@ -144,8 +155,12 @@ class MatrixElement : public FunctionGenerator { const std::vector& inputs() const { return _inputs; } /// The configured output quantities. const std::vector& outputs() const { return _outputs; } + /// Whether the `random_*` inputs are drawn internally. + bool sample_random_inputs() const { return _sample_random_inputs; } + /// The 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. @@ -166,6 +181,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/include/madspace/phasespace/three_particle.hpp b/madspace/include/madspace/phasespace/three_particle.hpp index 2ba277ccd3..f0a3cc3cdb 100644 --- a/madspace/include/madspace/phasespace/three_particle.hpp +++ b/madspace/include/madspace/phasespace/three_particle.hpp @@ -179,7 +179,7 @@ class TwoToThreeParticleScattering : public Mapping { * the momentum transfer @f$(p_2 - k)^2 \le 0@f$ * instead of @f$(p_2 + p_3)^2@f$; its absolute * value is sampled. With @p has_cut it is bounded - * by @f$|(p_2 - k)^2| \ge \texttt{etmin\_2}^2@f$, + * by @f$|(p_2 - k)^2| \ge \texttt{etmin\_2}^2 - m_2^2@f$, * the `drcut` and `s23_min_cut` conditions are * not used. Requires @p p12_condition. AmpliCol * builds a side of two particles this way, with @@ -225,11 +225,12 @@ class TwoToThreeParticleScattering : public Mapping { split_conditions(FunctionBuilder& fb, const NamedVector& conditions) const; // Sampled and kinematic range of the second invariant: s23 itself, or // |(p2 - k)^2| for an incoming recoil. Arguments: the kinematic range of - // s23 and the conditions. + // s23, the squared mass of p2 and the conditions. std::array sampling_ranges( FunctionBuilder& fb, Value s23_min, Value s23_max, + Value m2_squared, const NamedVector& conditions ) const; Invariant _t_invariant; diff --git a/madspace/include/madspace/phasespace/weighting_step.hpp b/madspace/include/madspace/phasespace/weighting_step.hpp new file mode 100644 index 0000000000..a21f1b36d0 --- /dev/null +++ b/madspace/include/madspace/phasespace/weighting_step.hpp @@ -0,0 +1,116 @@ +#pragma once + +#include +#include + +#include "madspace/phasespace/base.hpp" +#include "madspace/phasespace/matrix_element.hpp" + +namespace madspace { + +/** + * One re- or unweighting step after the first unweighting of an + * @ref EventGenerator. + * + * The step evaluates its own matrix element on the events that survived the + * previous steps and multiplies their weights by the ratio + * @f$r = |M_\text{step}|^2 / |M_\text{previous}|^2@f$. All other factors of the + * weight (PDFs, flux, Jacobian, channel weights, bias) are those of the first + * step and cancel in the ratio. An unweighting step then unweights the events + * with respect to the product of the ratios since the last unweighting step, + * times the overweight factor carried from it, with its own maximum weight. A + * reweighting step only multiplies the weights. + * + * Instead of the matrix element, the matrix elements of a step can return the + * ratio directly, with the output `matrix_element_ratio_out`. Its matrix element + * is then the ratio times the one of the previous step, which is the reference + * for the ratio of the next step. + * + * Typical uses are a fast surrogate of the matrix element in the integrand, + * followed by an unweighting step with the full matrix element, or a + * leading-color integrand reweighted to full color. + */ +struct WeightingStep { + /// What the step does with the ratio of the matrix elements. + enum Mode { + unweight, ///< Unweight with respect to the ratio, with its own maximum. + reweight ///< Multiply the event weight by the ratio. + }; + + /// One matrix element per matrix element of the integrand, called for the + /// events with that `matrix_element_index`. The random inputs are passed + /// in, so they must not be sampled internally. + std::vector matrix_elements; + /// Unweighting or reweighting. + Mode mode = unweight; + /// Like @ref GeneratorConfig::max_overweight_truncation, for this step. + double max_overweight_truncation = 0.01; + /// Like @ref GeneratorConfig::freeze_max_weight_after, for this step. + std::size_t freeze_max_weight_after = 10000; + /// Color, helicity or diagram index outputs that replace the values of the + /// earlier steps. Without, an index is taken from the first step whose + /// matrix element returns it. + std::vector override_outputs; + + /// Whether the matrix elements return the ratio to the previous step + /// instead of the matrix element. + bool returns_ratio() const; +}; + +/** + * Evaluates the matrix elements of a @ref WeightingStep. + * + * `batch` is the leading batch dimension. + * + * **Arguments** + * - `momenta` – `float`, shape `(batch, particle_count, 4)`. + * - `alpha_s` – `float`, shape `(batch,)`. + * - `flavor` – `int`, shape `(batch,)` – flavor index of the matrix element. + * - `matrix_element_index` – `int`, shape `(batch,)` – which matrix element to + * call, only with several matrix elements. + * - `color_index`, `helicity_index`, `diagram_index` – `int`, shape `(batch,)` + * – the indices of the earlier steps. + * - `random` – `float`, shape `(batch, 4 step_count)` – four random numbers + * per step: for the color, helicity and diagram sampling of the matrix + * element and for the unweighting. + * + * **Returns** + * - `matrix_element` – `float`, shape `(batch,)` – the matrix element, or the + * ratio to the previous step if the step returns ratios. + * - `color_index`, `helicity_index`, `diagram_index` – `int`, shape `(batch,)` + * – the indices the step writes, otherwise the input values. + * - `unweight_random` – `float`, shape `(batch,)` – the random number of this + * step for the unweighting. + */ +class WeightingStepFunction : public FunctionGenerator { +public: + /// Random numbers per event and step: color, helicity, diagram, unweighting. + static constexpr std::size_t random_count = 4; + + /** + * @param step The step. + * @param particle_count Number of particles of the events. + * @param write_indices Whether the color, helicity and diagram indices + * returned by the matrix elements are written. + * @param step_index Position of the step, from 0. + * @param step_count Number of steps, which sets the shape of `random`. + */ + WeightingStepFunction( + const WeightingStep& step, + std::size_t particle_count, + const std::array& write_indices, + std::size_t step_index = 0, + std::size_t step_count = 1 + ); + +private: + NamedVector build_function_impl( + FunctionBuilder& fb, const NamedVector& args + ) const override; + + std::vector _matrix_elements; + std::array _write_indices; + std::size_t _step_index; +}; + +} // namespace madspace diff --git a/madspace/include/madspace/umami.h b/madspace/include/madspace/umami.h index 91f3adc413..b50e37578c 100644 --- a/madspace/include/madspace/umami.h +++ b/madspace/include/madspace/umami.h @@ -30,7 +30,7 @@ extern "C" { * Minor version number of the UMAMI interface. Between minor versions, new keys for * errors, devices, metadata, inputs and outputs can be added. */ -#define UMAMI_MINOR_VERSION 1 +#define UMAMI_MINOR_VERSION 2 typedef enum { /** operation was executed successfully */ @@ -138,10 +138,13 @@ typedef enum { UMAMI_OUT_PDF, /** value of a custom event bias, type: `double`, shape: `()` */ UMAMI_OUT_BIAS, + /** ratio of the matrix element to a reference matrix element defined by the + * caller, type: `double`, shape: `()` */ + UMAMI_OUT_MATRIX_ELEMENT_RATIO, } UmamiOutputKey; /** Number of values in `UmamiOutputKey` */ -#define UMAMI_OUTPUT_KEY_COUNT (UMAMI_OUT_BIAS + 1) +#define UMAMI_OUTPUT_KEY_COUNT (UMAMI_OUT_MATRIX_ELEMENT_RATIO + 1) /** Implementation-defined pointer to a matrix element instance */ typedef void* UmamiHandle; diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index f78a3c0ee1..19b1627022 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -3114,6 +3114,20 @@ collect_channel_weights: desc: "weight per channel" desc: "Sums squared amplitudes into per-channel weights. Building block of @ref PropagatorChannelWeights." +density_channel_weights: + inputs: + - name: densities + type: [float, n] + desc: "sampling density of each channel at the event, from its inverse mapping" + - name: mask + type: [float, n] + desc: "1 for the channels that integrate the flavor of the event, else 0" + outputs: + - name: channel_weights + type: [float, n] + desc: "weight per channel" + desc: "Normalised channel weights proportional to the masked channel densities, uniform over the masked channels if none is positive. Building block of @ref MappingChannelWeights." + interpolate_pdf: inputs: - name: x diff --git a/madspace/madspace/python_function.py b/madspace/madspace/python_function.py index d2b11b47bd..b5a21be261 100644 --- a/madspace/madspace/python_function.py +++ b/madspace/madspace/python_function.py @@ -39,6 +39,7 @@ "fact_scale2": DataType.float, "pdf": DataType.float, "bias": DataType.float, + "matrix_element_ratio": DataType.float, } diff --git a/madspace/src/compgraphs/instruction.cpp b/madspace/src/compgraphs/instruction.cpp index 543e4d0277..5481a398d2 100644 --- a/madspace/src/compgraphs/instruction.cpp +++ b/madspace/src/compgraphs/instruction.cpp @@ -999,6 +999,7 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const { case UMAMI_OUT_FACT_SCALE_2: case UMAMI_OUT_PDF: case UMAMI_OUT_BIAS: + case UMAMI_OUT_MATRIX_ELEMENT_RATIO: dtype = DataType::dt_float; shape = {}; break; diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 5685f7d5f5..5c795fdc02 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -117,6 +117,7 @@ void op_matrix_element( // implemented at LUMI hackathon where the coffee was indeed terrible switch (output_keys[i]) { case UMAMI_OUT_MATRIX_ELEMENT: + case UMAMI_OUT_MATRIX_ELEMENT_RATIO: device.submit([=]() mutable { double* ptr = static_cast(output_ptrs[i]); std::fill(ptr, ptr + batch_size, 1.); diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index fbdbf7ffb7..67ddb491fd 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -1,5 +1,6 @@ #include "madspace/driver/channel_generator.hpp" #include "madspace/driver/random.hpp" +#include "madspace/phasespace/weighting_step.hpp" #include "madspace/util.hpp" using namespace madspace; @@ -81,6 +82,18 @@ DerivedSeed generate_phase_seed( ); } +// An empty path means the events are only kept in memory. +std::optional open_event_file( + const std::string& path, const DataLayout& layout, std::size_t particle_count +) { + if (path.empty()) { + return std::nullopt; + } + return std::make_optional( + path, layout, particle_count, EventFile::create, true + ); +} + } // namespace ChannelEventGenerator::ChannelEventGenerator( @@ -123,13 +136,10 @@ ChannelEventGenerator::ChannelEventGenerator( ) ), _event_file( - event_file, - _event_file_layout, - integrand.particle_count(), - EventFile::create, - true + open_event_file(event_file, _event_file_layout, integrand.particle_count()) ), - _weight_file(weight_file, weight_file_layout, 0, EventFile::create, true), + _weight_file(open_event_file(weight_file, weight_file_layout, 0)), + _weight_file_path(weight_file), _batch_size(config.start_batch_size), _particle_count(integrand.particle_count()), _integrand_channel_function(IntegrandChannelPart(integrand).function()), @@ -155,6 +165,17 @@ ChannelEventGenerator::ChannelEventGenerator( init_used_globals(); init_runtimes(); init_field_indices(); + auto& me_outputs = integrand.diff_xs().at(0).matrix_element().outputs(); + for (auto [provided, output] : + zip(_provided_indices, + std::array{ + MatrixElement::color_index_out, + MatrixElement::helicity_index_out, + MatrixElement::diagram_index_out + })) { + provided = + std::find(me_outputs.begin(), me_outputs.end(), output) != me_outputs.end(); + } if (const auto& grid_name = integrand.vegas_grid_name(); grid_name) { _vegas_optimizer = VegasGridOptimizer(contexts, grid_name.value(), config.vegas_damping); @@ -250,10 +271,9 @@ ChannelEventGenerator::ChannelEventGenerator( ParticleRecord::f_particle_data | _particle_layout_extra_flags ) ), - _event_file( - event_file, _event_file_layout, particle_count, EventFile::create, true - ), - _weight_file(weight_file, weight_file_layout, 0, EventFile::create, true), + _event_file(open_event_file(event_file, _event_file_layout, particle_count)), + _weight_file(open_event_file(weight_file, weight_file_layout, 0)), + _weight_file_path(weight_file), _batch_size(config.start_batch_size), _particle_count(particle_count), _integrand_channel_function(integrand_channel_function), @@ -299,6 +319,13 @@ void ChannelEventGenerator::init_runtimes() { } void ChannelEventGenerator::init_field_indices() { + auto& unweighter_outputs = _unweighter_function.outputs(); + for (auto [key, value] : zip(unweighter_outputs.keys(), unweighter_outputs)) { + // with the common batch size, to compare and combine channels + _event_types.push_back( + key, Type(value.type.dtype, madspace::batch_size, value.type.shape) + ); + } auto index_map = _integrand_common_function.outputs().index_map(); _field_indices.weight = index_map.at("weight"); _field_indices.momenta = index_map.at("momenta"); @@ -357,18 +384,38 @@ void ChannelEventGenerator::init_field_indices() { _field_indices.subprocess_index = -1; } _field_indices.bias = index_map.contains("bias") ? index_map.at("bias") : -1; + _field_indices.matrix_element = + index_map.contains("matrix_element") ? index_map.at("matrix_element") : -1; _field_indices.random = index_map.at("random"); _field_indices.rest = _field_indices.random + 1; } +EventFile& ChannelEventGenerator::event_file() { + if (!_event_file) { + throw std::logic_error("Channel was created without an event file"); + } + return *_event_file; +} + +EventFile& ChannelEventGenerator::weight_file() { + if (!_weight_file) { + throw std::logic_error("Channel was created without a weight file"); + } + return *_weight_file; +} + void ChannelEventGenerator::unweight_file(MixMaxRandom& rand_gen) { + if (!_step_settings.empty()) { + unweight_step_file(rand_gen); + return; + } std::size_t buf_size = 1000000; EventBuffer buffer(0, 0, weight_file_layout); std::size_t accept_count = _unweighted_accept_count; - for (std::size_t i = _unweighted_count; i < _weight_file.event_count(); - i += buf_size) { - _weight_file.seek(i); - _weight_file.read(buffer, buf_size); + auto& w_file = weight_file(); + for (std::size_t i = _unweighted_count; i < w_file.event_count(); i += buf_size) { + w_file.seek(i); + w_file.read(buffer, buf_size); for (std::size_t j = 0; j < buffer.event_count(); ++j) { auto weight = buffer.event(j).weight(); if (std::abs(weight.value()) / _max_weight < rand_gen.generate_double()) { @@ -380,13 +427,12 @@ void ChannelEventGenerator::unweight_file(MixMaxRandom& rand_gen) { ++accept_count; } } - _weight_file.seek(i); - _weight_file.write(buffer); + w_file.seek(i); + w_file.write(buffer); } - _unweighted_count = _weight_file.event_count(); + _unweighted_count = w_file.event_count(); _unweighted_accept_count = accept_count; _status.count_unweighted = accept_count; - _expected_unweighted = accept_count; } void ChannelEventGenerator::integrate(const GeneratorBatchJob& job) { @@ -406,11 +452,7 @@ void ChannelEventGenerator::integrate(const GeneratorBatchJob& job) { _cross_section.push(weight); _abs_cross_section.push(std::abs(w_view[i])); } - _status.mean = _cross_section.mean(); - _status.error = _cross_section.error(); - _status.mean_abs = _abs_cross_section.mean(); - _status.error_abs = _abs_cross_section.error(); - _status.rel_std_dev = _abs_cross_section.rel_std_dev(); + update_status_integrals(); _status.count += w_view.size(); _status.count_opt += w_view.size(); _count_requested_opt += job.requested_event_count; @@ -462,12 +504,13 @@ void ChannelEventGenerator::optimize_vegas(const GeneratorBatchJob& job) { double ChannelEventGenerator::channel_weight_sum(std::size_t event_count) { std::size_t buf_size = 1000000; - EventBuffer buffer(0, 0, weight_file_layout); + EventBuffer buffer(0, weight_record_count(), weight_layout()); double weight_sum = 0; - _weight_file.seek(0); + auto& w_file = weight_file(); + w_file.seek(0); std::size_t unweighted_count = 0; - for (std::size_t i = 0; i < _weight_file.event_count(); i += buf_size) { - _weight_file.read(buffer, buf_size); + for (std::size_t i = 0; i < w_file.event_count(); i += buf_size) { + w_file.read(buffer, buf_size); bool done = false; for (std::size_t j = 0; j < buffer.event_count(); ++j) { if (unweighted_count == event_count) { @@ -478,7 +521,9 @@ double ChannelEventGenerator::channel_weight_sum(std::size_t event_count) { if (weight == 0.) { continue; } - weight_sum += std::abs(weight) / _max_weight; + weight_sum += _step_settings.empty() + ? std::abs(weight) / _max_weight + : std::abs(unweighted_weight(buffer, j)); ++unweighted_count; } if (done) { @@ -640,6 +685,7 @@ void ChannelEventGenerator::start_job( void ChannelEventGenerator::prepare_unweight_job(GeneratorBatchJob& job) const { job.max_weight = _max_weight; + job.count_target = _status.count_target; } void ChannelEventGenerator::submit_unweight_job( @@ -662,6 +708,30 @@ void ChannelEventGenerator::submit_unweight_job( ); unweighter_args.push_back(Tensor(job.max_weight, context->device())); TensorVec unw_events = runtimes.unweighter->run(unweighter_args); + if (!_step_settings.empty()) { + // the events stay on the device for the steps, with the random + // numbers of the steps drawn per event + if (job.rng_seed) { + runtimes.step_random->set_seed(DerivedSeed( + job.rng_seed, + DerivedSeed::weighting_step_random, + job.rng_job_index, + job.channel_index, + job.rng_is_survey ? 1 + job.rng_survey_pass : 0 + )); + } + auto random = runtimes.step_random->run({unw_events.at(0)}); + job.step_host.push_back(unw_events.at(_field_indices.weight).cpu()); + job.step_host.push_back( + unw_events.at(_field_indices.matrix_element).cpu() + ); + if (_field_indices.bias != -1) { + job.step_host.push_back(unw_events.at(_field_indices.bias).cpu()); + } + job.unweighted_events = std::move(unw_events); + job.unweighted_events.push_back(random.at(0)); + return; + } for (auto& item : unw_events) { job.unweighted_events.push_back(item.cpu()); } @@ -684,7 +754,6 @@ std::size_t ChannelEventGenerator::next_vegas_batch_size() { void ChannelEventGenerator::clear_events() { _status.count_unweighted = 0; - _expected_unweighted = 0; if (!_max_weight_fixed) { _max_weight = 0; } @@ -693,23 +762,36 @@ void ChannelEventGenerator::clear_events() { _status.count_opt = 0; _count_requested_opt = 0; _status.count_after_cuts_opt = 0; - _event_file.clear(); - _weight_file.clear(); + if (_event_file) { + _event_file->clear(); + } + if (_weight_file) { + _weight_file->clear(); + } _cross_section.reset(); _abs_cross_section.reset(); _large_weights.clear(); + _first_step_count = 0.; + _freeze_count = 0.; + _final_integral = {}; + ++_event_epoch; + for (auto& state : _step_states) { + if (!state.max_weight_fixed) { + state.max_weight = 0.; + } + state.large_weights.clear(); + state.count_accepted = 0.; + state.freeze_count = 0.; + } for (auto& hist : _histograms) { std::fill(hist.bin_values.begin(), hist.bin_values.end(), 0.); std::fill(hist.bin_errors.begin(), hist.bin_errors.end(), 0.); } } -void ChannelEventGenerator::update_max_weight(Tensor weights) { - // Freeze on the expected unweighted count of the committed generate stages - // rather than on count_unweighted: the latter is advanced by the unweight - // stages, whose commits interleave with the generate stages depending on - // thread timing, which made the freeze point (and the run) timing-dependent. - if (_max_weight_fixed || _expected_unweighted > _config.freeze_max_weight_after) { +void ChannelEventGenerator::update_max_weight(Tensor weights, bool unweight) { + // not the written events, they are committed in a thread dependent order + if (_max_weight_fixed || _freeze_count > _config.freeze_max_weight_after) { return; } @@ -729,11 +811,9 @@ void ChannelEventGenerator::update_max_weight(Tensor weights) { } std::sort(_large_weights.begin(), _large_weights.end(), std::greater{}); apply_truncation_budget(); - // expected accepted events of this job, against the max weight it will be - // unweighted with (see prepare_unweight_job) - if (_max_weight > 0) { + if (unweight && _max_weight > 0) { for (std::size_t i = 0; i < w_view.size(); ++i) { - _expected_unweighted += std::min(std::abs(w_view[i]) / _max_weight, 1.); + _freeze_count += std::min(1., std::abs(w_view[i]) / _max_weight); } } } @@ -747,7 +827,7 @@ void ChannelEventGenerator::set_target_count(std::size_t target_count) { // channel that has met its smaller target is not revisited. Tighten the cap // now; unweight_all() re-unweights the whole file against it at the end. if (target_count < old_target && _max_weight != 0 && !_max_weight_fixed && - _expected_unweighted <= _config.freeze_max_weight_after) { + _freeze_count <= _config.freeze_max_weight_after) { apply_truncation_budget(); } } @@ -768,7 +848,11 @@ void ChannelEventGenerator::apply_truncation_budget() { if (w_sum / w > max_truncation) { if (_max_weight < w) { _status.count_unweighted *= _max_weight / w; - _expected_unweighted *= _max_weight / w; + _first_step_count *= _max_weight / w; + _freeze_count *= _max_weight / w; + for (auto& state : _step_states) { + state.count_accepted *= _max_weight / w; + } _max_weight = w; _unweighted_count = 0; _unweighted_accept_count = 0; @@ -782,6 +866,21 @@ void ChannelEventGenerator::apply_truncation_budget() { void ChannelEventGenerator::write_events( const std::vector& unweighted_events, double job_max_weight ) { + EventBuffer event_buffer(0, _particle_count, _event_file_layout); + EventBuffer weight_buffer(0, 0, weight_file_layout); + fill_event_buffers(unweighted_events, event_buffer, weight_buffer); + event_file().write(event_buffer); + weight_file().write(weight_buffer); + _status.count_unweighted += event_buffer.event_count() * + (job_max_weight > 0 ? job_max_weight / _max_weight : 1); + _status.done = _status.count_unweighted >= _status.count_target; +} + +void ChannelEventGenerator::fill_event_buffers( + const TensorVec& unweighted_events, + EventBuffer& event_buffer, + EventBuffer& weight_buffer +) const { auto w_view = unweighted_events.at(_field_indices.weight).view(); auto mom_view = unweighted_events.at(_field_indices.momenta).view(); auto colors_view = @@ -796,10 +895,8 @@ void ChannelEventGenerator::write_events( unweighted_events.at(_field_indices.ren_scale).view(); auto alphas_view = unweighted_events.at(_field_indices.alpha_qcd).view(); - EventBuffer event_buffer( - w_view.size(), _event_file.particle_count(), _event_file_layout - ); - EventBuffer weight_buffer(w_view.size(), 0, weight_file_layout); + event_buffer.resize(w_view.size()); + weight_buffer.resize(w_view.size()); for (std::size_t i = 0; i < w_view.size(); ++i) { weight_buffer.event(i).weight() = w_view[i]; auto event = event_buffer.event(i); @@ -915,12 +1012,6 @@ void ChannelEventGenerator::write_events( event.subprocess_index() = subproc_view[i]; } } - - _event_file.write(event_buffer); - _weight_file.write(weight_buffer); - _status.count_unweighted += - w_view.size() * (job_max_weight > 0 ? job_max_weight / _max_weight : 1); - _status.done = _status.count_unweighted >= _status.count_target; } void ChannelEventGenerator::save(const std::string& file_name) const { @@ -933,13 +1024,22 @@ std::string ChannelEventGenerator::to_json(bool include_estimates) const { j = *this; if (include_estimates) { j["max_weight"] = _max_weight; - double error = _abs_cross_section.error(); - if (_abs_cross_section.count() > 1 && std::isfinite(error) && error > 0) { - j["abs_integral"] = {{"mean", _abs_cross_section.mean()}, {"error", error}}; + auto integral = result_integral(); + double error = std::sqrt(integral.variance_abs / integral.count); + if (integral.count > 1 && std::isfinite(error) && error > 0) { + j["abs_integral"] = {{"mean", integral.mean_abs}, {"error", error}}; j["integral"] = { - {"mean", _cross_section.mean()}, {"error", _cross_section.error()} + {"mean", integral.mean}, + {"error", std::sqrt(integral.variance / integral.count)} }; } + if (!_step_states.empty()) { + json step_max_weights = json::array(); + for (auto& state : _step_states) { + step_max_weights.push_back(state.max_weight); + } + j["step_max_weights"] = step_max_weights; + } } return j.dump(); } @@ -1002,6 +1102,14 @@ ChannelEventGenerator ChannelEventGenerator::load_json( if (channel.contains("max_weight") && channel.at("max_weight").get() > 0) { generator.set_fixed_max_weight(channel.at("max_weight").get()); } + if (channel.contains("provided_indices")) { + generator._provided_indices = + channel.at("provided_indices").get>(); + } + if (channel.contains("step_max_weights")) { + generator._loaded_step_max_weights = + channel.at("step_max_weights").get>(); + } if (channel.contains("abs_integral")) { generator.set_abs_integral_prior( channel.at("abs_integral").at("mean").get(), @@ -1041,24 +1149,26 @@ void ChannelEventGenerator::set_integral_prior(double mean, double error) { _integral_prior = {mean, error}; } -bool ChannelEventGenerator::uses_integral_prior() const { - double error = _abs_cross_section.error(); +bool ChannelEventGenerator::uses_integral_prior(bool final_step) const { + auto integral = final_step ? result_integral() : first_step_integral(); + double error = std::sqrt(integral.variance_abs / integral.count); return _abs_integral_prior && - (_abs_cross_section.count() < 2 || !std::isfinite(error) || !(error > 0)); + (integral.count < 2 || !std::isfinite(error) || !(error > 0)); } std::pair ChannelEventGenerator::integral_prior() const { return _integral_prior ? *_integral_prior : *_abs_integral_prior; } -double ChannelEventGenerator::abs_integral_estimate() const { - double mean = _abs_cross_section.mean(); +double ChannelEventGenerator::abs_integral_estimate(bool final_step) const { + auto integral = final_step ? result_integral() : first_step_integral(); + double mean = integral.mean_abs; if (!_abs_integral_prior) { return mean; } auto [prior_mean, prior_error] = *_abs_integral_prior; - double error = _abs_cross_section.error(); - if (uses_integral_prior()) { + double error = std::sqrt(integral.variance_abs / integral.count); + if (uses_integral_prior(final_step)) { return prior_mean; } double prior_weight = 1. / (prior_error * prior_error); @@ -1066,6 +1176,431 @@ double ChannelEventGenerator::abs_integral_estimate() const { return (prior_mean * prior_weight + mean * weight) / (prior_weight + weight); } +ChannelEventGenerator::IntegralEstimate +ChannelEventGenerator::first_step_integral() const { + return { + .mean = _cross_section.mean(), + .variance = _cross_section.variance(), + .mean_abs = _abs_cross_section.mean(), + .variance_abs = _abs_cross_section.variance(), + .count = _cross_section.count(), + }; +} + +ChannelEventGenerator::IntegralEstimate ChannelEventGenerator::result_integral() const { + if (_step_settings.empty() || _final_integral.count == 0) { + return first_step_integral(); + } + auto& integral = _final_integral; + double count = static_cast(integral.count); + auto variance = [&](double sum2, double mean) { + return integral.count > 1 + ? std::max(0., (sum2 - count * mean * mean) / (count - 1)) + : 0.; + }; + double mean = integral.sum / count; + double mean_abs = integral.abs_sum / count; + return { + .mean = mean, + .variance = variance(integral.sum2, mean), + .mean_abs = mean_abs, + .variance_abs = variance(integral.abs_sum2, mean_abs), + .count = integral.count, + }; +} + +void ChannelEventGenerator::update_status_integrals() { + auto integral = result_integral(); + _status.mean = integral.mean; + _status.error = std::sqrt(integral.variance / integral.count); + _status.mean_abs = integral.mean_abs; + _status.error_abs = std::sqrt(integral.variance_abs / integral.count); + _status.rel_std_dev = std::sqrt(integral.variance_abs) / integral.mean_abs; +} + +void ChannelEventGenerator::set_weighting_steps( + const std::vector& steps +) { + if (steps.empty()) { + throw std::invalid_argument("at least one weighting step is needed"); + } + if (!_step_settings.empty()) { + bool same = + steps.size() == _step_settings.size() && + std::equal( + steps.begin(), + steps.end(), + _step_settings.begin(), + [](auto& a, auto& b) { + return a.unweight == b.unweight && a.ratio == b.ratio && + a.max_overweight_truncation == b.max_overweight_truncation && + a.freeze_max_weight_after == b.freeze_max_weight_after; + } + ); + if (!same) { + throw std::logic_error( + "the weighting steps of a channel cannot be changed" + ); + } + return; + } + if (_field_indices.matrix_element == -1) { + throw std::invalid_argument( + "a channel with weighting steps needs an integrand that returns the " + "matrix element" + ); + } + if (_status.count > 0) { + throw std::logic_error( + "the weighting steps must be set before events are generated" + ); + } + _step_settings = steps; + _step_states.assign(steps.size(), {}); + if (!_loaded_step_max_weights.empty()) { + if (_loaded_step_max_weights.size() != steps.size()) { + throw std::invalid_argument( + "the stored maximum weights do not match the weighting steps" + ); + } + for (std::size_t step = 1; double max_weight : _loaded_step_max_weights) { + if (max_weight > 0) { + set_fixed_step_max_weight(step, max_weight); + } + ++step; + } + } + if (_weight_file) { + _weight_file.reset(); + _weight_file = open_event_file( + _weight_file_path, step_weight_file_layout, steps.size() + 1 + ); + } + auto random_count = + static_cast(WeightingStepFunction::random_count * steps.size()); + FunctionBuilder fb( + NamedVector{{"weight", batch_float}}, + NamedVector{{"random", batch_float_array(random_count)}} + ); + fb.output(0, fb.random(fb.batch_size({fb.input(0)}), random_count)); + Function random_function = fb.function(); + for (auto [context, runtimes] : zip(_contexts, _runtimes)) { + runtimes.step_random = build_runtime(random_function, context, false); + } +} + +void ChannelEventGenerator::set_first_pass_settings( + std::optional max_overweight_truncation, + std::optional freeze_max_weight_after +) { + if (max_overweight_truncation) { + _config.max_overweight_truncation = *max_overweight_truncation; + } + if (freeze_max_weight_after) { + _config.freeze_max_weight_after = *freeze_max_weight_after; + _first_pass_freeze_max_weight_after = freeze_max_weight_after; + } +} + +void ChannelEventGenerator::set_fixed_step_max_weight( + std::size_t step, double max_weight +) { + if (!(max_weight > 0)) { + throw std::invalid_argument("The fixed maximum weight must be positive"); + } + auto& state = _step_states.at(step - 1); + state.max_weight = max_weight; + state.max_weight_fixed = true; + _unweighted_count = 0; + _unweighted_accept_count = 0; +} + +void ChannelEventGenerator::update_step_max_weight( + std::size_t step, const std::vector& values, std::size_t count_target +) { + auto& state = _step_states.at(step - 1); + auto& settings = _step_settings.at(step - 1); + // not count_accepted, it depends on the first maximum at the time of the call + if (state.max_weight_fixed || + state.freeze_count > settings.freeze_max_weight_after) { + return; + } + double min_nonzero = 0.; + for (double value : values) { + double abs_value = std::abs(value); + if (!std::isfinite(abs_value)) { + continue; + } + if (abs_value != 0 && (min_nonzero == 0 || abs_value < min_nonzero)) { + min_nonzero = abs_value; + } + if (abs_value > state.max_weight) { + state.large_weights.push_back(abs_value); + } + } + if (state.max_weight == 0) { + state.max_weight = min_nonzero; + } + std::sort(state.large_weights.begin(), state.large_weights.end(), std::greater{}); + apply_step_truncation_budget(step, count_target); + if (state.max_weight > 0) { + for (double value : values) { + state.freeze_count += std::min(1., std::abs(value) / state.max_weight); + } + } +} + +void ChannelEventGenerator::apply_step_truncation_budgets() { + for (std::size_t step = 1; + auto [state, settings] : zip(_step_states, _step_settings)) { + if (settings.unweight && !state.max_weight_fixed && state.max_weight != 0 && + state.freeze_count <= settings.freeze_max_weight_after) { + apply_step_truncation_budget(step, _status.count_target); + } + ++step; + } +} + +void ChannelEventGenerator::apply_step_truncation_budget( + std::size_t step, std::size_t count_target +) { + auto& state = _step_states.at(step - 1); + auto& settings = _step_settings.at(step - 1); + auto& large_weights = state.large_weights; + // same truncation as for the first step, see apply_truncation_budget() + double max_truncation = settings.max_overweight_truncation * + static_cast(std::min(count_target, settings.freeze_max_weight_after)); + double sum = 0; + std::size_t count = 0; + for (double value : large_weights) { + if (value < state.max_weight) { + break; + } + sum += value; + ++count; + if (sum / value > max_truncation) { + if (state.max_weight < value) { + // the events of this and the later steps + double factor = state.max_weight / value; + _status.count_unweighted *= factor; + for (std::size_t i = step - 1; i < _step_states.size(); ++i) { + _step_states.at(i).count_accepted *= factor; + } + state.freeze_count *= factor; + state.max_weight = value; + _unweighted_count = 0; + _unweighted_accept_count = 0; + } + break; + } + } + large_weights.erase(large_weights.begin() + count, large_weights.end()); +} + +void ChannelEventGenerator::add_step_accepted( + std::size_t step, + double count, + double max_weight, + const std::vector& max_weights +) { + // the fraction left with the current maxima, like in write_step_events() + double factor = max_weight > 0 ? max_weight / _max_weight : 1.; + for (std::size_t i = 0; i + 1 < step; ++i) { + if (_step_settings.at(i).unweight && max_weights.at(i) > 0) { + factor *= max_weights.at(i) / _step_states.at(i).max_weight; + } + } + _step_states.at(step - 1).count_accepted += count * factor; +} + +void ChannelEventGenerator::add_first_step_events( + std::size_t count, double job_max_weight +) { + _first_step_count += + count * (job_max_weight > 0 ? job_max_weight / _max_weight : 1); +} + +StepIntegral ChannelEventGenerator::step_integral(const StepEvents& step_events) const { + std::size_t step_count = _step_settings.size(); + // the final weights in units of the cross section + double weight_unit = step_events.max_weight; + for (auto [settings, job_max_weight] : + zip(_step_settings, step_events.max_weights)) { + if (settings.unweight) { + weight_unit *= job_max_weight; + } + } + StepIntegral integral{.count = step_events.generated_count}; + std::vector ratios(step_count); + for (std::size_t i = 0; i < step_events.weights.size(); ++i) { + for (std::size_t k = 0; k < step_count; ++k) { + ratios.at(k) = step_events.ratios.at(k).at(i); + } + double weight = weight_unit * + evaluate_step_chain( + step_events.weights.at(i), + step_events.max_weight, + ratios.data(), + step_events.max_weights.data(), + _step_settings + ) + .weight; + double xs_weight = + step_events.biases.empty() ? weight : weight / step_events.biases.at(i); + integral.sum += xs_weight; + integral.sum2 += xs_weight * xs_weight; + integral.abs_sum += std::abs(weight); + integral.abs_sum2 += weight * weight; + } + return integral; +} + +void ChannelEventGenerator::add_step_integral(const StepIntegral& integral) { + _final_integral.count += integral.count; + _final_integral.sum += integral.sum; + _final_integral.sum2 += integral.sum2; + _final_integral.abs_sum += integral.abs_sum; + _final_integral.abs_sum2 += integral.abs_sum2; + update_status_integrals(); +} + +void ChannelEventGenerator::fill_step_buffers( + const StepEvents& step_events, EventBuffer& event_buffer, EventBuffer& weight_buffer +) const { + fill_event_buffers(step_events.events, event_buffer, weight_buffer); + for (std::size_t i = 0; i < step_events.weights.size(); ++i) { + weight_buffer.event(i).weight() = step_events.weights.at(i); + auto first = weight_buffer.particle(i, 0); + first.step_ratio() = 1.; + first.step_max_weight() = step_events.max_weight; + for (std::size_t k = 0; k < _step_settings.size(); ++k) { + auto record = weight_buffer.particle(i, k + 1); + record.step_ratio() = step_events.ratios.at(k).at(i); + record.step_max_weight() = step_events.max_weights.at(k); + } + } +} + +void ChannelEventGenerator::write_step_events(const StepEvents& step_events) { + add_step_integral(step_integral(step_events)); + std::size_t count = step_events.weights.size(); + if (count > 0) { + EventBuffer event_buffer(0, _particle_count, _event_file_layout); + EventBuffer weight_buffer(0, weight_record_count(), step_weight_file_layout); + fill_step_buffers(step_events, event_buffer, weight_buffer); + event_file().write(event_buffer); + weight_file().write(weight_buffer); + } + // the fraction left by unweighting with the current maxima. Like in + // write_events(), it is rescaled when a maximum rises, so it does not + // depend on the order of the updates. + double count_factor = + step_events.max_weight > 0 ? step_events.max_weight / _max_weight : 1.; + for (auto [state, settings, job_max_weight] : + zip(_step_states, _step_settings, step_events.max_weights)) { + if (settings.unweight && job_max_weight > 0) { + count_factor *= job_max_weight / state.max_weight; + } + } + _status.count_unweighted += count * count_factor; + _status.done = _status.count_unweighted >= _status.count_target; +} + +StepChainResult +ChannelEventGenerator::step_chain(EventBuffer& weight_buffer, std::size_t index) const { + std::vector ratios, max_weights; + for (std::size_t k = 0; auto& state : _step_states) { + ratios.push_back(weight_buffer.particle(index, k + 1).step_ratio()); + max_weights.push_back(state.max_weight); + ++k; + } + return evaluate_step_chain( + weight_buffer.event(index).weight(), + _max_weight, + ratios.data(), + max_weights.data(), + _step_settings + ); +} + +double ChannelEventGenerator::unweighted_weight( + EventBuffer& weight_buffer, std::size_t index +) const { + double weight = weight_buffer.event(index).weight(); + if (_step_settings.empty()) { + return std::copysign(std::max(1., std::abs(weight) / _max_weight), weight); + } + return weight == 0. ? 0. : step_chain(weight_buffer, index).weight; +} + +void ChannelEventGenerator::unweight_step_file(MixMaxRandom& rand_gen) { + std::size_t buf_size = 100000; + std::size_t step_count = _step_settings.size(); + EventBuffer buffer(0, step_count + 1, step_weight_file_layout); + std::vector ratios(step_count), old_max_weights(step_count), + max_weights(step_count); + for (auto [max_weight, state] : zip(max_weights, _step_states)) { + max_weight = state.max_weight; + } + std::size_t accept_count = _unweighted_accept_count; + auto& w_file = weight_file(); + for (std::size_t i = _unweighted_count; i < w_file.event_count(); i += buf_size) { + w_file.seek(i); + w_file.read(buffer, buf_size); + for (std::size_t j = 0; j < buffer.event_count(); ++j) { + auto weight = buffer.event(j).weight(); + double random = rand_gen.generate_double(); + if (weight.value() == 0.) { + continue; + } + for (std::size_t k = 0; k < step_count; ++k) { + auto record = buffer.particle(j, k + 1); + ratios.at(k) = record.step_ratio(); + old_max_weights.at(k) = record.step_max_weight(); + } + // accepted with the stored maxima: keep with the ratio of the + // acceptances, which only decrease with larger maxima + double old_acceptance = + evaluate_step_chain( + weight.value(), + buffer.particle(j, 0).step_max_weight(), + ratios.data(), + old_max_weights.data(), + _step_settings + ) + .acceptance; + double acceptance = + evaluate_step_chain( + weight.value(), + _max_weight, + ratios.data(), + max_weights.data(), + _step_settings + ) + .acceptance; + if (!(acceptance > random * old_acceptance)) { + weight = 0; + continue; + } + weight = std::copysign( + std::max(std::abs(weight.value()), _max_weight), weight.value() + ); + buffer.particle(j, 0).step_max_weight() = _max_weight; + for (std::size_t k = 0; k < step_count; ++k) { + if (_step_settings.at(k).unweight) { + buffer.particle(j, k + 1).step_max_weight() = max_weights.at(k); + } + } + ++accept_count; + } + w_file.seek(i); + w_file.write(buffer); + } + _unweighted_count = w_file.event_count(); + _unweighted_accept_count = accept_count; + _status.count_unweighted = accept_count; +} + void madspace::to_json(nlohmann::json& j, const ChannelEventGenerator& channel) { json histograms = json::array(); json histogram_function; @@ -1092,5 +1627,6 @@ void madspace::to_json(nlohmann::json& j, const ChannelEventGenerator& channel) {"subprocess_index", channel._status.subprocess}, {"name", channel._status.name}, {"histograms", histograms}, + {"provided_indices", channel._provided_indices}, }; } diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index da22f8d241..2ce20a971f 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -8,6 +8,7 @@ #include #include +#include "madspace/driver/event_utils.hpp" #include "madspace/driver/logger.hpp" #include "madspace/driver/random.hpp" #include "madspace/util.hpp" @@ -21,7 +22,8 @@ EventGenerator::EventGenerator( const std::vector>& channels, std::uint64_t seed, std::shared_ptr status_file, - const GeneratorConfig& config + const GeneratorConfig& config, + const std::vector& groups ) : _config(config), _status{ @@ -48,6 +50,7 @@ EventGenerator::EventGenerator( _channel_optimizing(channels.size()), _channel_integral_fractions(channels.size(), 1.), _context_job_counts(contexts.size()), + _warned_invalid_step_values(channels.size()), _channel_batch_pending(channels.size()), _channel_batch_dispatch_done(channels.size()), _channel_gen_order(channels.size()), @@ -55,6 +58,11 @@ EventGenerator::EventGenerator( _channel_unweight_order(channels.size()), _channel_unweight_ready(channels.size()), _context_unweight_queue(contexts.size()), + _pipeline( + groups.empty() + ? nullptr + : std::make_unique(contexts, channels, groups, config) + ), _seed(seed), _status_file(status_file) {} @@ -68,7 +76,7 @@ void EventGenerator::commit_generate_job(GeneratorBatchJob& job) { } channel->integrate(job); update_integral_status(); - channel->update_max_weight(job.weights); + channel->update_max_weight(job.weights, job.unweight); --channel_job_count; --_context_job_counts.at(job.context_index); if (job.unweight) { @@ -83,7 +91,11 @@ void EventGenerator::commit_generate_job(GeneratorBatchJob& job) { void EventGenerator::commit_unweight_job(GeneratorBatchJob& job) { auto& channel = _channels.at(job.channel_index); - channel->write_events(job.unweighted_events, job.max_weight); + if (_pipeline && _pipeline->contains(job.channel_index)) { + _pipeline->add_events(job); + } else { + channel->write_events(job.unweighted_events, job.max_weight); + } update_counts(); --_channel_job_counts.at(job.channel_index); --_context_job_counts.at(job.context_index); @@ -126,7 +138,8 @@ void EventGenerator::run_generate(GenerateMode mode) { if (channel->status().count_opt == 0 || mode == GenerateMode::fix_max_weights) { channel->set_target_count( mode == GenerateMode::fix_max_weights - ? _config.freeze_max_weight_after + ? channel->first_pass_freeze_max_weight_after() + .value_or(_config.freeze_max_weight_after) : _config.target_count ); } @@ -185,6 +198,16 @@ void EventGenerator::run_generate(GenerateMode mode) { // Wait for the round to fully commit before resyncing cross-channel fractions. while (round_in_flight > 0) { std::size_t job_id = wait_for_result(round_in_flight); + if (WeightingPipeline::owns_job(job_id)) { + _pipeline->commit_job(job_id); + update_counts(); + print_gen_update(false); + std::size_t job_id_refill = _job_id; + std::size_t refill_unweight_dispatched = start_jobs(); + round_in_flight += (_job_id - job_id_refill) + + refill_unweight_dispatched + start_pipeline_jobs(); + continue; + } auto& job = _running_jobs.at(job_id); if (job.unweighted_events.size() == 0) { auto& ready = _channel_ready_gen.at(job.channel_index); @@ -212,8 +235,10 @@ void EventGenerator::run_generate(GenerateMode mode) { } std::size_t job_id_refill = _job_id; std::size_t refill_unweight_dispatched = start_jobs(); - round_in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched; + round_in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched + + start_pipeline_jobs(); } + drain_pipeline(); if (mode == GenerateMode::optimize_only) { if (std::none_of(_channels.begin(), _channels.end(), [](auto& channel) { @@ -225,6 +250,10 @@ void EventGenerator::run_generate(GenerateMode mode) { } update_integral_fractions(); + // only with the steps drained, they are not in the commit order of the jobs + for (auto& channel : _channels) { + channel->apply_step_truncation_budgets(); + } update_counts(); if (_status.done) { @@ -235,6 +264,7 @@ void EventGenerator::run_generate(GenerateMode mode) { } } print_gen_update(true); + warn_invalid_step_values(); } // Same commit-ordering approach as generate(). @@ -261,15 +291,27 @@ void EventGenerator::survey(std::size_t survey_pass) { } print_survey_init(); + auto iterates = [&](auto& channel, std::size_t iter) { + return channel->status().iterations <= iter && + !(iter >= min_iters && + channel->abs_cross_section().rel_error() < target_precision); + }; + // the events for the steps, until it is known whether the channel iterates + // again and clears them + std::vector> deferred_jobs(_channels.size()); + auto release_deferred = [&](std::size_t channel_index, bool keep) { + auto& jobs = deferred_jobs.at(channel_index); + if (keep) { + for (auto& job : jobs) { + _pipeline->add_events(job); + } + } + jobs.clear(); + }; std::size_t iter = 0; for (; !done && iter < max_iters; ++iter) { for (std::size_t i = 0; auto channel : _channels) { - if (channel->status().iterations > iter) { - ++i; - continue; - } - if (iter >= min_iters && - channel->abs_cross_section().rel_error() < target_precision) { + if (!iterates(channel, iter)) { ++i; continue; } @@ -291,6 +333,15 @@ void EventGenerator::survey(std::size_t survey_pass) { done = true; while (in_flight > 0) { std::size_t job_id = wait_for_result(in_flight); + if (WeightingPipeline::owns_job(job_id)) { + _pipeline->commit_job(job_id); + update_counts(); + std::size_t job_id_refill = _job_id; + std::size_t refill_unweight_dispatched = start_jobs(); + in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched + + start_pipeline_jobs(); + continue; + } auto& job = _running_jobs.at(job_id); if (job.unweighted_events.size() == 0) { // Generate-stage completion, globally ordered via @@ -306,7 +357,7 @@ void EventGenerator::survey(std::size_t survey_pass) { } channel->integrate(gen_job); update_integral(); - channel->update_max_weight(gen_job.weights); + channel->update_max_weight(gen_job.weights, gen_job.unweight); --channel_job_count; --_context_job_counts.at(gen_job.context_index); // Global order is ascending job id and so is the per-channel one, @@ -339,7 +390,17 @@ void EventGenerator::survey(std::size_t survey_pass) { order.pop_front(); auto& uw_job = _running_jobs.at(commit_id); auto& channel = _channels.at(uw_job.channel_index); - channel->write_events(uw_job.unweighted_events, uw_job.max_weight); + bool defer = false; + if (_pipeline && _pipeline->contains(uw_job.channel_index)) { + defer = iter + 1 < max_iters; + if (!defer) { + _pipeline->add_events(uw_job); + } + } else { + channel->write_events( + uw_job.unweighted_events, uw_job.max_weight + ); + } update_counts(); auto& channel_job_count = _channel_job_counts.at(uw_job.channel_index); @@ -353,16 +414,55 @@ void EventGenerator::survey(std::size_t survey_pass) { channel->optimize_vegas(uw_job); done_event_count += uw_job.batch_event_count; } + std::size_t channel_index = uw_job.channel_index; + if (defer) { + deferred_jobs.at(channel_index).push_back(std::move(uw_job)); + } _running_jobs.erase(commit_id); + // otherwise it depends on the other channels, see below + if (defer && channel_job_count == 0 && + (!done || !iterates(channel, iter + 1))) { + release_deferred(channel_index, !iterates(channel, iter + 1)); + } } } std::size_t job_id_refill = _job_id; std::size_t refill_unweight_dispatched = start_jobs(); - in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched; + in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched + + start_pipeline_jobs(); print_survey_update(false, done_event_count, total_event_count, iter); } + for (std::size_t channel_index = 0; channel_index < _channels.size(); + ++channel_index) { + release_deferred( + channel_index, + done || iter + 1 >= max_iters || + !iterates(_channels.at(channel_index), iter + 1) + ); + } + update_counts(); + drain_pipeline(); } print_survey_update(true, done_event_count, total_event_count, iter - 1); + warn_invalid_step_values(); +} + +void EventGenerator::warn_invalid_step_values() { + for (auto [channel, warned] : zip(_channels, _warned_invalid_step_values)) { + std::size_t count = channel->status().count_invalid_step_values; + if (count > warned) { + Logger::warning( + std::format( + "channel '{}': {} events got the weight 0 in the weighting steps, " + "as their matrix element ratio was not finite or the earlier " + "matrix element was 0", + channel->status().name, + count - warned + ) + ); + warned = count; + } + } } // Every job counted in *in_flight* posts exactly one result (see @@ -392,6 +492,9 @@ std::size_t EventGenerator::wait_for_result(std::size_t& in_flight) { void EventGenerator::reset_jobs() { _running_jobs.clear(); + if (_pipeline) { + _pipeline->reset(); + } _ready_jobs.clear(); _ready_job_rr_cursor = 0; _ready_gen.clear(); @@ -525,110 +628,53 @@ std::size_t EventGenerator::start_jobs() { return unweight_dispatched; } -void EventGenerator::update_integral() { - update_integral_status(); - update_integral_fractions(); +std::size_t EventGenerator::start_pipeline_jobs() { + if (!_pipeline) { + return 0; + } + // once no more events are dispatched, partial batches use the idle threads + bool flush = _ready_jobs.empty() && + std::all_of(_context_unweight_queue.begin(), + _context_unweight_queue.end(), + [](auto& queue) { return queue.empty(); }); + return _pipeline->start_jobs(_result_queue, flush, _context_job_counts); } -void EventGenerator::update_integral_status() { - double total_mean = 0., total_var = 0.; - double total_mean_abs = 0., total_var_abs = 0.; - std::size_t total_count = 0, total_count_opt = 0; - std::size_t total_count_after_cuts = 0, total_count_after_cuts_opt = 0; - std::size_t total_integ_count = 0; - std::size_t iterations = 0; - bool optimized = true; - const ChannelEventGenerator* bad_channel = nullptr; - for (auto& channel : _channels) { - auto& status = channel->status(); - auto& cross_section = channel->cross_section(); - auto& abs_cross_section = channel->abs_cross_section(); - // special case for channels with 0 samples, as they have nan variance - if (bad_channel == nullptr && cross_section.count() > 0 && - (!std::isfinite(cross_section.mean()) || - !std::isfinite(cross_section.variance()) || - !std::isfinite(abs_cross_section.mean()) || - !std::isfinite(abs_cross_section.variance()))) { - bad_channel = channel.get(); - } - if (channel->uses_integral_prior()) { - // not sampled in this run (no events assigned): use the earlier estimate - auto [mean, error] = channel->integral_prior(); - auto [mean_abs, error_abs] = channel->abs_integral_prior(); - total_mean += mean; - total_var += error * error; - total_mean_abs += mean_abs; - total_var_abs += error_abs * error_abs; - } else { - total_mean += cross_section.mean(); - total_var += cross_section.variance() / cross_section.count(); - total_mean_abs += abs_cross_section.mean(); - total_var_abs += abs_cross_section.variance() / abs_cross_section.count(); - } - total_count += status.count; - total_count_opt += status.count_opt; - total_count_after_cuts += status.count_after_cuts; - total_count_after_cuts_opt += status.count_after_cuts_opt; - total_integ_count += cross_section.count(); - iterations = std::max(status.iterations, iterations); - if (!channel->status().optimized) { - optimized = false; - } +// Batches no longer have to be full: the steps don't depend on the batching. +void EventGenerator::drain_pipeline() { + if (!_pipeline) { + return; } - if (bad_channel != nullptr || !std::isfinite(total_mean) || - !std::isfinite(total_mean_abs)) { - std::string where = bad_channel != nullptr - ? std::format( - "channel '{}' after {} samples (mean={}, variance={})", - bad_channel->status().name, - bad_channel->cross_section().count(), - bad_channel->cross_section().mean(), - bad_channel->cross_section().variance() - ) - : std::format("the sum over channels (mean={})", total_mean); - throw std::runtime_error( - std::format( - "non-finite integral in {}. A non-finite weight cannot be recovered " - "from by sampling further, so the integration is aborted. This usually " - "means the matrix element or the phase-space mapping returned nan/inf " - "for " - "some phase-space point -- check the process for a zero or negative " - "width, " - "a parameter point outside the model's validity, or a kinematic " - "configuration at a threshold.", - where - ) - ); + std::size_t in_flight = 0; + while (_pipeline->busy()) { + in_flight += _pipeline->start_jobs(_result_queue, true); + if (in_flight == 0) { + throw std::logic_error("weighting steps stuck without running jobs"); + } + _pipeline->commit_job(wait_for_result(in_flight)); + update_counts(); } + update_integral_status(); +} - _status.mean = total_mean; - _status.error = std::sqrt(total_var); - _status.mean_abs = total_mean_abs; - _status.error_abs = std::sqrt(total_var_abs); - _status.rel_std_dev = std::sqrt(total_var_abs * total_integ_count) / total_mean_abs; - _status.count = total_count; - _status.count_opt = total_count_opt; - _status.count_after_cuts = total_count_after_cuts; - _status.count_after_cuts_opt = total_count_after_cuts_opt; - _status.iterations = iterations; - _status.optimized = optimized; +void EventGenerator::update_integral() { + update_integral_status(); + update_integral_fractions(false); } -void EventGenerator::update_integral_fractions() { - // With an integral known from an earlier run, the sum of the channel estimates - // normalizes the fractions; without, this is _status.mean_abs. - double total_estimate = _status.mean_abs; - if (std::any_of(_channels.begin(), _channels.end(), [](auto& channel) { - return channel->has_abs_integral_prior(); - })) { - total_estimate = 0.; - for (auto& channel : _channels) { - total_estimate += channel->abs_integral_estimate(); - } +void EventGenerator::update_integral_status() { + sum_channel_status(_channels, _status); +} + +void EventGenerator::update_integral_fractions(bool final_step) { + // the sum of the channel estimates normalizes the fractions + double total_estimate = 0.; + for (auto& channel : _channels) { + total_estimate += channel->abs_integral_estimate(final_step); } for (auto [channel, integral_fraction] : zip(_channels, _channel_integral_fractions)) { - integral_fraction = channel->abs_integral_estimate() / total_estimate; + integral_fraction = channel->abs_integral_estimate(final_step) / total_estimate; } // Distribute events between channels, ensure sum is exactly target_count @@ -658,7 +704,15 @@ void EventGenerator::update_integral_fractions() { for (auto [channel, count, fraction] : zip(_channels, counts, _channel_integral_fractions)) { if (_generate_mode == GenerateMode::fix_max_weights) { - count = fraction > 0 ? _config.freeze_max_weight_after : 0; + // the maxima of all steps have to freeze + std::size_t freeze_count = + channel->first_pass_freeze_max_weight_after().value_or( + _config.freeze_max_weight_after + ); + for (auto& step : channel->weighting_steps()) { + freeze_count = std::max(freeze_count, step.freeze_max_weight_after); + } + count = fraction > 0 ? freeze_count : 0; } channel->set_target_count(count); } @@ -696,7 +750,7 @@ void EventGenerator::combine_to_compact_npy( _systematics = systematics; _event_histograms = histograms; MixMaxRandom select_rng(DerivedSeed(_seed, DerivedSeed::combine_select)); - auto [channel_data, particle_count, norm_factor] = init_combine(); + auto [channel_data, particle_count, norm_factors] = init_combine(); DataLayout layout = combined_layout(); DataLayout out_layout = output_layout( EventRecord::f_weight | EventRecord::f_subproc_index | @@ -712,7 +766,7 @@ void EventGenerator::combine_to_compact_npy( print_combine_init(); while (true) { _abort_check_function(); - read_and_combine(channel_data, buffer, norm_factor, select_rng); + read_and_combine(channel_data, buffer, norm_factors, select_rng); if (buffer.event_count() == 0) { break; } @@ -745,7 +799,7 @@ void EventGenerator::combine_to_lhe_npy( std::uint64_t seed = _seed; MixMaxRandom select_rng(DerivedSeed(seed, DerivedSeed::combine_select)); MixMaxRandom rand_gen(DerivedSeed(seed, DerivedSeed::lhe_complete)); - auto [channel_data, particle_count, norm_factor] = init_combine(); + auto [channel_data, particle_count, norm_factors] = init_combine(); DataLayout in_layout = combined_layout(); DataLayout out_layout = output_layout(EventRecord::f_lhe_event, ParticleRecord::f_lhe_particle); @@ -761,7 +815,7 @@ void EventGenerator::combine_to_lhe_npy( print_combine_init(); while (true) { _abort_check_function(); - read_and_combine(channel_data, buffer, norm_factor, select_rng); + read_and_combine(channel_data, buffer, norm_factors, select_rng); if (buffer.event_count() == 0) { break; } @@ -769,20 +823,18 @@ void EventGenerator::combine_to_lhe_npy( process_combined_batch(buffer, syst_weights); buffer_out.resize(buffer.event_count()); for (std::size_t i = 0; i < buffer.event_count(); ++i) { - fill_lhe_event(lhe_completer, lhe_event, buffer, i, rand_gen); + fill_lhe_event( + lhe_completer, + lhe_event, + buffer, + i, + rand_gen, + _channels.at(0)->event_layout_extra_flags(), + _channels.at(0)->particle_layout_extra_flags() + ); buffer_out.event(i).from_lhe_event(lhe_event); copy_event_fields(buffer, i, buffer_out, i, syst_weights); - std::size_t j = 0; - for (; j < lhe_event.particles.size(); ++j) { - buffer_out.particle(i, j).from_lhe_particle( - lhe_event.particles[j], lhe_event.has_cluster_scales - ); - } - for (; j < lhe_completer.max_particle_count(); ++j) { - buffer_out.particle(i, j).from_lhe_particle( - LHEParticle{}, lhe_event.has_cluster_scales - ); - } + write_lhe_particles(lhe_event, buffer_out, i); } event_file.write(buffer_out); if (event_count - last_update_count > 10000) { @@ -808,7 +860,7 @@ void EventGenerator::combine_to_lhe( std::uint64_t seed = _seed; MixMaxRandom select_rng(DerivedSeed(seed, DerivedSeed::combine_select)); ThreadPool pool(_config.combine_thread_count); - auto [channel_data, particle_count, norm_factor] = init_combine(); + auto [channel_data, particle_count, norm_factors] = init_combine(); struct CombineBuffers { EventBuffer in_buffer; std::string out_buffer; @@ -851,7 +903,7 @@ void EventGenerator::combine_to_lhe( while (idle_buffers.size() > 0 && !done) { std::size_t slot = idle_buffers.back(); auto& [in_buffer, out_buffer, syst_weights] = buffers.at(slot); - read_and_combine(channel_data, in_buffer, norm_factor, select_rng); + read_and_combine(channel_data, in_buffer, norm_factors, select_rng); if (in_buffer.event_count() == 0) { done = true; break; @@ -877,7 +929,13 @@ void EventGenerator::combine_to_lhe( process_combined_batch(in_buffer, syst_weights); for (std::size_t i = 0; i < in_buffer.event_count(); ++i) { fill_lhe_event( - lhe_completer, lhe_event, in_buffer, i, rand_gen + lhe_completer, + lhe_event, + in_buffer, + i, + rand_gen, + _channels.at(0)->event_layout_extra_flags(), + _channels.at(0)->particle_layout_extra_flags() ); fill_systematics( lhe_event, in_buffer, i, syst_weights, syst_ids @@ -916,43 +974,10 @@ void EventGenerator::combine_to_lhe( _event_histograms = nullptr; } -namespace { - -// The MLM alpha_s vertex scales and pdf ladder of any channel: a merged run -// mixes a 2 -> 1 sample, which has no vertex, with the jet samples, which do, -// and channel 0 may be either. -template -int any_alphas_scales(const Channels& channels) { - int flags = 0; - for (auto& channel : channels) { - flags |= channel->particle_layout_extra_flags() & - (ParticleRecord::f_alphas_scales | ParticleRecord::f_pdf_ladder); - } - return flags; -} - -// the pdf ladder fields of one particle, or zeros (an inactive slot) -void copy_ladder(ParticleRecord in, ParticleRecord out, bool from_in) { - out.ladder_x() = from_in ? in.ladder_x().value() : 0.; - out.ladder_pdg() = from_in ? in.ladder_pdg().value() : 0.; - out.ladder_q() = from_in ? in.ladder_q().value() : 0.; - out.ladder_info() = from_in ? in.ladder_info().value() : 0.; - out.ladder_lead_q() = from_in ? in.ladder_lead_q().value() : 0.; -} - -} // namespace - DataLayout EventGenerator::combined_layout() const { - return DataLayout( - EventRecord::layout( - EventRecord::f_weight | EventRecord::f_subproc_index | - EventRecord::f_event_data | - (_channels.at(0)->event_layout_extra_flags() & ~EventRecord::f_bias) - ), - ParticleRecord::layout( - ParticleRecord::f_particle_data | - _channels.at(0)->particle_layout_extra_flags() | any_alphas_scales(_channels) - ) + return combined_event_layout( + _channels.at(0)->event_layout_extra_flags(), + _channels.at(0)->particle_layout_extra_flags() | mlm_particle_flags(_channels) ); } @@ -1172,12 +1197,16 @@ std::vector EventGenerator::histograms() const { return hists; } -std::tuple, std::size_t, double> +std::tuple< + std::vector, + std::size_t, + std::vector> EventGenerator::init_combine() { std::vector channel_data; std::size_t count_sum = 0; std::size_t particle_count = 0; double weight_sum = 0.; + std::vector channel_weight_sums; for (auto& channel : _channels) { particle_count = std::max(particle_count, channel->event_file().particle_count()); @@ -1186,42 +1215,48 @@ EventGenerator::init_combine() { std::size_t count = channel->status().count_target; count_sum += count; channel->event_file().seek(0); - weight_sum += channel->channel_weight_sum(count); + channel_weight_sums.push_back(channel->channel_weight_sum(count)); + weight_sum += channel_weight_sums.back(); channel->weight_file().seek(0); channel_data.push_back({ .cum_count = count_sum, .event_buffer = EventBuffer( 0, channel->event_file().particle_count(), channel->event_file_layout() ), - .weight_buffer = EventBuffer(0, 0, weight_file_layout), + .weight_buffer = EventBuffer( + 0, channel->weight_record_count(), channel->weight_layout() + ), .buffer_index = 0, }); } - return {channel_data, particle_count, _status.mean_abs * count_sum / weight_sum}; + std::vector norm_factors; + for (auto [channel, channel_weight_sum] : zip(_channels, channel_weight_sums)) { + // After reweighting steps the mean event weights differ between + // channels, so the events of each channel are normalized separately + norm_factors.push_back( + _pipeline + ? (channel_weight_sum > 0 + ? channel->status().mean_abs * count_sum / channel_weight_sum + : 0.) + : _status.mean_abs * count_sum / weight_sum + ); + } + return {channel_data, particle_count, norm_factors}; } void EventGenerator::read_and_combine( std::vector& channel_data, EventBuffer& buffer, - double norm_factor, + const std::vector& norm_factors, MixMaxRandom& rand_gen ) { std::size_t batch_size = 1000; std::size_t event_count = std::min(batch_size, channel_data.back().cum_count); buffer.resize(event_count); - bool has_beam1 = _channels.at(0)->event_layout_extra_flags() & EventRecord::f_beam1; - bool has_beam2 = _channels.at(0)->event_layout_extra_flags() & EventRecord::f_beam2; - bool has_partial = - _channels.at(0)->event_layout_extra_flags() & EventRecord::f_partial_weights; - bool has_clustering = - _channels.at(0)->particle_layout_extra_flags() & ParticleRecord::f_clustering; - int mlm_flags = any_alphas_scales(_channels); - bool has_alphas_scales = mlm_flags & ParticleRecord::f_alphas_scales; - bool has_ladder = mlm_flags & ParticleRecord::f_pdf_ladder; - bool has_subproc_index = - _channels.at(0)->event_layout_extra_flags() & EventRecord::f_subproc_index; - bool has_bias = _channels.at(0)->event_layout_extra_flags() & EventRecord::f_bias; + int extra_flags = _channels.at(0)->event_layout_extra_flags(); + int particle_flags = + _channels.at(0)->particle_layout_extra_flags() | mlm_particle_flags(_channels); for (std::size_t event_index = 0; event_index < event_count; ++event_index) { std::size_t random_index = rand_gen.generate_int(channel_data.back().cum_count); @@ -1234,11 +1269,8 @@ void EventGenerator::read_and_combine( std::for_each(sampled_chan, channel_data.end(), [](auto& chan) { --chan.cum_count; }); - auto& channel = _channels.at(sampled_chan - channel_data.begin()); - bool channel_has_alphas_scales = - channel->particle_layout_extra_flags() & ParticleRecord::f_alphas_scales; - bool channel_has_ladder = - channel->particle_layout_extra_flags() & ParticleRecord::f_pdf_ladder; + std::size_t channel_index = sampled_chan - channel_data.begin(); + auto& channel = _channels.at(channel_index); double weight = 0.; while (true) { @@ -1256,140 +1288,21 @@ void EventGenerator::read_and_combine( ++sampled_chan->buffer_index; } - auto event_in = sampled_chan->event_buffer.event(sampled_chan->buffer_index); - auto event_out = buffer.event(event_index); - event_out.weight() = - std::copysign( - std::max(1., std::abs(weight) / channel->max_weight()), weight - ) * - norm_factor; - if (has_bias) { - // unweighted with the bias b, so the event carries 1 / b - event_out.weight() = event_out.weight() / event_in.bias().value(); - } - event_out.subprocess_index() = has_subproc_index - ? event_in.subprocess_index().value() - : static_cast(channel->status().subprocess); - event_out.diagram_index() = event_in.diagram_index(); - event_out.color_index() = event_in.color_index(); - event_out.flavor_index() = event_in.flavor_index(); - event_out.helicity_index() = event_in.helicity_index(); - event_out.ren_scale() = event_in.ren_scale(); - event_out.alpha_qcd() = event_in.alpha_qcd(); - - if (has_beam1) { - event_out.x1() = event_in.x1(); - event_out.fact_scale1() = event_in.fact_scale1(); - } - if (has_beam2) { - event_out.x2() = event_in.x2(); - event_out.fact_scale2() = event_in.fact_scale2(); - } - if (has_partial) { - event_out.partial_weight_product() = event_in.partial_weight_product(); - } - - std::size_t i = 0; - for (; i < sampled_chan->event_buffer.particle_count(); ++i) { - auto particle_in = - sampled_chan->event_buffer.particle(sampled_chan->buffer_index, i); - auto particle_out = buffer.particle(event_index, i); - particle_out.energy() = particle_in.energy(); - particle_out.px() = particle_in.px(); - particle_out.py() = particle_in.py(); - particle_out.pz() = particle_in.pz(); - if (has_clustering) { - particle_out.cluster_scale() = particle_in.cluster_scale(); - } - if (has_alphas_scales) { - // zero from a channel without vertices (2 -> 1) - particle_out.alphas_scale() = channel_has_alphas_scales - ? particle_in.alphas_scale().value() - : 0.; - } - if (has_ladder) { - copy_ladder(particle_in, particle_out, channel_has_ladder); - } - } - for (; i < buffer.particle_count(); ++i) { - auto particle_out = buffer.particle(event_index, i); - particle_out.energy() = 0.; - particle_out.px() = 0.; - particle_out.py() = 0.; - particle_out.pz() = 0.; - // the buffer is reused: a padding slot must not keep the scales of - // an earlier, higher-multiplicity event - if (has_clustering) { - particle_out.cluster_scale() = 0.; - } - if (has_alphas_scales) { - particle_out.alphas_scale() = 0.; - } - if (has_ladder) { - copy_ladder(particle_out, particle_out, false); - } - } - ++sampled_chan->buffer_index; - } -} - -void EventGenerator::fill_lhe_event( - LHECompleter& lhe_completer, - LHEEvent& lhe_event, - EventBuffer& buffer, - std::size_t event_index, - MixMaxRandom& rand_gen -) { - EventRecord event_in = buffer.event(event_index); - lhe_event.weight = event_in.weight(); - lhe_event.process_id = 0; - lhe_event.scale = event_in.ren_scale(); - lhe_event.alpha_qed = 0; // TODO: populate this - lhe_event.alpha_qcd = event_in.alpha_qcd(); - // Both scales, as LHEF3 attributes: SCALUP holds only one number - // and the codes do not agree on which, so write them out explicitly. - int scale_flags = _channels.at(0)->event_layout_extra_flags(); - lhe_event.mu_r = event_in.ren_scale(); - lhe_event.mu_f1 = (scale_flags & EventRecord::f_beam1) - ? double(event_in.fact_scale1()) - : double(event_in.ren_scale()); - lhe_event.mu_f2 = (scale_flags & EventRecord::f_beam2) - ? double(event_in.fact_scale2()) - : lhe_event.mu_f1; - lhe_event.has_scales = true; - lhe_event.particles.clear(); - for (std::size_t i = 0; i < buffer.particle_count(); ++i) { - auto particle_in = buffer.particle(event_index, i); - if (particle_in.energy() == 0.) { - break; - } - lhe_event.particles.push_back( - LHEParticle{ - .px = particle_in.px(), - .py = particle_in.py(), - .pz = particle_in.pz(), - .energy = particle_in.energy(), - } + copy_channel_event( + sampled_chan->event_buffer, + sampled_chan->buffer_index, + buffer, + event_index, + channel->unweighted_weight( + sampled_chan->weight_buffer, sampled_chan->buffer_index + ) * norm_factors.at(channel_index), + static_cast(channel->status().subprocess), + extra_flags, + channel->particle_layout_extra_flags(), + particle_flags ); + ++sampled_chan->buffer_index; } - bool has_clustering = - _channels.at(0)->particle_layout_extra_flags() & ParticleRecord::f_clustering; - lhe_event.has_cluster_scales = has_clustering; - if (has_clustering) { - for (std::size_t i = 0; i < lhe_event.particles.size(); ++i) { - lhe_event.particles.at(i).cluster_scale = - buffer.particle(event_index, i).cluster_scale(); - } - } - lhe_completer.complete_event_data( - lhe_event, - event_in.subprocess_index(), - event_in.diagram_index(), - event_in.color_index(), - event_in.flavor_index(), - event_in.helicity_index(), - rand_gen - ); } void EventGenerator::init_status(const std::string& status) { @@ -1606,6 +1519,27 @@ void EventGenerator::print_gen_init() { "Unweighted events:", "Run time:"} ); + if (_pipeline) { + _pretty_box_upper = PrettyBox( + _generate_mode == GenerateMode::fix_max_weights + ? "Integration and max. weight estimation" + : "Integration and unweighting", + 8, + {19, 0}, + offset + ); + _pretty_box_upper.set_column( + 0, + {"Result:", + "Rel. error:", + "Rel. stddev:", + "Number of events:", + "Unweighting eff.:", + " per pass:", + "Unweighted events:", + "Run time:"} + ); + } } _pretty_box_upper.print_first(); if (_channels.size() > 1) { @@ -1613,6 +1547,33 @@ void EventGenerator::print_gen_init() { } } +std::tuple> +EventGenerator::pass_efficiencies() const { + // events accepted per pass, a channel with fewer steps passes the others + std::vector accepted; + for (auto& channel : _channels) { + auto counts = channel->step_accepted_counts(); + if (counts.size() > accepted.size()) { + accepted.resize(counts.size(), 0.); + } + for (std::size_t i = 0; i < accepted.size(); ++i) { + accepted.at(i) += counts.at(std::min(i, counts.size() - 1)); + } + } + // the last pass ends with the counts of the final events, so that the + // product is the overall efficiency + accepted.back() = _status.count_unweighted; + std::vector later_passes; + for (std::size_t i = 1; i < accepted.size(); ++i) { + later_passes.push_back(accepted.at(i) / accepted.at(i - 1)); + } + return { + accepted.at(0) / _status.count_opt, + accepted.at(0) / _status.count_after_cuts_opt, + later_passes + }; +} + void EventGenerator::print_gen_update(bool done) { if (done) { add_timing_data("generate"); @@ -1634,7 +1595,7 @@ void EventGenerator::print_gen_update_pretty(bool done) { _last_print_time = now; bool optimize_only = _generate_mode == GenerateMode::optimize_only; - std::string int_str, rel_str, rsd_str, uweff_str, count_str; + std::string int_str, rel_str, rsd_str, uweff_str, pass_str, count_str; count_str = std::format( "{} before cuts, {} after", format_si_prefix(_status.count), @@ -1653,6 +1614,15 @@ void EventGenerator::print_gen_update_pretty(bool done) { _status.count_unweighted / _status.count_opt, _status.count_unweighted / _status.count_after_cuts_opt ); + if (_pipeline) { + auto [first, first_after_cuts, later] = pass_efficiencies(); + pass_str = + std::format("1: {:.5f} ({:.5f} after cuts)", first, first_after_cuts); + for (std::size_t i = 0; double eff : later) { + pass_str += std::format(", {}: {:.5f}", i + 2, eff); + ++i; + } + } } // the progress of the run: unweighted events, or optimized channels std::string progress_str; @@ -1686,6 +1656,18 @@ void EventGenerator::print_gen_update_pretty(bool done) { _pretty_box_upper.set_column( 1, {int_str, rel_str, rsd_str, count_str, progress_str, time_str} ); + } else if (_pipeline) { + _pretty_box_upper.set_column( + 1, + {int_str, + rel_str, + rsd_str, + count_str, + uweff_str, + pass_str, + progress_str, + time_str} + ); } else { _pretty_box_upper.set_column( 1, {int_str, rel_str, rsd_str, count_str, uweff_str, progress_str, time_str} @@ -1793,11 +1775,23 @@ void EventGenerator::print_gen_update_log(bool done) { ) ); } else { + std::string pass_efficiency_str; + if (_pipeline) { + auto [first, first_after_cuts, later] = pass_efficiencies(); + pass_efficiency_str = std::format( + ", unw. eff. per pass: {:.5f} ({:.5f} after cuts)", + first, + first_after_cuts + ); + for (double eff : later) { + pass_efficiency_str += std::format(", {:.5f}", eff); + } + } Logger::info( std::format( "generating, events: {} / {}, integral: {}, rel. error: {}, " "RSD: {:.3f}, samps: {}, samps. after cuts: {}, " - "unw. eff.: {:.5f}, unw. eff. after cuts: {:.5f}, time: {:%H:%M:%S}", + "unw. eff.: {:.5f}, unw. eff. after cuts: {:.5f}{}, time: {:%H:%M:%S}", format_si_prefix(_status.count_unweighted), format_si_prefix(_status.count_target), format_with_error(_status.mean, _status.error), @@ -1807,6 +1801,7 @@ void EventGenerator::print_gen_update_log(bool done) { format_si_prefix(_status.count_after_cuts), _status.count_unweighted / _status.count_opt, _status.count_unweighted / _status.count_after_cuts_opt, + pass_efficiency_str, std::chrono::round(now - _start_time) ) ); diff --git a/madspace/src/driver/event_stream.cpp b/madspace/src/driver/event_stream.cpp new file mode 100644 index 0000000000..28becf694a --- /dev/null +++ b/madspace/src/driver/event_stream.cpp @@ -0,0 +1,548 @@ +#include "madspace/driver/event_stream.hpp" + +#include +#include +#include +#include + +#include "madspace/driver/event_utils.hpp" +#include "madspace/util.hpp" + +using namespace madspace; +using namespace std::chrono_literals; + +namespace { + +// separates the stream's random streams from those of the combine pass +constexpr std::size_t stream_rng_index = 1; + +} // namespace + +EventStream::EventStream( + ContextPtr context, + const std::vector>& channels, + std::uint64_t seed, + const LHECompleter& lhe_completer, + const GeneratorConfig& config, + const std::vector& groups +) : + _context(context), + _channels(channels), + _seed(seed), + _lhe_completer(lhe_completer), + _config(config), + _combined_layout([&] { + if (channels.empty()) { + throw std::invalid_argument("EventStream needs at least one channel"); + } + return combined_event_layout( + channels.at(0)->event_layout_extra_flags(), + channels.at(0)->particle_layout_extra_flags() | mlm_particle_flags(channels) + ); + }()), + _select_rng(DerivedSeed(seed, DerivedSeed::combine_select, 0, 0, stream_rng_index)), + _lhe_rng(DerivedSeed(seed, DerivedSeed::lhe_complete, 0, 0, stream_rng_index)), + _channel_states(channels.size()), + _lookahead( + context->device()->device_type() == DeviceType::cpu + ? config.cpu_batch_size + : config.gpu_batch_size + ) { + if (!groups.empty()) { + _pipeline = std::make_unique( + std::vector{context}, + channels, + groups, + config, + [this](StepEvents& step_events) { close_step_events(step_events); } + ); + } + bool any_active = false; + for (auto [channel, state] : zip(_channels, _channel_states)) { + state.active = channel->max_weight_fixed() && channel->step_max_weights_fixed(); + any_active |= state.active; + } + if (!any_active) { + throw std::invalid_argument( + "EventStream needs channels with a fixed maximum weight" + ); + } + _coordinator = std::thread([this] { run_coordinator(); }); +} + +EventStream::~EventStream() { close(); } + +void EventStream::close() { + { + std::unique_lock lock(_mutex); + _stopping = true; + } + _coordinator_cv.notify_all(); + _consumer_cv.notify_all(); + if (_coordinator.joinable()) { + _coordinator.join(); + } +} + +void EventStream::run_coordinator() { + try { + // dispatches new jobs when the consumer asked for them, stops on close() + auto poll = [this] { + std::unique_lock lock(_mutex); + if (_stopping) { + throw StopRequest{}; + } + if (_dispatch_requested) { + _dispatch_requested = false; + dispatch_jobs(); + } + }; + while (true) { + { + std::unique_lock lock(_mutex); + if (_stopping) { + break; + } + _dispatch_requested = false; + dispatch_jobs(); + if (_pipeline) { + // partial batches only once no events are coming anymore + _results_pending += + _pipeline->start_jobs(_result_queue, _running_jobs.empty()); + } + if (_results_pending == 0) { + _coordinator_cv.wait(lock, [this] { + return _stopping || _dispatch_requested; + }); + continue; + } + } + auto result = _result_queue.wait(poll, 10ms); + --_results_pending; + if (result.exception) { + std::rethrow_exception(result.exception); + } + handle_result(result.id); + } + } catch (StopRequest&) { + } catch (...) { + std::unique_lock lock(_mutex); + _error = std::current_exception(); + } + _consumer_cv.notify_all(); + _result_queue.discard(_results_pending, nullptr); + _results_pending = 0; + _running_jobs.clear(); + if (_pipeline) { + _pipeline->reset(); + } + _step_jobs.clear(); +} + +double EventStream::expected_job_events(std::size_t channel_index) const { + auto& state = _channel_states.at(channel_index); + if (state.completed_jobs > 0) { + return std::max(state.completed_events, 1.) / state.completed_jobs; + } + // before the first job, assume a single event to stay on the safe side + return 1.; +} + +// Called with _mutex held. Which jobs are started depends on timing, but not +// what they contain: that is fixed by the channel's job sequence number. +void EventStream::dispatch_jobs() { + std::size_t max_jobs = 2 * _context->thread_pool().thread_count(); + if (_probabilities_dirty) { + update_probabilities(); + } + while (_running_jobs.size() < max_jobs) { + // the channel furthest below its share of the lookahead + std::size_t best_index = _channels.size(); + double best_ratio = 1.; + double prev_prob = 0.; + for (std::size_t i = 0; i < _channels.size(); ++i) { + double prob = _cum_probabilities.at(i) - prev_prob; + prev_prob = _cum_probabilities.at(i); + auto& state = _channel_states.at(i); + if (!state.active) { + continue; + } + // bounds the memory of channels that rarely yield events + if (state.in_flight + state.unopened_jobs >= max_jobs) { + continue; + } + double available = + state.buffered_events + state.in_flight * expected_job_events(i); + double target = prob * _lookahead; + double ratio; + if (state.buffered_events == 0 && state.in_flight == 0 && + state.unopened_jobs == 0) { + // every channel keeps at least one job buffered or running + ratio = -1.; + } else if (available < target) { + ratio = available / target; + } else { + continue; + } + if (ratio < best_ratio) { + best_ratio = ratio; + best_index = i; + } + } + if (best_index == _channels.size()) { + break; + } + + std::size_t job_id = _job_id++; + auto& job = + _running_jobs + .emplace( + job_id, + GeneratorBatchJob{ + .channel_index = best_index, + .unweight = true, + .batch_event_count = 0, + .split_job_count = 1, + .context_index = 0, + .job_id = job_id, + .is_vegas_batch = false, + } + ) + .first->second; + auto& state = _channel_states.at(best_index); + state.launch_order.push_back(job_id); + ++state.in_flight; + _channels.at(best_index)->start_job(job, _result_queue, _seed, false, 0); + ++_results_pending; + } +} + +void EventStream::handle_result(std::size_t job_id) { + if (WeightingPipeline::owns_job(job_id)) { + _pipeline->commit_job(job_id); + return; + } + auto& job = _running_jobs.at(job_id); + auto& channel = _channels.at(job.channel_index); + if (job.unweighted_events.empty()) { + // generation stage done, the max weight is fixed so unweight right away + std::unique_lock lock(_mutex); + channel->start_unweight_job(job, _result_queue); + ++_results_pending; + return; + } + if (_pipeline && _pipeline->contains(job.channel_index)) { + // keep what open_job() needs until the events are through the steps + auto& step_job = _step_jobs[job_id]; + step_job.channel_index = job.channel_index; + step_job.weights = job.weights; + step_job.biases = job.biases; + step_job.requested_event_count = job.requested_event_count; + _pipeline->add_events(job); + _running_jobs.erase(job_id); + return; + } + + ClosedJob closed_job{ + .job = {}, + .events = + EventBuffer(0, channel->particle_count(), channel->event_file_layout()), + .weights = EventBuffer(0, 0, weight_file_layout), + }; + channel->fill_event_buffers( + job.unweighted_events, closed_job.events, closed_job.weights + ); + // only the weights are needed to integrate the job later + closed_job.job.channel_index = job.channel_index; + closed_job.job.weights = job.weights; + closed_job.job.biases = job.biases; + closed_job.job.requested_event_count = job.requested_event_count; + std::size_t channel_index = job.channel_index; + _running_jobs.erase(job_id); + close_job(job_id, channel_index, std::move(closed_job)); +} + +void EventStream::close_step_events(StepEvents& step_events) { + auto& channel = _channels.at(step_events.channel_index); + auto node = _step_jobs.extract(step_events.job_id); + ClosedJob closed_job{ + .job = std::move(node.mapped()), + .events = + EventBuffer(0, channel->particle_count(), channel->event_file_layout()), + .weights = + EventBuffer(0, channel->weight_record_count(), channel->weight_layout()), + .step_integral = channel->step_integral(step_events), + }; + channel->fill_step_buffers(step_events, closed_job.events, closed_job.weights); + for (std::size_t i = 0; i < closed_job.weights.event_count(); ++i) { + closed_job.abs_weight_sum += + std::abs(channel->unweighted_weight(closed_job.weights, i)); + } + close_job(step_events.job_id, step_events.channel_index, std::move(closed_job)); +} + +void EventStream::close_job( + std::size_t job_id, std::size_t channel_index, ClosedJob&& closed_job +) { + std::unique_lock lock(_mutex); + auto& state = _channel_states.at(channel_index); + ++state.completed_jobs; + state.completed_events += closed_job.events.event_count(); + state.finished.emplace(job_id, std::move(closed_job)); + // close the jobs in launch order + bool closed_any = false; + while (!state.launch_order.empty()) { + auto node = state.finished.extract(state.launch_order.front()); + if (node.empty()) { + break; + } + state.launch_order.pop_front(); + --state.in_flight; + state.buffered_events += node.mapped().events.event_count(); + ++state.unopened_jobs; + state.closed.push_back(std::move(node.mapped())); + closed_any = true; + } + if (closed_any) { + _consumer_cv.notify_all(); + } +} + +// Called with _mutex held. +void EventStream::update_probabilities() { + _cum_probabilities.resize(_channels.size()); + double total = 0.; + for (std::size_t i = 0; i < _channels.size(); ++i) { + auto& state = _channel_states.at(i); + double prob = 0.; + if (state.active) { + // divided by the mean weight of the unweighted events (> 1 with + // overweight events), so that the weight sums are proportional to + // the cross sections + double mean_weight = state.opened_capped_sum > 0 + ? state.opened_abs_sum / state.opened_capped_sum + : 1.; + prob = _channels.at(i)->abs_integral_estimate() / mean_weight; + if (!std::isfinite(prob) || prob < 0) { + prob = 0.; + } + } + total += prob; + _cum_probabilities.at(i) = total; + } + if (total > 0) { + for (auto& prob : _cum_probabilities) { + prob /= total; + } + } + _probabilities_dirty = false; +} + +std::size_t EventStream::select_channel() { + if (_probabilities_dirty) { + update_probabilities(); + } + if (_cum_probabilities.back() <= 0) { + throw std::runtime_error("EventStream: all channels have a zero cross section"); + } + // the first channel whose cumulative probability exceeds r, so channels + // with zero probability are never picked + double r = _select_rng.generate_double() * _cum_probabilities.back(); + auto it = std::upper_bound(_cum_probabilities.begin(), _cum_probabilities.end(), r); + if (it == _cum_probabilities.end()) { + it = std::lower_bound( + _cum_probabilities.begin(), + _cum_probabilities.end(), + _cum_probabilities.back() + ); + } + return it - _cum_probabilities.begin(); +} + +// Called with _mutex held. +void EventStream::open_job(std::size_t channel_index, ClosedJob& closed_job) { + auto& channel = _channels.at(channel_index); + auto& state = _channel_states.at(channel_index); + channel->integrate(closed_job.job); + if (!channel->weighting_steps().empty()) { + channel->add_step_integral(closed_job.step_integral); + // the mean weight of the events after the steps + state.opened_abs_sum += closed_job.abs_weight_sum; + state.opened_capped_sum += closed_job.events.event_count(); + } else { + double max_weight = channel->max_weight(); + auto w_view = closed_job.job.weights.view(); + for (std::size_t i = 0; i < w_view.size(); ++i) { + double w = std::abs(w_view[i]); + state.opened_abs_sum += w; + state.opened_capped_sum += std::min(w, max_weight); + } + } + ++state.opened_jobs; + --state.unopened_jobs; + closed_job.opened = true; + _probabilities_dirty = true; +} + +void EventStream::wait_for_events( + std::size_t channel_index, std::unique_lock& lock +) { + _dispatch_requested = true; + _coordinator_cv.notify_all(); + auto& state = _channel_states.at(channel_index); + while (state.closed.empty()) { + if (_error) { + std::rethrow_exception(_error); + } + if (_stopping) { + throw std::runtime_error("EventStream was closed"); + } + if (!_consumer_cv.wait_for(lock, 100ms, [&] { + return !state.closed.empty() || _error || _stopping; + })) { + lock.unlock(); + _abort_check_function(); + lock.lock(); + } + } +} + +// The first job of the channel's buffer that still has events, opened. +// Called with _mutex held. +EventStream::ClosedJob& +EventStream::next_job(std::size_t channel_index, std::unique_lock& lock) { + auto& state = _channel_states.at(channel_index); + while (true) { + if (state.closed.empty()) { + wait_for_events(channel_index, lock); + } + auto& front = state.closed.front(); + if (!front.opened) { + open_job(channel_index, front); + } + if (front.next_event < front.events.event_count()) { + return front; + } + state.closed.pop_front(); + } +} + +std::vector EventStream::next_events(std::size_t count) { + std::unique_lock consumer_lock(_consumer_mutex); + EventBuffer buffer(count, max_particle_count(), _combined_layout); + { + std::unique_lock lock(_mutex); + if (_error) { + std::rethrow_exception(_error); + } + if (_stopping) { + throw std::runtime_error("EventStream was closed"); + } + _lookahead = std::max(_lookahead, 2 * count); + if (!_started) { + // every channel's estimate starts from one job (also those + // without an integral prior), so the first events are already + // drawn with all channels known + for (std::size_t i = 0; i < _channels.size(); ++i) { + if (_channel_states.at(i).active) { + next_job(i, lock); + } + } + _started = true; + } + for (std::size_t i = 0; i < count; ++i) { + std::size_t channel_index = select_channel(); + auto& channel = _channels.at(channel_index); + auto& state = _channel_states.at(channel_index); + auto& job = next_job(channel_index, lock); + std::size_t event_index = job.next_event++; + --state.buffered_events; + ++state.returned_events; + copy_channel_event( + job.events, + event_index, + buffer, + i, + channel->unweighted_weight(job.weights, event_index), + static_cast(channel->status().subprocess), + channel->event_layout_extra_flags(), + channel->particle_layout_extra_flags(), + _channels.at(0)->particle_layout_extra_flags() | + mlm_particle_flags(_channels) + ); + // keep the buffers filled while the batch is assembled + if (state.buffered_events == 0) { + _dispatch_requested = true; + _coordinator_cv.notify_all(); + } + } + _event_count += count; + _dispatch_requested = true; + } + _coordinator_cv.notify_all(); + + std::vector events(count); + for (std::size_t i = 0; i < count; ++i) { + fill_lhe_event( + _lhe_completer, + events.at(i), + buffer, + i, + _lhe_rng, + _channels.at(0)->event_layout_extra_flags(), + _channels.at(0)->particle_layout_extra_flags() + ); + } + return events; +} + +GeneratorStatus EventStream::status() const { + std::unique_lock lock(_mutex); + GeneratorStatus status{}; + // the inactive channels return no events + std::vector> active_channels; + for (auto [channel, state] : zip(_channels, _channel_states)) { + if (state.active) { + active_channels.push_back(channel); + } + } + sum_channel_status(active_channels, status); + status.count_unweighted = _event_count; + status.count_target = 0; + status.optimized = true; + status.done = false; + return status; +} + +std::vector EventStream::channel_status() const { + std::unique_lock lock(_mutex); + std::vector ret; + for (auto [channel, state] : zip(_channels, _channel_states)) { + auto status = channel->status(); + status.count_unweighted = state.returned_events; + status.count_target = 0; + ret.push_back(status); + } + return ret; +} + +std::vector EventStream::channel_probabilities() const { + std::unique_lock lock(_mutex); + // const: compute from scratch instead of updating the cache + auto& self = const_cast(*this); + if (self._probabilities_dirty) { + self.update_probabilities(); + } + std::vector probs; + double prev = 0.; + for (double cum_prob : _cum_probabilities) { + probs.push_back(cum_prob - prev); + prev = cum_prob; + } + return probs; +} + +std::size_t EventStream::event_count() const { + std::unique_lock lock(_mutex); + return _event_count; +} diff --git a/madspace/src/driver/event_utils.cpp b/madspace/src/driver/event_utils.cpp new file mode 100644 index 0000000000..c2c314e134 --- /dev/null +++ b/madspace/src/driver/event_utils.cpp @@ -0,0 +1,289 @@ +#include "madspace/driver/event_utils.hpp" + +#include +#include +#include + +using namespace madspace; + +DataLayout +madspace::combined_event_layout(int extra_event_flags, int extra_particle_flags) { + return DataLayout( + EventRecord::layout( + EventRecord::f_weight | EventRecord::f_subproc_index | + EventRecord::f_event_data | (extra_event_flags & ~EventRecord::f_bias) + ), + ParticleRecord::layout(ParticleRecord::f_particle_data | extra_particle_flags) + ); +} + +int madspace::mlm_particle_flags( + const std::vector>& channels +) { + int flags = 0; + for (auto& channel : channels) { + flags |= channel->particle_layout_extra_flags() & + (ParticleRecord::f_alphas_scales | ParticleRecord::f_pdf_ladder); + } + return flags; +} + +namespace { + +// the pdf ladder fields of one particle, or zeros (an inactive slot) +void copy_ladder(ParticleRecord in, ParticleRecord out, bool from_in) { + out.ladder_x() = from_in ? in.ladder_x().value() : 0.; + out.ladder_pdg() = from_in ? in.ladder_pdg().value() : 0.; + out.ladder_q() = from_in ? in.ladder_q().value() : 0.; + out.ladder_info() = from_in ? in.ladder_info().value() : 0.; + out.ladder_lead_q() = from_in ? in.ladder_lead_q().value() : 0.; +} + +} // namespace + +void madspace::copy_channel_event( + EventBuffer& in_buffer, + std::size_t in_index, + EventBuffer& out_buffer, + std::size_t out_index, + double weight, + int subprocess_index, + int extra_flags, + int in_particle_flags, + int out_particle_flags +) { + bool has_clustering = out_particle_flags & ParticleRecord::f_clustering; + bool has_alphas_scales = out_particle_flags & ParticleRecord::f_alphas_scales; + bool has_ladder = out_particle_flags & ParticleRecord::f_pdf_ladder; + bool channel_has_alphas_scales = + in_particle_flags & ParticleRecord::f_alphas_scales; + bool channel_has_ladder = in_particle_flags & ParticleRecord::f_pdf_ladder; + auto event_in = in_buffer.event(in_index); + auto event_out = out_buffer.event(out_index); + // unweighted with the bias b, so the event carries 1 / b + event_out.weight() = + extra_flags & EventRecord::f_bias ? weight / event_in.bias().value() : weight; + event_out.subprocess_index() = extra_flags & EventRecord::f_subproc_index + ? event_in.subprocess_index().value() + : subprocess_index; + event_out.diagram_index() = event_in.diagram_index(); + event_out.color_index() = event_in.color_index(); + event_out.flavor_index() = event_in.flavor_index(); + event_out.helicity_index() = event_in.helicity_index(); + event_out.ren_scale() = event_in.ren_scale(); + event_out.alpha_qcd() = event_in.alpha_qcd(); + + if (extra_flags & EventRecord::f_beam1) { + event_out.x1() = event_in.x1(); + event_out.fact_scale1() = event_in.fact_scale1(); + } + if (extra_flags & EventRecord::f_beam2) { + event_out.x2() = event_in.x2(); + event_out.fact_scale2() = event_in.fact_scale2(); + } + if (extra_flags & EventRecord::f_partial_weights) { + event_out.partial_weight_product() = event_in.partial_weight_product(); + } + + std::size_t i = 0; + for (; i < in_buffer.particle_count(); ++i) { + auto particle_in = in_buffer.particle(in_index, i); + auto particle_out = out_buffer.particle(out_index, i); + particle_out.energy() = particle_in.energy(); + particle_out.px() = particle_in.px(); + particle_out.py() = particle_in.py(); + particle_out.pz() = particle_in.pz(); + if (has_clustering) { + particle_out.cluster_scale() = particle_in.cluster_scale(); + } + if (has_alphas_scales) { + // zero from a channel without vertices (2 -> 1) + particle_out.alphas_scale() = + channel_has_alphas_scales ? particle_in.alphas_scale().value() : 0.; + } + if (has_ladder) { + copy_ladder(particle_in, particle_out, channel_has_ladder); + } + } + for (; i < out_buffer.particle_count(); ++i) { + auto particle_out = out_buffer.particle(out_index, i); + particle_out.energy() = 0.; + particle_out.px() = 0.; + particle_out.py() = 0.; + particle_out.pz() = 0.; + // the buffer is reused: a padding slot must not keep the scales of + // an earlier, higher-multiplicity event + if (has_clustering) { + particle_out.cluster_scale() = 0.; + } + if (has_alphas_scales) { + particle_out.alphas_scale() = 0.; + } + if (has_ladder) { + copy_ladder(particle_out, particle_out, false); + } + } +} + +void madspace::fill_lhe_event( + LHECompleter& lhe_completer, + LHEEvent& lhe_event, + EventBuffer& buffer, + std::size_t event_index, + MixMaxRandom& rand_gen, + int extra_event_flags, + int extra_particle_flags +) { + EventRecord event_in = buffer.event(event_index); + lhe_event.weight = event_in.weight(); + lhe_event.process_id = 0; + lhe_event.scale = event_in.ren_scale(); + lhe_event.alpha_qed = 0; // TODO: populate this + lhe_event.alpha_qcd = event_in.alpha_qcd(); + // Both scales, as LHEF3 attributes: SCALUP holds only one number + // and the codes do not agree on which, so write them out explicitly. + lhe_event.mu_r = event_in.ren_scale(); + lhe_event.mu_f1 = (extra_event_flags & EventRecord::f_beam1) + ? double(event_in.fact_scale1()) + : double(event_in.ren_scale()); + lhe_event.mu_f2 = (extra_event_flags & EventRecord::f_beam2) + ? double(event_in.fact_scale2()) + : lhe_event.mu_f1; + lhe_event.has_scales = true; + lhe_event.particles.clear(); + for (std::size_t i = 0; i < buffer.particle_count(); ++i) { + auto particle_in = buffer.particle(event_index, i); + if (particle_in.energy() == 0.) { + break; + } + lhe_event.particles.push_back( + LHEParticle{ + .px = particle_in.px(), + .py = particle_in.py(), + .pz = particle_in.pz(), + .energy = particle_in.energy(), + } + ); + } + bool has_clustering = extra_particle_flags & ParticleRecord::f_clustering; + lhe_event.has_cluster_scales = has_clustering; + if (has_clustering) { + for (std::size_t i = 0; i < lhe_event.particles.size(); ++i) { + lhe_event.particles.at(i).cluster_scale = + buffer.particle(event_index, i).cluster_scale(); + } + } + lhe_completer.complete_event_data( + lhe_event, + event_in.subprocess_index(), + event_in.diagram_index(), + event_in.color_index(), + event_in.flavor_index(), + event_in.helicity_index(), + rand_gen + ); +} + +void madspace::write_lhe_particles( + const LHEEvent& lhe_event, EventBuffer& buffer, std::size_t event_index +) { + std::size_t i = 0; + for (; i < lhe_event.particles.size(); ++i) { + buffer.particle(event_index, i) + .from_lhe_particle(lhe_event.particles[i], lhe_event.has_cluster_scales); + } + for (; i < buffer.particle_count(); ++i) { + buffer.particle(event_index, i) + .from_lhe_particle(LHEParticle{}, lhe_event.has_cluster_scales); + } +} + +void madspace::sum_channel_status( + const std::vector>& channels, + GeneratorStatus& status +) { + double total_mean = 0., total_var = 0.; + double total_mean_abs = 0., total_var_abs = 0.; + std::size_t total_count = 0, total_count_opt = 0; + std::size_t total_count_after_cuts = 0, total_count_after_cuts_opt = 0; + std::size_t total_integ_count = 0; + std::size_t iterations = 0; + std::size_t count_invalid_step_values = 0; + bool optimized = true; + const ChannelEventGenerator* bad_channel = nullptr; + for (auto& channel : channels) { + auto& chan_status = channel->status(); + auto integral = channel->result_integral(); + // special case for channels with 0 samples, as they have nan variance + if (bad_channel == nullptr && integral.count > 0 && + (!std::isfinite(integral.mean) || !std::isfinite(integral.variance) || + !std::isfinite(integral.mean_abs) || + !std::isfinite(integral.variance_abs))) { + bad_channel = channel.get(); + } + if (channel->uses_integral_prior()) { + // not sampled in this run (no events assigned): use the earlier estimate + auto [mean, error] = channel->integral_prior(); + auto [mean_abs, error_abs] = channel->abs_integral_prior(); + total_mean += mean; + total_var += error * error; + total_mean_abs += mean_abs; + total_var_abs += error_abs * error_abs; + } else { + total_mean += integral.mean; + total_mean_abs += integral.mean_abs; + if (integral.count > 0) { + total_var += integral.variance / integral.count; + total_var_abs += integral.variance_abs / integral.count; + } + } + total_count += chan_status.count; + total_count_opt += chan_status.count_opt; + total_count_after_cuts += chan_status.count_after_cuts; + total_count_after_cuts_opt += chan_status.count_after_cuts_opt; + total_integ_count += integral.count; + iterations = std::max(chan_status.iterations, iterations); + count_invalid_step_values += chan_status.count_invalid_step_values; + if (!chan_status.optimized) { + optimized = false; + } + } + if (bad_channel != nullptr || !std::isfinite(total_mean) || + !std::isfinite(total_mean_abs)) { + std::string where = bad_channel != nullptr + ? std::format( + "channel '{}' after {} samples (mean={}, variance={})", + bad_channel->status().name, + bad_channel->result_integral().count, + bad_channel->result_integral().mean, + bad_channel->result_integral().variance + ) + : std::format("the sum over channels (mean={})", total_mean); + throw std::runtime_error( + std::format( + "non-finite integral in {}. A non-finite weight cannot be recovered " + "from by sampling further, so the integration is aborted. This usually " + "means the matrix element or the phase-space mapping returned nan/inf " + "for " + "some phase-space point -- check the process for a zero or negative " + "width, " + "a parameter point outside the model's validity, or a kinematic " + "configuration at a threshold.", + where + ) + ); + } + + status.mean = total_mean; + status.error = std::sqrt(total_var); + status.mean_abs = total_mean_abs; + status.error_abs = std::sqrt(total_var_abs); + status.rel_std_dev = std::sqrt(total_var_abs * total_integ_count) / total_mean_abs; + status.count = total_count; + status.count_opt = total_count_opt; + status.count_after_cuts = total_count_after_cuts; + status.count_after_cuts_opt = total_count_after_cuts_opt; + status.iterations = iterations; + status.count_invalid_step_values = count_invalid_step_values; + status.optimized = optimized; +} diff --git a/madspace/src/driver/generator_data.cpp b/madspace/src/driver/generator_data.cpp index 3c1a0da788..b2d3696767 100644 --- a/madspace/src/driver/generator_data.cpp +++ b/madspace/src/driver/generator_data.cpp @@ -76,6 +76,7 @@ void madspace::to_json(nlohmann::json& j, const GeneratorStatus& status) { {"iters_without_improvement", status.iters_without_improvement}, {"optimized", status.optimized}, {"done", status.done}, + {"count_invalid_step_values", status.count_invalid_step_values}, }; } @@ -88,3 +89,27 @@ void madspace::to_json(nlohmann::json& j, const Histogram& hist) { {"bin_errors", hist.bin_errors}, }; } + +StepChainResult madspace::evaluate_step_chain( + double weight0, + double max_weight0, + const double* ratios, + const double* max_weights, + const std::vector& steps +) { + double abs_weight0 = std::abs(weight0); + double acceptance = std::min(1., abs_weight0 / max_weight0); + double overweight = std::max(1., abs_weight0 / max_weight0); + double pending = std::copysign(1., weight0); + for (std::size_t i = 0; auto& step : steps) { + pending *= ratios[i]; + if (step.unweight) { + double unweight_value = overweight * std::abs(pending); + acceptance *= std::min(1., unweight_value / max_weights[i]); + overweight = std::max(1., unweight_value / max_weights[i]); + pending = std::copysign(1., pending); + } + ++i; + } + return {acceptance, overweight * pending}; +} diff --git a/madspace/src/driver/lhe_output.cpp b/madspace/src/driver/lhe_output.cpp index 6a93d5fcbc..a868c86d6c 100644 --- a/madspace/src/driver/lhe_output.cpp +++ b/madspace/src/driver/lhe_output.cpp @@ -3,6 +3,7 @@ #include "madspace/util.hpp" #include +#include #include #include @@ -240,12 +241,47 @@ LHECompleter::append_colors(const SubprocArgs& args, std::size_t particle_count) } void LHECompleter::append_pdg_ids(const SubprocArgs& args, std::size_t particle_count) { - for (auto& pdg_id_options : args.pdg_ids) { + if (args.pdg_id_weights.size() != 0 && + args.pdg_id_weights.size() != args.pdg_ids.size()) { + throw std::invalid_argument("Need PDG-id option weights for every flavor"); + } + for (std::size_t flavor = 0; auto& pdg_id_options : args.pdg_ids) { if (pdg_id_options.size() == 0) { throw std::invalid_argument( "Must provide at least one option per flavor index" ); } + std::span weights; + if (args.pdg_id_weights.size() != 0) { + weights = args.pdg_id_weights.at(flavor); + } + ++flavor; + if (weights.size() != 0) { + if (weights.size() != pdg_id_options.size()) { + throw std::invalid_argument("Need one weight per PDG-id option"); + } + double total = 0.; + for (double weight : weights) { + if (!(weight > 0.) || !std::isfinite(weight)) { + throw std::invalid_argument( + "PDG-id option weights must be positive and finite" + ); + } + total += weight; + } + // equal weights keep the uniform draw + if (std::any_of(weights.begin(), weights.end(), [&](double weight) { + return weight != weights[0]; + })) { + _pdg_id_weight_offsets[_pdg_id_and_count.size()] = + _pdg_id_cumulative.size(); + double sum = 0.; + for (double weight : weights) { + sum += weight; + _pdg_id_cumulative.push_back(sum / total); + } + } + } _pdg_id_and_count.push_back({_pdg_ids.size(), pdg_id_options.size()}); for (auto& pdg_ids : pdg_id_options) { if (pdg_ids.size() != particle_count) { @@ -579,7 +615,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) { @@ -607,9 +646,20 @@ void LHECompleter::complete_event_data( subproc_data.helicity_offset + subproc_data.particle_count * helicity_index; std::size_t mass_offset = subproc_data.mass_offset; - auto [pdg_index, pdg_count] = - _pdg_id_and_count.at(subproc_data.pdg_id_offset + flavor_index); - std::size_t pdg_random = rand_gen.generate_int(pdg_count); + std::size_t flavor_slot = subproc_data.pdg_id_offset + flavor_index; + auto [pdg_index, pdg_count] = _pdg_id_and_count.at(flavor_slot); + std::size_t pdg_random; + if (auto weights = _pdg_id_weight_offsets.find(flavor_slot); + weights != _pdg_id_weight_offsets.end()) { + auto begin = _pdg_id_cumulative.begin() + weights->second; + pdg_random = std::min( + std::upper_bound(begin, begin + pdg_count, rand_gen.generate_double()) - + begin, + pdg_count - 1 + ); + } else { + pdg_random = rand_gen.generate_int(pdg_count); + } std::size_t pdg_offset = pdg_index + subproc_data.particle_count * pdg_random; for (std::size_t particle_index = 0; auto& particle : event.particles) { @@ -640,6 +690,9 @@ void LHECompleter::complete_event_data( std::swap(event.particles.at(0), event.particles.at(1)); } + if (!has_diagrams) { + return; + } auto find_propagators = _propagator_index_and_count.find( cantor_pairing(subprocess_index, diagram_index, color_index) ); @@ -788,6 +841,8 @@ void madspace::to_json(nlohmann::json& j, const LHECompleter& lhe_completer) { {"helicities", lhe_completer._helicities}, {"pdg_id_and_count", lhe_completer._pdg_id_and_count}, {"pdg_ids", lhe_completer._pdg_ids}, + {"pdg_id_weight_offsets", lhe_completer._pdg_id_weight_offsets}, + {"pdg_id_cumulative", lhe_completer._pdg_id_cumulative}, {"propagator_index_and_count", propagator_index_and_count}, {"propagators", lhe_completer._propagators}, {"propagator_colors", lhe_completer._propagator_colors}, @@ -806,6 +861,12 @@ void madspace::from_json(const nlohmann::json& j, LHECompleter& lhe_completer) { lhe_completer._pdg_id_and_count = j.at("pdg_id_and_count").get>>(); lhe_completer._pdg_ids = j.at("pdg_ids").get>(); + // absent from completers saved before weighted options: all uniform + lhe_completer._pdg_id_weight_offsets = j.value( + "pdg_id_weight_offsets", std::unordered_map{} + ); + lhe_completer._pdg_id_cumulative = + j.value("pdg_id_cumulative", std::vector{}); lhe_completer._propagator_index_and_count = {}; for (auto& item : j.at("propagator_index_and_count")) { lhe_completer._propagator_index_and_count[item.at(0).get()] = diff --git a/madspace/src/driver/madnis_training.cpp b/madspace/src/driver/madnis_training.cpp index c577aa149f..614c310a82 100644 --- a/madspace/src/driver/madnis_training.cpp +++ b/madspace/src/driver/madnis_training.cpp @@ -5,6 +5,13 @@ using namespace madspace; +namespace { + +// maximum number of subprocess rows shown in the pretty training box +constexpr std::size_t max_pretty_subproc_rows = 10; + +} // namespace + MadnisTraining::MadnisTraining( ContextPtr generator_context, ContextPtr optimizer_context, @@ -506,7 +513,12 @@ void MadnisTraining::start_single_job( // buffer content depend on thread scheduling bool store_buffer = _batch_index >= _config.buffer_skip_batches; _generator_context->thread_pool().submit( - [this, channel_index, global_channel_index, batch_size, job_id, store_buffer, + [this, + channel_index, + global_channel_index, + batch_size, + job_id, + store_buffer, &job]() { auto& channel = _channels.at(channel_index); if (_seed) { @@ -1022,8 +1034,10 @@ void MultiMadnisTraining::print_progress_init() { } if (_training_args.size() > 1) { - _pretty_box_lower = - PrettyBox("Subprocesses", _training_args.size() + 1, {12, 12, 12, 0}); + std::size_t row_count = + std::min(_training_args.size(), max_pretty_subproc_rows); + _pretty_subproc_rows.assign(_training_args.size(), std::vector(4)); + _pretty_box_lower = PrettyBox("Subprocesses", row_count + 1, {12, 12, 12, 0}); _pretty_box_lower.set_row(0, {"Subprocess", "Loss", "Channels", "Batch"}); _pretty_box_upper = PrettyBox("MadNIS training", 2, {18, 0}, _pretty_box_lower.line_count()); @@ -1147,13 +1161,21 @@ void MultiMadnisTraining::print_progress_update( _pretty_box_upper.set_column( 1, {std::format("{:<15} {}", subproc_str, progress_bar_all), time_str} ); - _pretty_box_lower.set_row( - subproc_index + 1, - {std::format("{}", subproc_index), - std::format("{:>.4f}", loss), - std::format("{}", chan_count), - std::format("{:<15} {}", batch_str, progress_bar)} - ); + _pretty_subproc_rows.at(subproc_index) = { + std::format("{}", subproc_index), + std::format("{:>.4f}", loss), + std::format("{}", chan_count), + std::format("{:<15} {}", batch_str, progress_bar) + }; + // scroll so that the last rows up to the current subprocess are shown + std::size_t row_count = std::min(subproc_count, max_pretty_subproc_rows); + std::size_t first_row = + subproc_index + 1 > row_count ? subproc_index + 1 - row_count : 0; + for (std::size_t row = 0; row < row_count; ++row) { + _pretty_box_lower.set_row( + row + 1, _pretty_subproc_rows.at(first_row + row) + ); + } _pretty_box_upper.print_update(); _pretty_box_lower.print_update(); } diff --git a/madspace/src/driver/weighting_pipeline.cpp b/madspace/src/driver/weighting_pipeline.cpp new file mode 100644 index 0000000000..cd82c5994b --- /dev/null +++ b/madspace/src/driver/weighting_pipeline.cpp @@ -0,0 +1,844 @@ +#include "madspace/driver/weighting_pipeline.hpp" + +#include +#include +#include +#include +#include +#include + +#include "madspace/util.hpp" + +#include + +using namespace madspace; + +namespace { + +std::vector to_vector(const Tensor& tensor) { + auto view = tensor.view(); + std::vector ret(view.size()); + for (std::size_t i = 0; i < ret.size(); ++i) { + ret[i] = view[i]; + } + return ret; +} + +template +void compact_vector(std::vector& vec, const std::vector& indices) { + if (vec.empty()) { + return; + } + std::vector ret; + ret.reserve(indices.size()); + for (auto index : indices) { + ret.push_back(vec.at(index)); + } + vec = std::move(ret); +} + +Function gather_function(const NamedVector& types) { + NamedVector arg_types{{"indices", batch_int}}; + arg_types.insert_back(types); + FunctionBuilder fb(arg_types, types); + for (std::size_t i = 0; i < types.size(); ++i) { + fb.output(i, fb.batch_gather(fb.input(0), fb.input(i + 1))); + } + return fb.function(); +} + +Function concat_function(const NamedVector& types, std::size_t part_count) { + NamedVector arg_types; + for (std::size_t part = 1; part <= part_count; ++part) { + for (auto [key, type] : zip(types.keys(), types)) { + arg_types.push_back(std::format("arg{}_{}", part, key), type); + } + } + FunctionBuilder fb(arg_types, types); + int count = static_cast(types.size()); + for (int i = 0; i < count; ++i) { + ValueVec parts; + for (std::size_t part = 0; part < part_count; ++part) { + parts.push_back(fb.input(i + static_cast(part) * count)); + } + auto [cat, sizes] = fb.batch_cat(parts); + fb.output(i, cat); + } + return fb.function(); +} + +using json = nlohmann::json; + +json matrix_element_to_json(const MatrixElement& matrix_element) { + std::vector me_frame; + for (std::size_t i = 0; auto flag : matrix_element.frame_mask()) { + if (flag) { + me_frame.push_back(static_cast(i + 1)); + } + ++i; + } + // see the MatrixElement constructor: an empty reference mask with a frame + // means that the frame is the one of the incoming particles + auto& reference_mask = matrix_element.reference_mask(); + auto& incoming_mask = + reference_mask.empty() ? matrix_element.frame_mask() : reference_mask; + std::size_t incoming_count = incoming_mask.empty() + ? 2 + : std::count(incoming_mask.begin(), incoming_mask.end(), 1); + json inputs = json::array(), outputs = json::array(); + for (auto input : matrix_element.inputs()) { + inputs.push_back(static_cast(input)); + } + for (auto output : matrix_element.outputs()) { + outputs.push_back(static_cast(output)); + } + return { + {"index", matrix_element.matrix_element_index()}, + {"particle_count", matrix_element.particle_count()}, + {"inputs", inputs}, + {"outputs", outputs}, + {"diagram_count", matrix_element.diagram_count()}, + {"sample_random_inputs", matrix_element.sample_random_inputs()}, + {"me_frame", me_frame}, + {"incoming_count", incoming_count}, + {"channel", matrix_element.channel()}, + }; +} + +MatrixElement matrix_element_from_json(const json& j) { + std::vector inputs; + std::vector outputs; + for (int input : j.at("inputs")) { + inputs.push_back(static_cast(input)); + } + for (int output : j.at("outputs")) { + outputs.push_back(static_cast(output)); + } + return MatrixElement( + j.at("index").get(), + j.at("particle_count").get(), + inputs, + outputs, + j.at("diagram_count").get(), + j.at("sample_random_inputs").get(), + j.at("me_frame").get>(), + j.at("incoming_count").get(), + j.at("channel").get() + ); +} + +} // namespace + +std::string +madspace::weighting_groups_to_json(const std::vector& groups) { + json j = json::array(); + for (auto& group : groups) { + json steps = json::array(); + for (auto& step : group.steps) { + json matrix_elements = json::array(), override_outputs = json::array(); + for (auto& matrix_element : step.matrix_elements) { + matrix_elements.push_back(matrix_element_to_json(matrix_element)); + } + for (auto output : step.override_outputs) { + override_outputs.push_back(static_cast(output)); + } + steps.push_back({ + {"matrix_elements", matrix_elements}, + {"mode", + step.mode == WeightingStep::unweight ? "unweight" : "reweight"}, + {"max_overweight_truncation", step.max_overweight_truncation}, + {"freeze_max_weight_after", step.freeze_max_weight_after}, + {"override_outputs", override_outputs}, + }); + } + json group_json{{"channel_indices", group.channel_indices}, {"steps", steps}}; + if (group.first_pass_max_overweight_truncation) { + group_json["first_pass_max_overweight_truncation"] = + *group.first_pass_max_overweight_truncation; + } + if (group.first_pass_freeze_max_weight_after) { + group_json["first_pass_freeze_max_weight_after"] = + *group.first_pass_freeze_max_weight_after; + } + j.push_back(group_json); + } + return j.dump(); +} + +std::vector +madspace::weighting_groups_from_json(const std::string& text) { + std::vector groups; + for (auto& group_json : json::parse(text)) { + WeightingGroup group{ + .channel_indices = + group_json.at("channel_indices").get>() + }; + if (group_json.contains("first_pass_max_overweight_truncation")) { + group.first_pass_max_overweight_truncation = + group_json.at("first_pass_max_overweight_truncation").get(); + } + if (group_json.contains("first_pass_freeze_max_weight_after")) { + group.first_pass_freeze_max_weight_after = + group_json.at("first_pass_freeze_max_weight_after").get(); + } + for (auto& step_json : group_json.at("steps")) { + WeightingStep step{ + .mode = step_json.at("mode").get() == "unweight" + ? WeightingStep::unweight + : WeightingStep::reweight, + .max_overweight_truncation = + step_json.at("max_overweight_truncation").get(), + .freeze_max_weight_after = + step_json.at("freeze_max_weight_after").get(), + }; + for (auto& me_json : step_json.at("matrix_elements")) { + step.matrix_elements.push_back(matrix_element_from_json(me_json)); + } + for (int output : step_json.at("override_outputs")) { + step.override_outputs.push_back( + static_cast(output) + ); + } + group.steps.push_back(std::move(step)); + } + groups.push_back(std::move(group)); + } + return groups; +} + +void WeightingPipeline::Chunk::clear() { + source.reset(); + keep.reset(); + events.clear(); + cpu_events.clear(); + weights.clear(); + matrix_elements.clear(); + biases.clear(); + overweights.clear(); + ratio_products.clear(); + for (auto& step_ratios : ratios) { + step_ratios.clear(); + } +} + +void WeightingPipeline::Chunk::compact(const std::vector& indices) { + compact_vector(weights, indices); + compact_vector(matrix_elements, indices); + compact_vector(biases, indices); + compact_vector(overweights, indices); + compact_vector(ratio_products, indices); + for (auto& step_ratios : ratios) { + compact_vector(step_ratios, indices); + } +} + +WeightingPipeline::WeightingPipeline( + const std::vector& contexts, + const std::vector>& channels, + const std::vector& groups, + const GeneratorConfig& config, + Sink sink +) : + _contexts(contexts), + _channels(channels), + _config(config), + _channel_groups(channels.size(), -1), + _channel_sequences(channels.size()), + _next_sequences(channels.size()), + _ready(channels.size()), + _context_job_counts(contexts.size()), + _sink(sink) { + if (!_sink) { + _sink = [this](StepEvents& step_events) { + _channels.at(step_events.channel_index)->write_step_events(step_events); + }; + } + const std::array index_outputs{ + MatrixElement::color_index_out, + MatrixElement::helicity_index_out, + MatrixElement::diagram_index_out, + }; + for (std::size_t group_index = 0; auto& group_config : groups) { + if (group_config.steps.empty() || group_config.channel_indices.empty()) { + throw std::invalid_argument( + "a weighting group needs at least one channel and one step" + ); + } + for (auto channel_index : group_config.channel_indices) { + if (channel_index >= channels.size()) { + throw std::invalid_argument( + "channel index of weighting group too large" + ); + } + if (_channel_groups.at(channel_index) != -1) { + throw std::invalid_argument( + "a channel can only belong to one weighting group" + ); + } + _channel_groups.at(channel_index) = static_cast(group_index); + } + + auto& first_channel = channels.at(group_config.channel_indices.at(0)); + NamedVector types = first_channel->event_types(); + for (auto channel_index : group_config.channel_indices) { + auto& channel = channels.at(channel_index); + auto& channel_types = channel->event_types(); + if (channel_types.keys() != types.keys() || + !std::equal( + channel_types.begin(), channel_types.end(), types.begin() + ) || + channel->provided_indices() != first_channel->provided_indices()) { + throw std::invalid_argument( + "the channels of a weighting group need the same events" + ); + } + } + auto& index_map = types.index_map(); + if (!index_map.contains("matrix_element")) { + throw std::invalid_argument( + "a weighting group needs integrands that return the matrix element" + ); + } + std::size_t step_count = group_config.steps.size(); + types.push_back( + "step_random", + batch_float_array(WeightingStepFunction::random_count * step_count) + ); + + Group group{ + .momenta = static_cast(index_map.at("momenta")), + .alpha_s = static_cast(index_map.at("alpha_qcd")), + .flavor = static_cast(index_map.at("matrix_element_flavor")), + .me_index = index_map.contains("matrix_element_index") + ? static_cast(index_map.at("matrix_element_index")) + : -1, + .color = static_cast(index_map.at("color_index")), + .helicity = static_cast(index_map.at("helicity_index")), + .diagram = static_cast(index_map.at("diagram_index")), + .random = static_cast(types.size() - 1), + }; + + std::vector step_functions; + auto provided = first_channel->provided_indices(); + for (std::size_t step_index = 0; auto& step : group_config.steps) { + if (step.matrix_elements.size() > 1 && group.me_index == -1) { + throw std::invalid_argument( + "several matrix elements in a weighting step, but the integrand " + "has only one" + ); + } + group.multiple_matrix_elements.push_back(step.matrix_elements.size() > 1); + group.settings.push_back({ + .unweight = step.mode == WeightingStep::unweight, + .max_overweight_truncation = step.max_overweight_truncation, + .freeze_max_weight_after = step.freeze_max_weight_after, + .ratio = step.returns_ratio(), + }); + std::array write_indices{}; + for (std::size_t i = 0; i < index_outputs.size(); ++i) { + auto output = index_outputs.at(i); + auto& outputs = step.matrix_elements.at(0).outputs(); + bool returns = + std::find(outputs.begin(), outputs.end(), output) != outputs.end(); + bool overrides = + std::find( + step.override_outputs.begin(), + step.override_outputs.end(), + output + ) != step.override_outputs.end(); + write_indices.at(i) = returns && (!provided.at(i) || overrides); + provided.at(i) = provided.at(i) || returns; + } + step_functions.push_back( + WeightingStepFunction( + step, + first_channel->particle_count(), + write_indices, + step_index, + step_count + ) + .function() + ); + ++step_index; + } + + Function gather = gather_function(types); + for (auto& context : contexts) { + ContextRuntimes runtimes{ + .gather = build_runtime(gather, context, false), + }; + for (auto& function : step_functions) { + runtimes.steps.push_back(build_runtime(function, context, false)); + } + group.runtimes.push_back(std::move(runtimes)); + } + group.types = types; + group.buffers.assign( + contexts.size(), std::vector>(step_count + 1) + ); + group.buffer_sizes.assign( + contexts.size(), std::vector(step_count + 1) + ); + for (auto channel_index : group_config.channel_indices) { + channels.at(channel_index)->set_weighting_steps(group.settings); + channels.at(channel_index) + ->set_first_pass_settings( + group_config.first_pass_max_overweight_truncation, + group_config.first_pass_freeze_max_weight_after + ); + _next_sequences.at(channel_index).assign(step_count + 1, 0); + _ready.at(channel_index).resize(step_count + 1); + } + _groups.push_back(std::move(group)); + ++group_index; + } +} + +std::size_t WeightingPipeline::batch_size(std::size_t context_index) const { + return _contexts.at(context_index)->device()->device_type() == DeviceType::cpu + ? _config.cpu_batch_size + : _config.gpu_batch_size; +} + +void WeightingPipeline::add_events(GeneratorBatchJob& job) { + auto& channel = _channels.at(job.channel_index); + auto chunk = std::make_shared(); + chunk->channel_index = job.channel_index; + chunk->job_id = job.job_id; + chunk->sequence = _channel_sequences.at(job.channel_index)++; + chunk->epoch = channel->event_epoch(); + chunk->context_index = job.context_index; + chunk->generated_count = job.weights.size(0); + chunk->count_target = job.count_target; + chunk->source = std::make_shared(std::move(job.unweighted_events)); + chunk->weights = to_vector(job.step_host.at(0)); + chunk->matrix_elements = to_vector(job.step_host.at(1)); + if (job.step_host.size() > 2) { + chunk->biases = to_vector(job.step_host.at(2)); + } + chunk->max_weight = job.max_weight; + for (double weight : chunk->weights) { + chunk->overweights.push_back(std::max(1., std::abs(weight) / job.max_weight)); + chunk->ratio_products.push_back(std::copysign(1., weight)); + } + channel->add_first_step_events(chunk->size(), job.max_weight); + job.unweighted_events.clear(); + job.step_host.clear(); + ++_chunk_count; + advance(chunk, 1); +} + +void WeightingPipeline::advance(ChunkPtr chunk, std::size_t stage) { + if (chunk->size() == 0) { + set_ready(chunk, stage); + return; + } + auto& group = _groups.at(_channel_groups.at(chunk->channel_index)); + group.buffers.at(chunk->context_index).at(stage - 1).push_back(chunk); + group.buffer_sizes.at(chunk->context_index).at(stage - 1) += chunk->size(); +} + +void WeightingPipeline::set_ready(ChunkPtr chunk, std::size_t stage) { + std::size_t channel_index = chunk->channel_index; + auto& ready = _ready.at(channel_index).at(stage - 1); + auto& next_sequence = _next_sequences.at(channel_index).at(stage - 1); + ready.emplace(chunk->sequence, chunk); + while (!ready.empty() && ready.begin()->first == next_sequence) { + ChunkPtr next_chunk = ready.begin()->second; + ready.erase(ready.begin()); + ++next_sequence; + commit_stage(next_chunk, stage); + } +} + +void WeightingPipeline::commit_stage(ChunkPtr chunk, std::size_t stage) { + auto& channel = _channels.at(chunk->channel_index); + std::size_t step_count = channel->weighting_steps().size(); + // the events of the channel were cleared since + bool stale = chunk->epoch != channel->event_epoch(); + if (stale) { + chunk->clear(); + } + if (stage <= step_count) { + if (!stale && chunk->size() > 0) { + decide(*chunk, stage); + } else { + chunk->ratios.emplace_back(); + chunk->max_weights.push_back(0.); + } + advance(chunk, stage + 1); + return; + } + --_chunk_count; + if (!stale) { + StepEvents step_events{ + .channel_index = chunk->channel_index, + .job_id = chunk->job_id, + .events = std::move(chunk->cpu_events), + .weights = std::move(chunk->weights), + .max_weight = chunk->max_weight, + .ratios = std::move(chunk->ratios), + .max_weights = std::move(chunk->max_weights), + .biases = std::move(chunk->biases), + .generated_count = chunk->generated_count, + }; + _sink(step_events); + } +} + +void WeightingPipeline::decide(Chunk& chunk, std::size_t step) { + auto& channel = _channels.at(chunk.channel_index); + auto& settings = channel->weighting_steps().at(step - 1); + std::size_t count = chunk.size(); + std::vector ratios(count); + std::size_t invalid_count = 0; + for (std::size_t i = 0; i < count; ++i) { + double previous = chunk.matrix_elements.at(i); + double& value = chunk.step_matrix_elements.at(i); + double ratio = value; + if (!settings.ratio) { + ratio = previous != 0. ? value / previous + : value != 0. + ? std::numeric_limits::quiet_NaN() + : 0.; + } + // the event gets the weight 0, so that it cannot set the maximum + if (!std::isfinite(ratio)) { + ratio = 0.; + ++invalid_count; + } + if (settings.ratio) { + // the implied matrix element is the reference of the next step + value = ratio * previous; + } + ratios.at(i) = ratio; + chunk.ratio_products.at(i) *= ratio; + } + if (invalid_count > 0) { + channel->add_invalid_step_values(invalid_count); + } + std::vector keep; + if (settings.unweight) { + // the overweight of the last unweighting times the ratios since, see + // evaluate_step_chain() + std::vector values(count); + for (std::size_t i = 0; i < count; ++i) { + values.at(i) = + chunk.overweights.at(i) * std::abs(chunk.ratio_products.at(i)); + } + channel->update_step_max_weight(step, values, chunk.count_target); + double max_weight = channel->step_max_weight(step); + for (std::size_t i = 0; i < count; ++i) { + double value = values.at(i); + if (value > 0. && value > chunk.step_random.at(i) * max_weight) { + keep.push_back(static_cast(i)); + } + chunk.overweights.at(i) = std::max(1., value / max_weight); + chunk.ratio_products.at(i) = std::copysign(1., chunk.ratio_products.at(i)); + } + channel->add_step_accepted( + step, static_cast(keep.size()), chunk.max_weight, chunk.max_weights + ); + chunk.max_weights.push_back(max_weight); + } else { + for (std::size_t i = 0; i < count; ++i) { + if (ratios.at(i) != 0.) { + keep.push_back(static_cast(i)); + } + } + channel->add_step_accepted( + step, static_cast(keep.size()), chunk.max_weight, chunk.max_weights + ); + chunk.max_weights.push_back(0.); + } + chunk.ratios.push_back(std::move(ratios)); + chunk.matrix_elements = std::move(chunk.step_matrix_elements); + chunk.step_matrix_elements.clear(); + chunk.step_random.clear(); + if (keep.size() < count) { + chunk.compact(keep); + chunk.keep = std::move(keep); + } + if (chunk.size() == 0) { + chunk.source.reset(); + chunk.keep.reset(); + } +} + +std::size_t WeightingPipeline::start_jobs( + ResultQueue& result_queue, bool flush, const std::vector& busy_counts +) { + std::size_t started = 0; + for (std::size_t context_index = 0; context_index < _contexts.size(); + ++context_index) { + std::vector> stages; + for (std::size_t group_index = 0; auto& group : _groups) { + for (std::size_t stage = 1; stage <= group.settings.size() + 1; ++stage) { + if (!group.buffers.at(context_index).at(stage - 1).empty()) { + stages.emplace_back(group_index, stage); + } + } + ++group_index; + } + std::size_t free = 0, share = 1; + if (flush && !stages.empty()) { + // with flush, the buffers share the free threads + std::size_t thread_count = + _contexts.at(context_index)->thread_pool().thread_count(); + std::size_t busy = _context_job_counts.at(context_index) + + (busy_counts.empty() ? 0 : busy_counts.at(context_index)); + free = thread_count > busy ? thread_count - busy : 0; + share = std::max(1, free / stages.size()); + } + for (auto [group_index, stage] : stages) { + auto& group = _groups.at(group_index); + auto& buffer = group.buffers.at(context_index).at(stage - 1); + auto& size = group.buffer_sizes.at(context_index).at(stage - 1); + bool full = size >= batch_size(context_index); + if (!full && free == 0) { + continue; + } + std::size_t job_count = + std::clamp(std::min(share, free), 1, buffer.size()); + free -= std::min(free, job_count); + // split at the chunk boundaries into jobs of similar size + std::vector chunks; + std::size_t job_events = 0, done_events = 0; + for (std::size_t job_index = 0; auto& chunk : buffer) { + chunks.push_back(chunk); + job_events += chunk->size(); + if (job_events * (job_count - job_index) >= size - done_events || + &chunk == &buffer.back()) { + done_events += job_events; + job_events = 0; + start_job( + result_queue, + group_index, + context_index, + stage, + std::move(chunks) + ); + chunks.clear(); + ++job_index; + ++started; + } + } + buffer.clear(); + size = 0; + } + } + return started; +} + +void WeightingPipeline::start_job( + ResultQueue& result_queue, + std::size_t group_index, + std::size_t context_index, + std::size_t stage, + std::vector chunks +) { + std::size_t job_id = job_id_flag | _job_counter++; + auto& job = + _jobs + .emplace( + job_id, + Job{ + .group = group_index, + .stage = stage, + .context_index = context_index, + .chunks = std::move(chunks), + } + ) + .first->second; + ++_context_job_counts.at(context_index); + result_queue.submit( + _contexts.at(context_index)->thread_pool(), job_id, [this, &job] { + run_job(job); + } + ); +} + +void WeightingPipeline::gather_events(Chunk& chunk, ContextRuntimes& runtimes) { + auto& source = *chunk.source; + std::size_t count = chunk.size(); + if (!chunk.keep && chunk.offset == 0 && source.at(0).size(0) == count) { + chunk.events = source; + } else { + chunk.events = gather_rows(source, {&chunk}, chunk.context_index, runtimes); + } + chunk.source.reset(); + chunk.keep.reset(); + chunk.offset = 0; +} + +TensorVec WeightingPipeline::gather_rows( + const TensorVec& source, + const std::vector& chunks, + std::size_t context_index, + ContextRuntimes& runtimes +) { + std::size_t count = 0; + for (auto chunk : chunks) { + count += chunk->size(); + } + Tensor indices(DataType::dt_int, {count}); + auto view = indices.view(); + for (std::size_t row = 0; auto chunk : chunks) { + for (std::size_t i = 0; i < chunk->size(); ++i, ++row) { + view[row] = static_cast( + chunk->offset + (chunk->keep ? chunk->keep->at(i) : i) + ); + } + } + auto device = _contexts.at(context_index)->device(); + if (device != cpu_device()) { + Tensor device_indices(DataType::dt_int, {count}, device); + device_indices.copy_from(indices); + indices = device_indices; + } + TensorVec args{indices}; + args.insert(args.end(), source.begin(), source.end()); + return runtimes.gather->run(args); +} + +RuntimePtr& WeightingPipeline::concat_runtime( + std::size_t group_index, std::size_t context_index, std::size_t part_count +) { + std::lock_guard lock(_concat_mutex); + auto& group = _groups.at(group_index); + auto& runtime = group.runtimes.at(context_index).concats[part_count]; + if (!runtime) { + runtime = build_runtime( + concat_function(group.types, part_count), _contexts.at(context_index), false + ); + } + return runtime; +} + +void WeightingPipeline::run_job(Job& job) { + auto& group = _groups.at(job.group); + auto& runtimes = group.runtimes.at(job.context_index); + if (job.stage > group.settings.size()) { + // after the last step, the events go to the CPU to be written + for (auto& chunk : job.chunks) { + gather_events(*chunk, runtimes); + for (auto& tensor : chunk->events) { + chunk->cpu_events.push_back(tensor.cpu()); + } + chunk->events.clear(); + } + return; + } + + // one gather per source, then one concatenation of the sources + std::vector> source_chunks; + for (auto& chunk : job.chunks) { + auto found = + std::find_if(source_chunks.begin(), source_chunks.end(), [&](auto& chunks) { + return chunks.at(0)->source == chunk->source; + }); + if (found == source_chunks.end()) { + source_chunks.push_back({chunk}); + } else { + found->push_back(chunk); + } + } + job.chunks.clear(); + TensorVec events, concat_args; + for (auto& chunks : source_chunks) { + auto& first = *chunks.at(0); + auto& source = *first.source; + TensorVec part; + if (chunks.size() == 1 && !first.keep && first.offset == 0 && + source.at(0).size(0) == first.size()) { + part = source; + } else { + std::vector chunk_ptrs; + for (auto& chunk : chunks) { + chunk_ptrs.push_back(chunk.get()); + } + part = gather_rows(source, chunk_ptrs, job.context_index, runtimes); + } + concat_args.insert(concat_args.end(), part.begin(), part.end()); + for (auto& chunk : chunks) { + chunk->source.reset(); + chunk->keep.reset(); + chunk->offset = 0; + // the chunks in the order of the events + job.chunks.push_back(chunk); + } + } + events = source_chunks.size() == 1 + ? std::move(concat_args) + : concat_runtime(job.group, job.context_index, source_chunks.size()) + ->run(concat_args); + + TensorVec step_args{ + events.at(group.momenta), events.at(group.alpha_s), events.at(group.flavor) + }; + if (group.multiple_matrix_elements.at(job.stage - 1)) { + step_args.push_back(events.at(group.me_index)); + } + step_args.push_back(events.at(group.color)); + step_args.push_back(events.at(group.helicity)); + step_args.push_back(events.at(group.diagram)); + step_args.push_back(events.at(group.random)); + TensorVec outputs = runtimes.steps.at(job.stage - 1)->run(step_args); + auto matrix_elements = to_vector(outputs.at(0).cpu()); + auto random = to_vector(outputs.at(4).cpu()); + events.at(group.color) = outputs.at(1); + events.at(group.helicity) = outputs.at(2); + events.at(group.diagram) = outputs.at(3); + + auto shared_events = std::make_shared(std::move(events)); + std::size_t offset = 0; + for (auto& chunk : job.chunks) { + std::size_t count = chunk->size(); + chunk->source = shared_events; + chunk->offset = offset; + chunk->step_matrix_elements.assign( + matrix_elements.begin() + offset, matrix_elements.begin() + offset + count + ); + chunk->step_random.assign( + random.begin() + offset, random.begin() + offset + count + ); + offset += count; + } +} + +void WeightingPipeline::commit_job(std::size_t job_id) { + auto node = _jobs.extract(job_id); + auto& job = node.mapped(); + --_context_job_counts.at(job.context_index); + for (auto& chunk : job.chunks) { + set_ready(chunk, job.stage); + } +} + +void WeightingPipeline::reset() { + _jobs.clear(); + std::fill(_context_job_counts.begin(), _context_job_counts.end(), 0); + for (auto& group : _groups) { + for (auto& buffers : group.buffers) { + for (auto& buffer : buffers) { + buffer.clear(); + } + } + for (auto& sizes : group.buffer_sizes) { + std::fill(sizes.begin(), sizes.end(), 0); + } + } + for (auto& channel_ready : _ready) { + for (auto& ready : channel_ready) { + ready.clear(); + } + } + // the sequences continue, as the dropped chunks are never committed + for (std::size_t channel_index = 0; channel_index < _channels.size(); + ++channel_index) { + for (auto& next_sequence : _next_sequences.at(channel_index)) { + next_sequence = _channel_sequences.at(channel_index); + } + } + _chunk_count = 0; +} diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index 1b2624b72f..4574c58e61 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -76,6 +76,7 @@ void op_matrix_element( // implemented at LUMI hackathon where the coffee was indeed terrible switch (output_keys[i]) { case UMAMI_OUT_MATRIX_ELEMENT: + case UMAMI_OUT_MATRIX_ELEMENT_RATIO: thrust::fill_n( thrust_par.on(device.stream()), thrust::device_pointer_cast(static_cast(output_ptrs[i])), diff --git a/madspace/src/kernels/multichannel.hpp b/madspace/src/kernels/multichannel.hpp index 1c37f4a3c8..945264c12c 100644 --- a/madspace/src/kernels/multichannel.hpp +++ b/madspace/src/kernels/multichannel.hpp @@ -44,6 +44,35 @@ KERNELSPEC void kernel_collect_channel_weights( } } +template +KERNELSPEC void kernel_density_channel_weights( + FIn densities, FIn mask, FOut channel_weights +) { + // A density that is not a finite positive number (an inverse mapping that + // failed) counts as zero: that channel cannot have produced the event. + FVal norm(0.), mask_count(0.); + for (std::size_t i = 0; i < channel_weights.size(); ++i) { + auto density = densities[i]; + auto masked = mask[i] > 0.; + auto valid = (density > 0.) & (density < 1e300) & masked; + channel_weights[i] = where(valid, density, FVal(0.)); + norm = norm + channel_weights[i]; + mask_count = mask_count + where(masked, FVal(1.), FVal(0.)); + } + // if all inverses failed numerically, uniform over the masked channels, so + // that the weights stay a function of the event alone + auto degenerate = norm == 0.; + auto safe_norm = where(degenerate, FVal(1.), norm); + auto uniform = 1. / where(mask_count > 0., mask_count, FVal(1.)); + for (std::size_t i = 0; i < channel_weights.size(); ++i) { + channel_weights[i] = where( + degenerate, + where(mask[i] > 0., uniform, FVal(0.)), + channel_weights[i] / safe_norm + ); + } +} + template KERNELSPEC void kernel_sde2_channel_weights( FIn invariants, diff --git a/madspace/src/phasespace/channel_weights.cpp b/madspace/src/phasespace/channel_weights.cpp index dff2cdfe1a..e3f309fa71 100644 --- a/madspace/src/phasespace/channel_weights.cpp +++ b/madspace/src/phasespace/channel_weights.cpp @@ -2,6 +2,9 @@ #include "madspace/util.hpp" +#include +#include + using namespace madspace; PropagatorChannelWeights::PropagatorChannelWeights( @@ -211,3 +214,134 @@ NamedVector SubchannelWeights::build_function_impl( ); return {{"channel_weights_out", channel_weights_out}}; } + +MappingChannelWeights::MappingChannelWeights( + const std::vector& mappings, + const nested_vector2& flavor_masks, + const std::vector& diagram_channels, + const nested_vector2& flavor_diagram_masks +) : + FunctionGenerator( + "MappingChannelWeights", + [&] { + if (mappings.empty()) { + throw std::invalid_argument("MappingChannelWeights needs a channel"); + } + NamedVector args{ + {"momenta", batch_four_vec_array(mappings.at(0).particle_count())}, + {"x1", batch_float}, + {"x2", batch_float}, + {"flavor", batch_int}, + }; + if (mappings.at(0).mirror_beams()) { + args.push_back("mirror_index", batch_int); + } + return args; + }(), + [&] { + NamedVector ret{ + {"channel_weights", batch_float_array(mappings.size())} + }; + if (!diagram_channels.empty()) { + ret.push_back("diagram_index", batch_int); + } + return ret; + }() + ), + _mappings(mappings), + _flavor_masks(flavor_masks), + _diagram_channels(diagram_channels), + _flavor_diagram_masks(flavor_diagram_masks), + _mirror_beams(mappings.at(0).mirror_beams()) { + for (auto& mapping : mappings) { + if (mapping.channel_count() > 1) { + throw std::invalid_argument( + "MappingChannelWeights: the mappings cannot have permutations" + ); + } + if (mapping.particle_count() != mappings.at(0).particle_count() || + mapping.mirror_beams() != _mirror_beams) { + throw std::invalid_argument( + "MappingChannelWeights: the mappings must take the same momenta" + ); + } + } + for (auto& mask : flavor_masks) { + if (mask.size() != mappings.size()) { + throw std::invalid_argument( + "MappingChannelWeights: one flavor mask entry per channel expected" + ); + } + } + if (flavor_masks.empty()) { + throw std::invalid_argument("MappingChannelWeights needs a flavor"); + } + for (auto channel : diagram_channels) { + if (channel < 0 || std::cmp_greater_equal(channel, mappings.size())) { + throw std::invalid_argument( + "MappingChannelWeights: diagram channel out of range" + ); + } + } + if (!diagram_channels.empty()) { + if (flavor_diagram_masks.size() != flavor_masks.size()) { + throw std::invalid_argument( + "MappingChannelWeights: one diagram mask per flavor expected" + ); + } + for (auto& mask : flavor_diagram_masks) { + if (mask.size() != diagram_channels.size()) { + throw std::invalid_argument( + "MappingChannelWeights: one diagram mask entry per diagram " + "expected" + ); + } + } + } +} + +NamedVector MappingChannelWeights::build_function_impl( + FunctionBuilder& fb, const NamedVector& args +) const { + ValueVec conditions; + if (_mirror_beams) { + conditions.push_back(args.at("mirror_index")); + } + ValueVec densities; + for (auto& mapping : _mappings) { + auto inverse = mapping.build_inverse( + fb, {args.at("momenta"), args.at("x1"), args.at("x2")}, conditions + ); + // the discrete choices are sampled uniformly, each with probability 1/2 + Value density = fb.mul( + inverse.at("det"), + std::pow(0.5, static_cast(mapping.discrete_dim())) + ); + // the event is out of reach of the mapping if a random number is outside + // the unit hypercube + if (mapping.random_dim() > 0) { + density = fb.mul(density, fb.cut_all(inverse.at("random"), 0., 1.)); + } + densities.push_back(density); + } + Value flavor = args.at("flavor"); + Value channel_weights = fb.density_channel_weights( + fb.stack(densities), fb.gather_vector(flavor, _flavor_masks) + ); + NamedVector ret{{"channel_weights", channel_weights}}; + if (!_diagram_channels.empty()) { + Value probs = fb.mul( + fb.select(channel_weights, _diagram_channels), + fb.gather_vector(flavor, _flavor_diagram_masks) + ); + Value random = + fb + .unstack(fb.random( + fb.batch_size({args.at("momenta")}), static_cast(1) + )) + .at(0); + auto [diagram, det] = fb.sample_discrete_probs(random, probs); + ret.push_back("diagram_index", diagram); + } + return ret; +} diff --git a/madspace/src/phasespace/cross_section.cpp b/madspace/src/phasespace/cross_section.cpp index f5f46b31cd..1d97a319fc 100644 --- a/madspace/src/phasespace/cross_section.cpp +++ b/madspace/src/phasespace/cross_section.cpp @@ -11,7 +11,8 @@ DifferentialCrossSection::DifferentialCrossSection( const std::variant& pdf1, const std::variant& pdf2, bool input_momentum_fraction, - bool decay + bool decay, + bool return_matrix_element ) : FunctionGenerator( "DifferentialCrossSection", @@ -50,7 +51,13 @@ DifferentialCrossSection::DifferentialCrossSection( } return arg_types; }(), - matrix_element.return_types() + [&] { + NamedVector ret_types = matrix_element.return_types(); + if (return_matrix_element) { + ret_types.push_back("raw_matrix_element", batch_float); + } + return ret_types; + }() ), _pid_options(pid_options), _matrix_element(matrix_element), @@ -62,7 +69,16 @@ DifferentialCrossSection::DifferentialCrossSection( _e_cm(cm_energy), _energy_scale(energy_scale), _input_momentum_fraction(input_momentum_fraction), - _decay(decay) { + _decay(decay), + _return_matrix_element(return_matrix_element) { + auto& outputs = matrix_element.outputs(); + if (std::find( + outputs.begin(), outputs.end(), MatrixElement::matrix_element_ratio_out + ) != outputs.end()) { + throw std::invalid_argument( + "a matrix element returning a ratio can only be used in a weighting step" + ); + } if (decay) { if (_has_pdf.at(0) || _has_pdf.at(1)) { throw std::invalid_argument("a decay has no beams, so no PDFs"); @@ -157,14 +173,14 @@ NamedVector DifferentialCrossSection::build_function_impl( throw std::runtime_error("matrix element missing in return values"); } std::size_t me_index = search - _matrix_element.outputs().begin(); + Value raw_matrix_element = me_result.at(me_index); if (_decay) { // Differential decay rate: dGamma = |M|^2 dPhi_n / (2 M), in GeV. The // flux is a compile-time constant here (M is fixed by the decaying // particle), so this is a single multiply and needs no dedicated // instruction. Deliberately no GeV^-2 -> pb conversion: a width is not // a cross section. - me_result.at(me_index) = - fb.mul(me_result.at(me_index), 1. / (2. * _e_cm)); + me_result.at(me_index) = fb.mul(me_result.at(me_index), 1. / (2. * _e_cm)); } else { me_result.at(me_index) = fb.diff_cross_section( x1x2.at(0), @@ -175,5 +191,8 @@ NamedVector DifferentialCrossSection::build_function_impl( _e_cm * _e_cm ); } + if (_return_matrix_element) { + me_result.push_back("raw_matrix_element", raw_matrix_element); + } return me_result; } diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index 5533fc48f1..e99f931879 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -4,6 +4,7 @@ #include #include +#include using namespace madspace; @@ -16,12 +17,15 @@ std::size_t final_channel_count( const std::vector& second_chan_weight_remap, std::size_t second_remapped_chan_count, const std::optional& prop_chan_weights, - const std::optional& subchan_weights + const std::optional& subchan_weights, + const std::optional& mapping_chan_weights ) { if (second_chan_weight_remap.size() > 0) { return second_remapped_chan_count; } else if (subchan_weights) { return subchan_weights->channel_count(); + } else if (mapping_chan_weights && first_chan_weight_remap.size() == 0) { + return mapping_chan_weights->channel_count(); } else if (first_chan_weight_remap.size() > 0 || prop_chan_weights) { // with the denominators sde_strategy, prop_chan_weights already // produces first_remapped_chan_count weights without needing a remap @@ -31,6 +35,16 @@ std::size_t final_channel_count( } } +/// Position of `output` in the return values of a differential cross section +/// built on `matrix_element`, or -1 if it is not returned. +int me_output_index( + const MatrixElement& matrix_element, MatrixElement::MatrixElementOutput output +) { + auto& outputs = matrix_element.outputs(); + auto search = std::find(outputs.begin(), outputs.end(), output); + return search == outputs.end() ? -1 : static_cast(search - outputs.begin()); +} + // Whether the event carries the product of its beam densities for the // systematics to divide the PDF variations by. Not under MLM pdf reweighting: // there the densities are taken low on the clustering ladder (pdf_scale) and @@ -48,8 +62,8 @@ bool emits_alphas_scales( const std::optional& running_coupling, std::size_t particle_count ) { - return energy_scale && energy_scale->mlm_alphas_reweighting() && - running_coupling && particle_count > 3; + return energy_scale && energy_scale->mlm_alphas_reweighting() && running_coupling && + particle_count > 3; } // Whether the event carries the MLM pdf reweighting ladder, for the @@ -111,6 +125,7 @@ Integrand::Integrand( std::size_t compressed_channel_weight_count, const std::optional& custom_bias, const std::optional& pdf_grid2, + const std::optional& mapping_chan_weights, int mlm_history_qcd_power ) : FunctionGenerator( @@ -139,7 +154,8 @@ Integrand::Integrand( second_chan_weight_remap, second_remapped_chan_count, prop_chan_weights, - subchan_weights + subchan_weights, + mapping_chan_weights ); if (madnis_channel_count > 1 && madnis_channel_count * 8 > compressed_channel_weight_count * 12) { @@ -212,8 +228,7 @@ Integrand::Integrand( energy_scale, running_coupling, mapping.particle_count() )) { ret_types.push_back( - "alphas_scales", - batch_float_array(mapping.particle_count() - 3) + "alphas_scales", batch_float_array(mapping.particle_count() - 3) ); } if (emits_pdf_ladder( @@ -230,6 +245,13 @@ Integrand::Integrand( ret_types.push_back(name, batch_float_array(2)); } } + if (diff_xs_first.returns_matrix_element()) { + ret_types.push_back("matrix_element", batch_float); + ret_types.push_back("matrix_element_flavor", batch_int); + if (diff_xs.size() > 1) { + ret_types.push_back("matrix_element_index", batch_int); + } + } // before "random": the outputs after it feed the discrete histograms if (custom_bias) { ret_types.push_back("bias", batch_float); @@ -256,6 +278,7 @@ Integrand::Integrand( _running_coupling(running_coupling), _energy_scale(energy_scale), _prop_chan_weights(prop_chan_weights), + _mapping_chan_weights(mapping_chan_weights), _subchan_weights(subchan_weights), _chan_weight_net(chan_weight_net), _first_chan_weight_remap(first_chan_weight_remap), @@ -374,19 +397,15 @@ Integrand::Integrand( for (std::size_t f = 0; f < option_count; ++f) { _pdf_rw_table.at(0 * option_count + f) = index_of(21); for (std::size_t beam = 0; beam < 2; ++beam) { - _pdf_rw_table.at((1 + beam) * option_count + f) = index_of( - pid_options.size() > 0 ? pid_options.at(f).at(beam) : 21 - ); + _pdf_rw_table.at((1 + beam) * option_count + f) = + index_of(pid_options.size() > 0 ? pid_options.at(f).at(beam) : 21); } for (std::size_t c = 0; c < absolute.size(); ++c) { - _pdf_rw_table.at((3 + c) * option_count + f) = - index_of(absolute.at(c)); + _pdf_rw_table.at((3 + c) * option_count + f) = index_of(absolute.at(c)); } } for (std::size_t c = 0; c < class_count; ++c) { - _pdf_rw_class_offsets.push_back( - static_cast(c * option_count) - ); + _pdf_rw_class_offsets.push_back(static_cast(c * option_count)); } _pdf_rw = PartonDensity(*grid, rw_pids, true); _pdf_rw_pdgs.assign(rw_pids.begin(), rw_pids.end()); @@ -411,6 +430,14 @@ Integrand::Integrand( "an MLM clustering history per diagram needs a single matrix element" ); } + if (me_output_index( + diff_xs.at(0).matrix_element(), MatrixElement::diagram_amp2_out + ) == -1) { + throw std::invalid_argument( + "an MLM clustering history per diagram needs the diagram_amp2 " + "output of the matrix element" + ); + } auto& starts = energy_scale->mlm_diagram_start_states(); std::size_t diagram_count = diff_xs.at(0).matrix_element().diagram_count(); std::vector mask(diagram_count, 0.); @@ -448,6 +475,42 @@ Integrand::Integrand( _flavor_mirror_factors.push_back(mirror ? 2. : 1.); _has_mirror |= mirror; } + auto& me_first = diff_xs.at(0).matrix_element(); + for (auto& dxs : diff_xs) { + if (dxs.matrix_element().outputs() != me_first.outputs() || + dxs.matrix_element().external_inputs() != me_first.external_inputs() || + dxs.returns_matrix_element() != diff_xs.at(0).returns_matrix_element()) { + throw std::invalid_argument( + "all matrix elements of an integrand need the same inputs and outputs" + ); + } + } + for (auto input : me_first.external_inputs()) { + if (input != MatrixElement::momenta_in && input != MatrixElement::flavor_in && + input != MatrixElement::alpha_s_in) { + throw std::invalid_argument( + "the integrand only supplies momenta, flavor and alpha_s to the " + "matrix element, the random inputs have to be sampled internally" + ); + } + } + std::size_t channel_count = final_channel_count( + diff_xs, + first_chan_weight_remap, + first_remapped_chan_count, + second_chan_weight_remap, + second_remapped_chan_count, + prop_chan_weights, + subchan_weights, + mapping_chan_weights + ); + if (channel_count > 1 && !prop_chan_weights && !mapping_chan_weights && + me_output_index(me_first, MatrixElement::diagram_amp2_out) == -1) { + throw std::invalid_argument( + "diagram-based channel weights need the diagram_amp2 output of the " + "matrix element" + ); + } // The orientation is drawn before the mapping, so that the boost, the cuts // and the PDFs all see the orientation the event is written with. if (_has_mirror && !mapping.mirror_beams()) { @@ -455,6 +518,19 @@ Integrand::Integrand( "mirrored flavors need a PhaseSpaceMapping built with mirror_beams = true" ); } + if (mapping_chan_weights) { + if (prop_chan_weights) { + throw std::invalid_argument( + "propagator and mapping channel weights cannot be combined" + ); + } + if (mapping_chan_weights->mirror_beams() != mapping.mirror_beams()) { + throw std::invalid_argument( + "the channel weight mappings must have the beam orientation of the " + "integrand's mapping" + ); + } + } // evaluate_pdfs, which takes the densities again once the history is // picked per diagram, knows nothing of the event's orientation, so it // cannot pick a mirrored event's density from the other beam. @@ -527,6 +603,9 @@ NamedVector Integrand::compute_channel_part_ret_types() const { ret.push_back("x1_acc", acc_float); ret.push_back("x2_acc", acc_float); ret.push_back("flavor_id", acc_int); + if (_mapping_chan_weights && _mapping.mirror_beams()) { + ret.push_back("mirror_index_acc", acc_int); + } ret.push_back("weight_after_cuts", acc_float); if (_madnis_training && !std::holds_alternative(_discrete_flavor)) { ret.push_back("extra_weight_after_cuts", acc_float); @@ -553,7 +632,9 @@ NamedVector Integrand::compute_channel_part_ret_types() const { } if (_pdf_rw && _pid_options.size() > 0) { for (auto name : ladder_slot_names) { - ret.push_back(std::format("{}_acc", name), acc_float_array(particle_count - 2)); + ret.push_back( + std::format("{}_acc", name), acc_float_array(particle_count - 2) + ); } for (auto name : ladder_beam_names) { ret.push_back(std::format("{}_acc", name), acc_float_array(2)); @@ -683,9 +764,22 @@ NamedVector Integrand::build_channel_part( mapping_conditions.push_back("mirror_index", mirror_index); } - // Apply phase space mapping + // 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.input_types().keys(), {latent}}, mapping_conditions + fb, {_mapping.input_types().keys(), mapping_inputs}, mapping_conditions ); weights_before_cuts.push_back(mapping_result["det"]); Value momenta = mapping_result["momenta"]; @@ -727,10 +821,11 @@ NamedVector Integrand::build_channel_part( _pdfs.at(i).value().build_function(fb, {x_acc.at(i), pdf_scale}).at(0); if (_has_mirror && _pdfs_swapped.at(i)) { // leg i of a mirrored event comes from the other beam - auto pdf_swapped = _pdfs_swapped.at(i) - .value() - .build_function(fb, {x_acc.at(i), pdf_scale}) - .at(0); + auto pdf_swapped = + _pdfs_swapped.at(i) + .value() + .build_function(fb, {x_acc.at(i), pdf_scale}) + .at(0); pdf = fb.gather_vector( fb.batch_gather(indices_acc, mirror_index), fb.stack({pdf, pdf_swapped}) ); @@ -923,6 +1018,9 @@ NamedVector Integrand::build_channel_part( out.push_back("x1_acc", x_acc.at(0)); out.push_back("x2_acc", x_acc.at(1)); out.push_back("flavor_id", flavor_id); + if (_mapping_chan_weights && _mapping.mirror_beams()) { + out.push_back("mirror_index_acc", fb.batch_gather(indices_acc, mirror_index)); + } out.push_back("weight_after_cuts", weight_after_cuts); if (_madnis_training && extra_weight_after_cuts) { out.push_back("extra_weight_after_cuts", extra_weight_after_cuts); @@ -956,8 +1054,9 @@ NamedVector Integrand::build_channel_part( for (auto name : ladder_slot_names) { out.push_back( std::format("{}_acc", name), - mlm_ladder.size() > 0 ? mlm_ladder.at(name) - : fb.full({0., batch_size, slot_count}) + mlm_ladder.size() > 0 + ? mlm_ladder.at(name) + : fb.full({0., batch_size, slot_count}) ); } for (auto name : ladder_beam_names) { @@ -1034,9 +1133,8 @@ ValueVec Integrand::mlm_weights( auto flavor_classes = scales.at("pdf_rw_flavor"); std::size_t slot_count = flavor_classes.type.shape.at(0); double low = _energy_scale->min_scale(); - double high = _energy_scale->max_scale() > 0. - ? _energy_scale->max_scale() - : 1e30; + double high = + _energy_scale->max_scale() > 0. ? _energy_scale->max_scale() : 1e30; auto ones = fb.full({1., batch_size_acc}); Value factor, veto; ValueVec ladder_x, ladder_pdg, ladder_q, ladder_info; @@ -1052,10 +1150,7 @@ ValueVec Integrand::mlm_weights( auto beam = pick("pdf_rw_beam"); // x of the beam line: the beam's own momentum fraction times // everything the clustering z has taken off it since. - auto x_beam = fb.add( - x.at(0), - fb.mul(beam, fb.sub(x.at(1), x.at(0))) - ); + auto x_beam = fb.add(x.at(0), fb.mul(beam, fb.sub(x.at(1), x.at(0)))); auto x = fb.mul(x_beam, pick("pdf_rw_x")); auto scale_of = [&](const char* name) { auto q = pick(name); @@ -1123,7 +1218,8 @@ ValueVec Integrand::mlm_weights( // each beam's leading density, at the scale it was evaluated at // (clamped by apply_scale_range) ladder->push_back( - "ladder_lead_q", fb.stack({scales.at("pdf_scale1"), scales.at("pdf_scale2")}) + "ladder_lead_q", + fb.stack({scales.at("pdf_scale1"), scales.at("pdf_scale2")}) ); } } @@ -1205,12 +1301,32 @@ NamedVector Integrand::build_common_part( _second_chan_weight_remap, _second_remapped_chan_count, _prop_chan_weights, - _subchan_weights + _subchan_weights, + _mapping_chan_weights ); bool use_compressed_channel_weights = channel_count > 1 && channel_count * 8 > _compressed_channel_weight_count * 12; Value chan_weights_acc; - if (channel_count > 1 && _prop_chan_weights) { + Value mapping_diagram_acc; + if (_mapping_chan_weights && + (channel_count > 1 || _mapping_chan_weights->has_diagrams())) { + ValueVec cw_args{momenta_acc, x1_acc, x2_acc, flavor_id}; + if (_mapping_chan_weights->mirror_beams()) { + cw_args.push_back(args.at("mirror_index_acc")); + } + auto cw_result = _mapping_chan_weights->build_function(fb, cw_args); + chan_weights_acc = cw_result.at(0); + if (_mapping_chan_weights->has_diagrams()) { + mapping_diagram_acc = cw_result.at(1); + } + if (_first_chan_weight_remap.size() > 0) { + chan_weights_acc = fb.collect_channel_weights( + chan_weights_acc, + _first_chan_weight_remap.at(0), + _first_remapped_chan_count + ); + } + } else if (channel_count > 1 && _prop_chan_weights) { chan_weights_acc = _prop_chan_weights->build_function(fb, {momenta_acc}).at(0); if (_first_chan_weight_remap.size() > 0) { chan_weights_acc = fb.collect_channel_weights( @@ -1239,8 +1355,9 @@ NamedVector Integrand::build_common_part( if (_energy_scale && _energy_scale->is_mlm()) { cluster_scales_acc = args.at("cluster_scales_acc"); } - bool has_alphas_scales = - emits_alphas_scales(_energy_scale, _running_coupling, _mapping.particle_count()); + bool has_alphas_scales = emits_alphas_scales( + _energy_scale, _running_coupling, _mapping.particle_count() + ); Value alphas_scales_acc; if (has_alphas_scales) { alphas_scales_acc = args.at("alphas_scales_acc"); @@ -1257,11 +1374,19 @@ NamedVector Integrand::build_common_part( } // Evaluate differential cross section + auto& me_first = _diff_xs.at(0).matrix_element(); + int amp2_index = me_output_index(me_first, MatrixElement::diagram_amp2_out); + Value me_flavor_id = + _flavor_remap.size() > 0 ? flavor_table_lookup(_flavor_remap) : flavor_id; auto make_xs_args = [&](std::array& pdfs, Value alpha) { - ValueVec xs_args{ - momenta_acc, - _flavor_remap.size() > 0 ? flavor_table_lookup(_flavor_remap) : flavor_id, - }; + ValueVec xs_args; + for (auto input : me_first.external_inputs()) { + if (input == MatrixElement::momenta_in) { + xs_args.push_back(momenta_acc); + } else if (input == MatrixElement::flavor_in) { + xs_args.push_back(me_flavor_id); + } + } xs_args.push_back(x1_acc); xs_args.push_back(x2_acc); xs_args.push_back(flavor_id); @@ -1277,11 +1402,12 @@ NamedVector Integrand::build_common_part( ValueVec dxs_vec; Value ps_flavor_id; Value subproc_id; + Value dxs_index; if (_diff_xs.size() == 1) { dxs_vec = _diff_xs.at(0).build_function(fb, xs_args).values(); ps_flavor_id = flavor_id; - if (channel_count > 1 && !_prop_chan_weights) { - chan_weights_acc = dxs_vec.at(1); + if (channel_count > 1 && !_prop_chan_weights && !_mapping_chan_weights) { + chan_weights_acc = dxs_vec.at(amp2_index); if (_first_chan_weight_remap.size() > 0) { chan_weights_acc = fb.collect_channel_weights( chan_weights_acc, @@ -1291,7 +1417,7 @@ NamedVector Integrand::build_common_part( } } } else { - Value dxs_index = flavor_table_lookup(_flavor_diff_xs_indices); + dxs_index = flavor_table_lookup(_flavor_diff_xs_indices); ps_flavor_id = flavor_table_lookup(_flavor_per_subproc_remap); subproc_id = flavor_table_lookup(_flavor_subproc_indices); ValueVec split_indices = @@ -1309,8 +1435,8 @@ NamedVector Integrand::build_common_part( split_out.push_back(out); split_out.push_back(indices); } - if (channel_count > 1 && !_prop_chan_weights) { - Value split_cw = outputs.at(1); + if (channel_count > 1 && !_prop_chan_weights && !_mapping_chan_weights) { + Value split_cw = outputs.at(amp2_index); if (_first_chan_weight_remap.size() > 0) { split_cw = fb.collect_channel_weights( split_cw, @@ -1324,14 +1450,14 @@ NamedVector Integrand::build_common_part( ++i; } for (std::size_t i = 0; auto& split_out : split_outputs) { - if (i == 1) { + if (std::cmp_equal(i, amp2_index)) { dxs_vec.push_back(Value()); } else { dxs_vec.push_back(fb.batch_merge_by_index(split_out)); } ++i; } - if (channel_count > 1 && !_prop_chan_weights) { + if (channel_count > 1 && !_prop_chan_weights && !_mapping_chan_weights) { chan_weights_acc = fb.batch_merge_by_index(split_channel_weights); } } @@ -1354,7 +1480,7 @@ NamedVector Integrand::build_common_part( ); auto [diagram, diagram_det] = fb.sample_discrete_probs( fb.squeeze(fb.random(batch_size_acc, 1)), - fb.mul(dxs_vec.at(1), diagram_mask) + fb.mul(dxs_vec.at(amp2_index), diagram_mask) ); auto scales = _energy_scale->build_mlm_from_start_state( fb, @@ -1405,24 +1531,29 @@ NamedVector Integrand::build_common_part( for (int k = 0; k < _mlm_history_qcd_power; ++k) { factor = fb.mul(factor, alpha_ratio); } + Value alpha_factor = factor; for (std::size_t i = 0; i < 2; ++i) { if (_diff_xs.at(0).has_pdf(i)) { // a first density of exactly zero gave a zero weight, which // stays zero (the guard only keeps 0 * x / 0 from a NaN) - auto denominator = fb.add( - pdfs_first.at(i), fb.cut_one(pdfs_first.at(i), 0., 0.) - ); - factor = - fb.mul(factor, fb.div(pdfs_acc.at(i), denominator)); + auto denominator = + fb.add(pdfs_first.at(i), fb.cut_one(pdfs_first.at(i), 0., 0.)); + factor = fb.mul(factor, fb.div(pdfs_acc.at(i), denominator)); } } dxs_vec.at(0) = fb.mul(dxs_vec.at(0), factor); + if (_diff_xs.at(0).returns_matrix_element()) { + // the raw |M|^2 has no densities, only the alpha_s rescaling + auto& raw_me = dxs_vec.at(me_first.outputs().size()); + raw_me = fb.mul(raw_me, alpha_factor); + } } else { // several powers of alpha_s mix (gluon fusion and VBF in one // process, say): the matrix element has to be evaluated again - dxs_vec = _diff_xs.at(0) - .build_function(fb, make_xs_args(pdfs_acc, alpha_qcd_acc)) - .values(); + dxs_vec = + _diff_xs.at(0) + .build_function(fb, make_xs_args(pdfs_acc, alpha_qcd_acc)) + .values(); } } @@ -1595,9 +1726,27 @@ NamedVector Integrand::build_common_part( outputs.push_back("momenta", args.at("momenta")); auto zeros_int = fb.full({static_cast(0), batch_size_val}); auto zeros_float = fb.full({0., batch_size_val}); - outputs.push_back("color_index", scatter_or_drop(zeros_int, dxs_vec.at(2))); - outputs.push_back("helicity_index", scatter_or_drop(zeros_int, dxs_vec.at(3))); - outputs.push_back("diagram_index", scatter_or_drop(zeros_int, dxs_vec.at(4))); + // indices the matrix element does not provide are -1, a later + // reweighting step can fill them in + for (auto [name, me_output] : + {std::pair{"color_index", MatrixElement::color_index_out}, + std::pair{"helicity_index", MatrixElement::helicity_index_out}, + std::pair{"diagram_index", MatrixElement::diagram_index_out}}) { + int index = me_output_index(me_first, me_output); + if (index == -1 && me_output == MatrixElement::diagram_index_out && + mapping_diagram_acc) { + outputs.push_back( + name, scatter_or_drop(zeros_int, mapping_diagram_acc) + ); + continue; + } + outputs.push_back( + name, + index == -1 + ? optional_cut(fb.full({static_cast(-1), batch_size_val})) + : scatter_or_drop(zeros_int, dxs_vec.at(index)) + ); + } outputs.push_back("flavor_index", scatter_or_drop(zeros_int, ps_flavor_id)); if (subproc_id) { outputs.push_back( @@ -1605,9 +1754,7 @@ NamedVector Integrand::build_common_part( ); } - outputs.push_back( - "ren_scale", scatter_or_drop(zeros_float, ren_scale_acc) - ); + outputs.push_back("ren_scale", scatter_or_drop(zeros_float, ren_scale_acc)); outputs.push_back("alpha_qcd", scatter_or_drop(zeros_float, alpha_qcd_acc)); ValueVec pdf_vals; if (_diff_xs.at(0).has_pdf(0)) { @@ -1636,8 +1783,7 @@ NamedVector Integrand::build_common_part( outputs.push_back( "cluster_scales", scatter_or_drop( - fb.full({0., batch_size_val, outgoing_count}), - cluster_scales_acc + fb.full({0., batch_size_val, outgoing_count}), cluster_scales_acc ) ); } @@ -1646,8 +1792,21 @@ NamedVector Integrand::build_common_part( auto& value = ladder_acc.at(name); auto width = static_cast(value.type.shape.at(0)); outputs.push_back( - name, - scatter_or_drop(fb.full({0., batch_size_val, width}), value) + name, scatter_or_drop(fb.full({0., batch_size_val, width}), value) + ); + } + } + if (_diff_xs.at(0).returns_matrix_element()) { + outputs.push_back( + "matrix_element", + scatter_or_drop(zeros_float, dxs_vec.at(me_first.outputs().size())) + ); + outputs.push_back( + "matrix_element_flavor", scatter_or_drop(zeros_int, me_flavor_id) + ); + if (dxs_index) { + outputs.push_back( + "matrix_element_index", scatter_or_drop(zeros_int, dxs_index) ); } } diff --git a/madspace/src/phasespace/matrix_element.cpp b/madspace/src/phasespace/matrix_element.cpp index 174aa52fd3..d59f56d63e 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); @@ -149,6 +152,9 @@ MatrixElement::MatrixElement( case bias_out: ret_types.push_back("bias", batch_float); break; + case matrix_element_ratio_out: + ret_types.push_back("matrix_element_ratio", batch_float); + break; default: throw std::invalid_argument("unknown output type"); } @@ -162,6 +168,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") @@ -189,7 +196,17 @@ MatrixElement::MatrixElement( mask.clear(); } return mask; - }()) {} + }()) { + auto has_output = [&](MatrixElementOutput output) { + return std::find(outputs.begin(), outputs.end(), output) != outputs.end(); + }; + if (has_output(matrix_element_out) && has_output(matrix_element_ratio_out)) { + throw std::invalid_argument( + "a matrix element returns either matrix_element or matrix_element_ratio, " + "not both" + ); + } +} NamedVector MatrixElement::build_function_impl( FunctionBuilder& fb, const NamedVector& args @@ -267,6 +284,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: @@ -332,6 +351,9 @@ NamedVector MatrixElement::build_function_impl( case bias_out: output_key = UMAMI_OUT_BIAS; break; + case matrix_element_ratio_out: + output_key = UMAMI_OUT_MATRIX_ELEMENT_RATIO; + break; } matrix_args.push_back(static_cast(output_key)); matrix_args.push_back(size_arg); @@ -342,6 +364,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/phasespace/three_particle.cpp b/madspace/src/phasespace/three_particle.cpp index 98b0f4ac9c..a3871414bc 100644 --- a/madspace/src/phasespace/three_particle.cpp +++ b/madspace/src/phasespace/three_particle.cpp @@ -178,6 +178,7 @@ std::array TwoToThreeParticleScattering::sampling_ranges( FunctionBuilder& fb, Value s23_min, Value s23_max, + Value m2_squared, const NamedVector& conditions ) const { // returns the sampled and the kinematic range of the sampled variable @@ -188,8 +189,11 @@ std::array TwoToThreeParticleScattering::sampling_ranges( auto v_phys_min = negated(fb, s23_max), v_phys_max = negated(fb, s23_min); auto v_min = v_phys_min, v_max = v_phys_max; if (_has_cut) { - // |(p2 - k)^2| >= m_T(p2)^2 for a massless k along the beam axis - v_min = fb.max(v_min, fb.square(conditions.at(4))); + // |(p2 - k)^2| >= m_T(p2)^2 - m2^2 = p_T(p2)^2 for a massless k along + // the beam axis: the ETmin floor of a massive particle bounds m_T + v_min = fb.max( + v_min, fb.max(fb.sub(fb.square(conditions.at(4)), m2_squared), Value(0.)) + ); // an empty range keeps a small width; the weight is zero there, as // the point then lies outside the kinematic range auto width = fb.max(fb.mul(v_min, Value(1e-10)), Value(1e-300)); @@ -217,7 +221,7 @@ Mapping::Result TwoToThreeParticleScattering::build_forward_impl( std::array ranges; if (_incoming_recoil) { auto [lo, hi] = fb.s23_min_max(p_a, p_12, p_3, t1, m1, m2); - ranges = sampling_ranges(fb, lo, hi, conditions); + ranges = sampling_ranges(fb, lo, hi, fb.square(m2), conditions); } else if (_has_cut) { auto [s_min, s_max] = fb.s23_min_max_cut( p_a, @@ -303,7 +307,10 @@ Mapping::Result TwoToThreeParticleScattering::build_inverse_impl( if (_incoming_recoil) { auto [s23, lo, hi] = fb.s23_value_and_min_max(p_a, p_3, t1_abs, p1, p2); v = negated(fb, s23); - ranges = sampling_ranges(fb, lo, hi, conditions); + nested_vector2 factors_p2{{1.0}}; + auto m2_squared = + fb.unstack(fb.invariants_from_momenta(fb.stack({p2}), factors_p2)).at(0); + ranges = sampling_ranges(fb, lo, hi, m2_squared, conditions); } else if (_has_cut) { auto [s23, s_min, s_max] = fb.s23_value_and_min_max_cut( p_a, diff --git a/madspace/src/phasespace/weighting_step.cpp b/madspace/src/phasespace/weighting_step.cpp new file mode 100644 index 0000000000..8982b1bf24 --- /dev/null +++ b/madspace/src/phasespace/weighting_step.cpp @@ -0,0 +1,239 @@ +#include "madspace/phasespace/weighting_step.hpp" + +#include "madspace/util.hpp" + +using namespace madspace; + +namespace { + +const std::array index_outputs{ + MatrixElement::color_index_out, + MatrixElement::helicity_index_out, + MatrixElement::diagram_index_out, +}; + +int output_position( + const MatrixElement& matrix_element, MatrixElement::MatrixElementOutput output +) { + auto& outputs = matrix_element.outputs(); + auto search = std::find(outputs.begin(), outputs.end(), output); + return search == outputs.end() ? -1 : static_cast(search - outputs.begin()); +} + +} // namespace + +bool WeightingStep::returns_ratio() const { + return !matrix_elements.empty() && + output_position( + matrix_elements.at(0), MatrixElement::matrix_element_ratio_out + ) != -1; +} + +WeightingStepFunction::WeightingStepFunction( + const WeightingStep& step, + std::size_t particle_count, + const std::array& write_indices, + std::size_t step_index, + std::size_t step_count +) : + FunctionGenerator( + "WeightingStepFunction", + [&] { + NamedVector arg_types{ + {"momenta", batch_four_vec_array(particle_count)}, + {"alpha_s", batch_float}, + {"flavor", batch_int}, + }; + if (step.matrix_elements.size() > 1) { + arg_types.push_back("matrix_element_index", batch_int); + } + arg_types.push_back("color_index", batch_int); + arg_types.push_back("helicity_index", batch_int); + arg_types.push_back("diagram_index", batch_int); + arg_types.push_back("random", batch_float_array(random_count * step_count)); + return arg_types; + }(), + { + {"matrix_element", batch_float}, + {"color_index", batch_int}, + {"helicity_index", batch_int}, + {"diagram_index", batch_int}, + {"unweight_random", batch_float}, + } + ), + _matrix_elements(step.matrix_elements), + _write_indices(write_indices), + _step_index(step_index) { + if (step_index >= step_count) { + throw std::invalid_argument("step index out of range"); + } + if (_matrix_elements.empty()) { + throw std::invalid_argument("a weighting step needs a matrix element"); + } + for (auto& matrix_element : _matrix_elements) { + if (matrix_element.outputs() != _matrix_elements.at(0).outputs()) { + throw std::invalid_argument( + "the matrix elements of a weighting step need the same outputs" + ); + } + if (matrix_element.particle_count() != particle_count) { + throw std::invalid_argument( + "the matrix elements of a weighting step need the particle count " + "of the events" + ); + } + if (output_position(matrix_element, MatrixElement::matrix_element_out) == -1 && + output_position(matrix_element, MatrixElement::matrix_element_ratio_out) == + -1) { + throw std::invalid_argument( + "the matrix elements of a weighting step must return the matrix " + "element or its ratio" + ); + } + for (auto input : matrix_element.external_inputs()) { + switch (input) { + case MatrixElement::momenta_in: + case MatrixElement::alpha_s_in: + case MatrixElement::flavor_in: + case MatrixElement::random_color_in: + case MatrixElement::random_helicity_in: + case MatrixElement::random_diagram_in: + case MatrixElement::helicity_in: + case MatrixElement::diagram_in: + break; + default: + throw std::invalid_argument( + "unsupported input of a weighting step matrix element" + ); + } + } + 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" + ); + } + } +} + +NamedVector WeightingStepFunction::build_function_impl( + FunctionBuilder& fb, const NamedVector& args +) const { + ValueVec all_random = fb.unstack(args.at("random")); + ValueVec random( + all_random.begin() + random_count * _step_index, + all_random.begin() + random_count * (_step_index + 1) + ); + ValueVec indices{ + args.at("color_index"), args.at("helicity_index"), args.at("diagram_index") + }; + auto me_args = [&](const MatrixElement& matrix_element, auto&& arg) { + ValueVec ret; + for (auto input : matrix_element.inputs()) { + switch (input) { + case MatrixElement::momenta_in: + ret.push_back(arg(args.at("momenta"))); + break; + case MatrixElement::alpha_s_in: + ret.push_back(arg(args.at("alpha_s"))); + break; + case MatrixElement::flavor_in: + ret.push_back(arg(args.at("flavor"))); + break; + case MatrixElement::random_color_in: + ret.push_back(arg(random.at(0))); + break; + case MatrixElement::random_helicity_in: + ret.push_back(arg(random.at(1))); + break; + case MatrixElement::random_diagram_in: + ret.push_back(arg(random.at(2))); + break; + case MatrixElement::helicity_in: + ret.push_back(arg(indices.at(1))); + break; + case MatrixElement::diagram_in: + ret.push_back(arg(indices.at(2))); + break; + default: + break; + } + } + return ret; + }; + // matrix element plus the index outputs, missing ones are empty + auto select_outputs = + [&](const MatrixElement& matrix_element, const ValueVec& outputs) { + int value_position = + output_position(matrix_element, MatrixElement::matrix_element_out); + if (value_position == -1) { + value_position = output_position( + matrix_element, MatrixElement::matrix_element_ratio_out + ); + } + ValueVec ret{outputs.at(value_position)}; + for (auto output : index_outputs) { + int position = output_position(matrix_element, output); + ret.push_back(position == -1 ? Value() : outputs.at(position)); + } + return ret; + }; + + ValueVec results; + if (_matrix_elements.size() == 1) { + auto& matrix_element = _matrix_elements.at(0); + results = select_outputs( + matrix_element, + matrix_element + .build_function(fb, me_args(matrix_element, [](Value v) { return v; })) + .values() + ); + } else { + ValueVec split_indices = fb.batch_split_by_index( + args.at("matrix_element_index"), + static_cast(_matrix_elements.size()) + ); + std::vector split_results(1 + index_outputs.size()); + for (auto [matrix_element, split_index] : + zip(_matrix_elements, split_indices)) { + ValueVec outputs = select_outputs( + matrix_element, + matrix_element + .build_function( + fb, + me_args( + matrix_element, + [&](Value v) { return fb.batch_gather(split_index, v); } + ) + ) + .values() + ); + for (auto [output, split] : zip(outputs, split_results)) { + if (output) { + split.push_back(output); + split.push_back(split_index); + } + } + } + for (auto& split : split_results) { + results.push_back(split.empty() ? Value() : fb.batch_merge_by_index(split)); + } + } + + ValueVec outputs{results.at(0)}; + for (std::size_t i = 0; i < index_outputs.size(); ++i) { + Value result = results.at(i + 1); + outputs.push_back(result && _write_indices.at(i) ? result : indices.at(i)); + } + outputs.push_back(random.at(3)); + return {return_types().keys(), outputs}; +} diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 9a46c0c1a1..2f02f0ce1d 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1,4 +1,5 @@ #include +#include #include #include #include @@ -142,6 +143,65 @@ void named_vector_instance(py::module_& m, const char* name) { ); } +// One numpy array per LHE event field (shape (n,)) and per LHE particle field +// (shape (n, max_particle_count), zero padded), plus the particle counts. +py::dict lhe_events_to_numpy( + const std::vector& events, std::size_t max_particle_count +) { + py::ssize_t n = events.size(); + py::ssize_t m = max_particle_count; + auto event_field = [&](auto member) { + using T = std::remove_cvref_t; + py::array_t array(n); + auto view = array.template mutable_unchecked<1>(); + for (py::ssize_t i = 0; i < n; ++i) { + view(i) = events[i].*member; + } + return array; + }; + auto particle_field = [&](auto member) { + using T = std::remove_cvref_t; + py::array_t array({n, m}); + auto view = array.template mutable_unchecked<2>(); + for (py::ssize_t i = 0; i < n; ++i) { + auto& particles = events[i].particles; + for (py::ssize_t j = 0; j < m; ++j) { + view(i, j) = j < static_cast(particles.size()) + ? particles[j].*member + : T{}; + } + } + return array; + }; + py::array_t particle_count(n); + auto count_view = particle_count.mutable_unchecked<1>(); + for (py::ssize_t i = 0; i < n; ++i) { + count_view(i) = events[i].particles.size(); + } + + py::dict ret; + ret["process_id"] = event_field(&LHEEvent::process_id); + ret["weight"] = event_field(&LHEEvent::weight); + ret["scale"] = event_field(&LHEEvent::scale); + ret["alpha_qed"] = event_field(&LHEEvent::alpha_qed); + ret["alpha_qcd"] = event_field(&LHEEvent::alpha_qcd); + ret["particle_count"] = particle_count; + ret["pdg_id"] = particle_field(&LHEParticle::pdg_id); + ret["status_code"] = particle_field(&LHEParticle::status_code); + ret["mother1"] = particle_field(&LHEParticle::mother1); + ret["mother2"] = particle_field(&LHEParticle::mother2); + ret["color"] = particle_field(&LHEParticle::color); + ret["anti_color"] = particle_field(&LHEParticle::anti_color); + ret["px"] = particle_field(&LHEParticle::px); + ret["py"] = particle_field(&LHEParticle::py); + ret["pz"] = particle_field(&LHEParticle::pz); + ret["energy"] = particle_field(&LHEParticle::energy); + ret["mass"] = particle_field(&LHEParticle::mass); + ret["lifetime"] = particle_field(&LHEParticle::lifetime); + ret["spin"] = particle_field(&LHEParticle::spin); + return ret; +} + } // namespace PYBIND11_MODULE(_madspace_py, m) { @@ -515,6 +575,7 @@ PYBIND11_MODULE(_madspace_py, m) { {"fact_scale2", UMAMI_OUT_FACT_SCALE_2}, {"pdf", UMAMI_OUT_PDF}, {"bias", UMAMI_OUT_BIAS}, + {"matrix_element_ratio", UMAMI_OUT_MATRIX_ELEMENT_RATIO}, }) { umami_output_keys[name] = key; } @@ -1573,6 +1634,7 @@ PYBIND11_MODULE(_madspace_py, m) { {"fact_scale2_out", MatrixElement::fact_scale2_out}, {"pdf_out", MatrixElement::pdf_out}, {"bias_out", MatrixElement::bias_out}, + {"matrix_element_ratio_out", MatrixElement::matrix_element_ratio_out}, } ); matrix_element @@ -1585,7 +1647,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 @@ -1597,6 +1660,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( @@ -1606,13 +1670,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( @@ -1632,6 +1698,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, @@ -1766,6 +1833,37 @@ PYBIND11_MODULE(_madspace_py, m) { pydoc::doc("SubchannelWeights::channel_count") ); + py::classh( + m, "MappingChannelWeights", pydoc::doc("MappingChannelWeights") + ) + .def( + py::init< + const std::vector&, + const nested_vector2&, + const std::vector&, + const nested_vector2&>(), + py::arg("mappings"), + py::arg("flavor_masks"), + py::arg("diagram_channels") = std::vector{}, + py::arg("flavor_diagram_masks") = nested_vector2{}, + pydoc::doc("MappingChannelWeights::MappingChannelWeights") + ) + .def( + "channel_count", + &MappingChannelWeights::channel_count, + pydoc::doc("MappingChannelWeights::channel_count") + ) + .def( + "mirror_beams", + &MappingChannelWeights::mirror_beams, + pydoc::doc("MappingChannelWeights::mirror_beams") + ) + .def( + "has_diagrams", + &MappingChannelWeights::has_diagrams, + pydoc::doc("MappingChannelWeights::has_diagrams") + ); + py::classh( m, "MomentumPreprocessing", pydoc::doc("MomentumPreprocessing") ) @@ -1947,57 +2045,50 @@ PYBIND11_MODULE(_madspace_py, m) { {"goodjet", PartonLineScheme::goodjet}, } ); - mlm - .def( - py::init< - std::vector, - nested_vector3, - nested_vector2, - double, - JetScaleScheme, - ScaleScheme, - std::unordered_map, - double, - double, - double, - bool, - std::vector, - int, - PartonLineScheme, - AlphasScheme, - bool, - ClusteringMeasure, - ClusteringHistory, - bool>(), - py::arg("topologies"), - py::arg("permutations"), - py::arg("diagram_indices"), - py::arg("cm_energy"), - py::arg("jet_scale_scheme") = JetScaleScheme::production, - py::arg("scale_scheme") = ScaleScheme::clustering_mean, - py::arg("pdg_color_types") = std::unordered_map{}, - py::arg("xqcut") = 0., - py::arg("bw_cutoff") = 15., - py::arg("jet_radius") = 1.0, - py::arg("hadronic") = true, - py::arg("external_pdg_ids") = std::vector{}, - py::arg("max_jet_flavor") = 4, - py::arg("parton_line_scheme") = PartonLineScheme::goodjet, - py::arg("alphas_scheme") = AlphasScheme::per_vertex, - py::arg("pdf_reweighting") = true, - py::arg("clustering_measure") = ClusteringMeasure::fxfx, - py::arg("clustering_history") = ClusteringHistory::all_diagrams, - py::arg("clustering_mean_pdf_ladder") = true - ) - .def_property_readonly( - "pdf_absolute_pdgs", &MLMClustering::pdf_absolute_pdgs - ) - .def_property_readonly( - "clustering_measure", &MLMClustering::clustering_measure - ) - .def_property_readonly( - "clustering_history", &MLMClustering::clustering_history - ) + mlm.def( + py::init< + std::vector, + nested_vector3, + nested_vector2, + double, + JetScaleScheme, + ScaleScheme, + std::unordered_map, + double, + double, + double, + bool, + std::vector, + int, + PartonLineScheme, + AlphasScheme, + bool, + ClusteringMeasure, + ClusteringHistory, + bool>(), + py::arg("topologies"), + py::arg("permutations"), + py::arg("diagram_indices"), + py::arg("cm_energy"), + py::arg("jet_scale_scheme") = JetScaleScheme::production, + py::arg("scale_scheme") = ScaleScheme::clustering_mean, + py::arg("pdg_color_types") = std::unordered_map{}, + py::arg("xqcut") = 0., + py::arg("bw_cutoff") = 15., + py::arg("jet_radius") = 1.0, + py::arg("hadronic") = true, + py::arg("external_pdg_ids") = std::vector{}, + py::arg("max_jet_flavor") = 4, + py::arg("parton_line_scheme") = PartonLineScheme::goodjet, + py::arg("alphas_scheme") = AlphasScheme::per_vertex, + py::arg("pdf_reweighting") = true, + py::arg("clustering_measure") = ClusteringMeasure::fxfx, + py::arg("clustering_history") = ClusteringHistory::all_diagrams, + py::arg("clustering_mean_pdf_ladder") = true + ) + .def_property_readonly("pdf_absolute_pdgs", &MLMClustering::pdf_absolute_pdgs) + .def_property_readonly("clustering_measure", &MLMClustering::clustering_measure) + .def_property_readonly("clustering_history", &MLMClustering::clustering_history) .def_property_readonly( "diagram_start_states", &MLMClustering::diagram_start_states ) @@ -2363,6 +2454,7 @@ PYBIND11_MODULE(_madspace_py, m) { PdfGrid, DifferentialCrossSection::CachedPdf>&, bool, + bool, bool>(), py::arg("matrix_element"), py::arg("cm_energy"), @@ -2373,6 +2465,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("pdf2") = std::monostate{}, py::arg("input_momentum_fraction") = true, py::arg("decay") = false, + py::arg("return_matrix_element") = false, pydoc::doc("DifferentialCrossSection::DifferentialCrossSection") ) .def( @@ -2445,6 +2538,7 @@ PYBIND11_MODULE(_madspace_py, m) { std::size_t, const std::optional&, const std::optional&, + const std::optional&, int>(), py::arg("mapping"), py::arg("diff_xs"), @@ -2476,6 +2570,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("compressed_channel_weight_count") = 50, py::arg("custom_bias") = std::nullopt, py::arg("pdf_grid2") = std::nullopt, + py::arg("mapping_chan_weights") = std::nullopt, py::arg("mlm_history_qcd_power") = -1, pydoc::doc("Integrand::Integrand") ) @@ -2526,6 +2621,11 @@ PYBIND11_MODULE(_madspace_py, m) { &Integrand::energy_scale, pydoc::doc("Integrand::energy_scale") ) + .def( + "mapping_chan_weights", + &Integrand::mapping_chan_weights, + pydoc::doc("Integrand::mapping_chan_weights") + ) .def( "prop_chan_weights", &Integrand::prop_chan_weights, @@ -3035,6 +3135,11 @@ PYBIND11_MODULE(_madspace_py, m) { ) .def_readwrite( "done", &GeneratorStatus::done, pydoc::doc("GeneratorStatus::done") + ) + .def_readwrite( + "count_invalid_step_values", + &GeneratorStatus::count_invalid_step_values, + pydoc::doc("GeneratorStatus::count_invalid_step_values") ); py::classh(m, "Histogram", pydoc::doc("Histogram")) @@ -3264,7 +3369,8 @@ PYBIND11_MODULE(_madspace_py, m) { std::unordered_map, nested_vector2, nested_vector3, - nested_vector3>(), + nested_vector3, + nested_vector2>(), py::arg("process_id") = 0, py::arg("topologies") = std::vector{}, py::arg("permutations") = nested_vector3{}, @@ -3274,7 +3380,8 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("pdg_color_types") = std::unordered_map{}, py::arg("helicities") = nested_vector2{}, py::arg("pdg_ids") = nested_vector3{}, - py::arg("diagram_propagator_pdgs") = nested_vector3{} + py::arg("diagram_propagator_pdgs") = nested_vector3{}, + py::arg("pdg_id_weights") = nested_vector2{} ) .def_readwrite( "process_id", @@ -3325,6 +3432,11 @@ PYBIND11_MODULE(_madspace_py, m) { "diagram_propagator_pdgs", &LHECompleter::SubprocArgs::diagram_propagator_pdgs, pydoc::doc("LHECompleter::SubprocArgs::diagram_propagator_pdgs") + ) + .def_readwrite( + "pdg_id_weights", + &LHECompleter::SubprocArgs::pdg_id_weights, + pydoc::doc("LHECompleter::SubprocArgs::pdg_id_weights") ); py::classh(m, "MixMaxRandom", pydoc::doc("MixMaxRandom")) .def(py::init<>(), pydoc::doc("MixMaxRandom::MixMaxRandom")) @@ -3476,8 +3588,9 @@ PYBIND11_MODULE(_madspace_py, m) { &ChannelEventGenerator::load_json, py::arg("channel_json"), py::arg("contexts"), - py::arg("event_file"), - py::arg("weight_file"), + py::kw_only(), + py::arg("event_file") = "", + py::arg("weight_file") = "", py::arg("config"), pydoc::doc("ChannelEventGenerator::load_json") ) @@ -3523,7 +3636,160 @@ PYBIND11_MODULE(_madspace_py, m) { &ChannelEventGenerator::set_fixed_max_weight, py::arg("max_weight"), pydoc::doc("ChannelEventGenerator::set_fixed_max_weight") + ) + .def( + "max_weight", + &ChannelEventGenerator::max_weight, + pydoc::doc("ChannelEventGenerator::max_weight") + ) + .def( + "step_max_weight", + &ChannelEventGenerator::step_max_weight, + py::arg("step"), + pydoc::doc("ChannelEventGenerator::step_max_weight") + ) + .def( + "max_overweight_truncation", + &ChannelEventGenerator::max_overweight_truncation, + pydoc::doc("ChannelEventGenerator::max_overweight_truncation") + ) + .def( + "freeze_max_weight_after", + &ChannelEventGenerator::freeze_max_weight_after, + pydoc::doc("ChannelEventGenerator::freeze_max_weight_after") + ) + .def( + "set_target_count", + &ChannelEventGenerator::set_target_count, + py::arg("target_count"), + pydoc::doc("ChannelEventGenerator::set_target_count") + ) + .def( + "apply_step_truncation_budgets", + &ChannelEventGenerator::apply_step_truncation_budgets, + pydoc::doc("ChannelEventGenerator::apply_step_truncation_budgets") + ) + .def( + "set_fixed_step_max_weight", + &ChannelEventGenerator::set_fixed_step_max_weight, + py::arg("step"), + py::arg("max_weight"), + pydoc::doc("ChannelEventGenerator::set_fixed_step_max_weight") + ); + + py::classh weighting_step( + m, "WeightingStep", pydoc::doc("WeightingStep") + ); + add_enum( + weighting_step, + "Mode", + { + {"unweight", WeightingStep::unweight}, + {"reweight", WeightingStep::reweight}, + } + ); + weighting_step + .def( + py::init([](const std::vector& matrix_elements, + WeightingStep::Mode mode, + double max_overweight_truncation, + std::size_t freeze_max_weight_after, + const std::vector& + override_outputs) { + return WeightingStep{ + .matrix_elements = matrix_elements, + .mode = mode, + .max_overweight_truncation = max_overweight_truncation, + .freeze_max_weight_after = freeze_max_weight_after, + .override_outputs = override_outputs, + }; + }), + py::arg("matrix_elements"), + py::arg("mode") = WeightingStep::unweight, + py::arg("max_overweight_truncation") = 0.01, + py::arg("freeze_max_weight_after") = 10000, + py::arg("override_outputs") = + std::vector{}, + pydoc::doc("WeightingStep") + ) + .def_readwrite( + "matrix_elements", + &WeightingStep::matrix_elements, + pydoc::doc("WeightingStep::matrix_elements") + ) + .def_readwrite("mode", &WeightingStep::mode, pydoc::doc("WeightingStep::mode")) + .def_readwrite( + "max_overweight_truncation", + &WeightingStep::max_overweight_truncation, + pydoc::doc("WeightingStep::max_overweight_truncation") + ) + .def_readwrite( + "freeze_max_weight_after", + &WeightingStep::freeze_max_weight_after, + pydoc::doc("WeightingStep::freeze_max_weight_after") + ) + .def_readwrite( + "override_outputs", + &WeightingStep::override_outputs, + pydoc::doc("WeightingStep::override_outputs") + ) + .def( + "returns_ratio", + &WeightingStep::returns_ratio, + pydoc::doc("WeightingStep::returns_ratio") + ); + + py::classh(m, "WeightingGroup", pydoc::doc("WeightingGroup")) + .def( + py::init([](const std::vector& channel_indices, + const std::vector& steps, + std::optional first_pass_max_overweight_truncation, + std::optional first_pass_freeze_max_weight_after) { + return WeightingGroup{ + .channel_indices = channel_indices, + .steps = steps, + .first_pass_max_overweight_truncation = + first_pass_max_overweight_truncation, + .first_pass_freeze_max_weight_after = + first_pass_freeze_max_weight_after, + }; + }), + py::arg("channel_indices"), + py::arg("steps"), + py::arg("first_pass_max_overweight_truncation") = py::none(), + py::arg("first_pass_freeze_max_weight_after") = py::none(), + pydoc::doc("WeightingGroup") + ) + .def_readwrite( + "first_pass_max_overweight_truncation", + &WeightingGroup::first_pass_max_overweight_truncation, + pydoc::doc("WeightingGroup::first_pass_max_overweight_truncation") + ) + .def_readwrite( + "first_pass_freeze_max_weight_after", + &WeightingGroup::first_pass_freeze_max_weight_after, + pydoc::doc("WeightingGroup::first_pass_freeze_max_weight_after") + ) + .def_readwrite( + "channel_indices", + &WeightingGroup::channel_indices, + pydoc::doc("WeightingGroup::channel_indices") + ) + .def_readwrite( + "steps", &WeightingGroup::steps, pydoc::doc("WeightingGroup::steps") ); + m.def( + "weighting_groups_to_json", + &weighting_groups_to_json, + py::arg("groups"), + pydoc::doc("weighting_groups_to_json") + ); + m.def( + "weighting_groups_from_json", + &weighting_groups_from_json, + py::arg("text"), + pydoc::doc("weighting_groups_from_json") + ); py::classh(m, "PdfMemberSpec", pydoc::doc("PdfMemberSpec")) .def( @@ -3715,9 +3981,7 @@ PYBIND11_MODULE(_madspace_py, m) { "is_scale", &Variation::is_scale, pydoc::doc("Variation::is_scale") ) .def_property_readonly( - "is_emission", - &Variation::is_emission, - pydoc::doc("Variation::is_emission") + "is_emission", &Variation::is_emission, pydoc::doc("Variation::is_emission") ); py::classh(m, "PdfGroupInfo", pydoc::doc("PdfGroupInfo")) .def_readonly( @@ -3862,10 +4126,12 @@ PYBIND11_MODULE(_madspace_py, m) { int particle_flags = (momenta.empty() ? ParticleRecord::f_none : ParticleRecord::f_particle_data) | - (alphas_scales.empty() ? ParticleRecord::f_none - : ParticleRecord::f_alphas_scales) | - (ladder_lead.empty() ? ParticleRecord::f_none - : ParticleRecord::f_pdf_ladder); + (alphas_scales.empty() + ? ParticleRecord::f_none + : ParticleRecord::f_alphas_scales) | + (ladder_lead.empty() + ? ParticleRecord::f_none + : ParticleRecord::f_pdf_ladder); DataLayout layout( EventRecord::layout( EventRecord::f_weight | EventRecord::f_subproc_index | @@ -3900,7 +4166,8 @@ PYBIND11_MODULE(_madspace_py, m) { particle.py() = p.at(2); particle.pz() = p.at(3); } - for (std::size_t j = 0; !alphas_scales.empty() && j < particle_count; + for (std::size_t j = 0; + !alphas_scales.empty() && j < particle_count; ++j) { auto& row = alphas_scales.at(i); buffer.particle(i, j).alphas_scale() = @@ -3912,8 +4179,8 @@ PYBIND11_MODULE(_madspace_py, m) { auto& lead = ladder_lead.at(i); particle.ladder_lead_q() = j < 2 ? lead.at(j) : 0.; bool has_step = !ladder.empty() && j < ladder.at(i).size(); - auto step = has_step ? ladder.at(i).at(j) - : std::vector(4, 0.); + auto step = + has_step ? ladder.at(i).at(j) : std::vector(4, 0.); particle.ladder_x() = step.at(0); particle.ladder_pdg() = step.at(1); particle.ladder_q() = step.at(2); @@ -4081,7 +4348,8 @@ PYBIND11_MODULE(_madspace_py, m) { const std::vector>&, std::uint64_t, std::shared_ptr, - const GeneratorConfig&>(), + const GeneratorConfig&, + const std::vector&>(), py::arg("contexts"), py::arg("channels"), py::arg("seed"), @@ -4091,6 +4359,7 @@ PYBIND11_MODULE(_madspace_py, m) { EventGenerator::default_config, "EventGenerator.default_config" ), + py::arg("groups") = std::vector{}, pydoc::doc("EventGenerator::EventGenerator") ) .def( @@ -4109,11 +4378,13 @@ PYBIND11_MODULE(_madspace_py, m) { .def( "optimize", &EventGenerator::optimize, + py::call_guard(), pydoc::doc("EventGenerator::optimize") ) .def( "fix_max_weights", &EventGenerator::fix_max_weights, + py::call_guard(), pydoc::doc("EventGenerator::fix_max_weights") ) .def( @@ -4168,6 +4439,90 @@ PYBIND11_MODULE(_madspace_py, m) { pydoc::doc("EventGenerator::channels") ); + py::classh(m, "EventStream", pydoc::doc("EventStream")) + .def( + py::init( + [](ContextPtr context, + const std::vector>& channels, + std::uint64_t seed, + const LHECompleter& lhe_completer, + const GeneratorConfig& config, + const std::vector& groups) { + // Destroying the stream waits for its running jobs, which can + // need the GIL, e.g. for matrix elements in Python + return std::shared_ptr( + new EventStream( + context, channels, seed, lhe_completer, config, groups + ), + [](EventStream* stream) { + if (PyGILState_Check()) { + py::gil_scoped_release release; + delete stream; + } else { + delete stream; + } + } + ); + } + ), + py::arg("context"), + py::arg("channels"), + py::arg("seed"), + py::arg("lhe_completer"), + py::arg("config"), + py::arg("groups") = std::vector{}, + pydoc::doc("EventStream::EventStream") + ) + .def( + "next_events", + &EventStream::next_events, + py::arg("count"), + py::call_guard(), + pydoc::doc("EventStream::next_events") + ) + .def( + "next_batch", + [](EventStream& stream, std::size_t count) { + std::vector events; + { + py::gil_scoped_release release; + events = stream.next_events(count); + } + return lhe_events_to_numpy(events, stream.max_particle_count()); + }, + py::arg("count"), + "Like next_events, as a dict of numpy arrays: the LHE event fields " + "with shape (count,), the LHE particle fields with shape (count, " + "max_particle_count), padded with zeros, and particle_count." + ) + .def( + "close", + &EventStream::close, + py::call_guard(), + pydoc::doc("EventStream::close") + ) + .def("status", &EventStream::status, pydoc::doc("EventStream::status")) + .def( + "channel_status", + &EventStream::channel_status, + pydoc::doc("EventStream::channel_status") + ) + .def( + "channel_probabilities", + &EventStream::channel_probabilities, + pydoc::doc("EventStream::channel_probabilities") + ) + .def( + "event_count", + &EventStream::event_count, + pydoc::doc("EventStream::event_count") + ) + .def( + "max_particle_count", + &EventStream::max_particle_count, + pydoc::doc("EventStream::max_particle_count") + ); + py::classh logger(m, "Logger", pydoc::doc("Logger")); add_enum( logger, @@ -4251,5 +4606,12 @@ PYBIND11_MODULE(_madspace_py, m) { } }; EventGenerator::set_abort_check_function(abort_check_function); + // EventStream waits with the GIL released + EventStream::set_abort_check_function([] { + py::gil_scoped_acquire gil; + if (PyErr_CheckSignals() != 0) { + throw py::error_already_set(); + } + }); MadnisTraining::set_abort_check_function(abort_check_function); } diff --git a/madspace/tests/test_color_ordered_massive_side.py b/madspace/tests/test_color_ordered_massive_side.py new file mode 100644 index 0000000000..19b42d7346 --- /dev/null +++ b/madspace/tests/test_color_ordered_massive_side.py @@ -0,0 +1,74 @@ +"""A side of two massive particles in ColorOrderedMapping with a cut elsewhere. + +g g > t t~ g in the color order t g g g t~ (0 2 4 1 3 in particle numbers): the +gluon is alone on one side, the top pair is the other, built by the 2->3 step +off the opposite beam. Any cut turns on the transverse-energy floors of that +step, which are m_T >= m for a top without a pt cut. The floor on the momentum +transfer to the beam was etmin^2 = m_T^2, right only for massless particles: +the bound is m_T^2 - m^2 = pt^2. It removed the tops with |t| < m_t^2, about 6% +of the volume and 20% of it at small top pt, and the AmpliCol channels of +g g > t t~ g came out 25% low. +""" + +import math + +import pytest +import torch +from madspace_testing import generator, rand, randint + +import madspace as ms + +N = 400_000 +CM_ENERGY = 600.0 +ORDER = [2, 0, 4, 1, 3] +MASSES = [173.0, 173.0, 0.0] +PT_MIN = 20.0 + + +def sample(pt_min, seed, masses=MASSES): + mapping = ms.ColorOrderedMapping( + ms.ColorOrderedMapping.Options(ORDER), pt_min=pt_min + ) + g = generator(seed) + r = rand((N, mapping.random_dim()), generator=g) + d = randint(0, 2, (N, mapping.discrete_dim()), generator=g, dtype=torch.int32) + inputs = [r[:, i].contiguous() for i in range(r.shape[1])] + inputs += [d[:, j].contiguous() for j in range(d.shape[1])] + conditions = [torch.full((N,), CM_ENERGY)] + [torch.full((N,), m) for m in masses] + *momenta, det = mapping.map_forward(inputs, conditions) + det = torch.nan_to_num(det) * 2.0 ** mapping.discrete_dim() + return mapping, inputs, conditions, momenta, r, det + + +def pt(p): + return torch.sqrt(p[:, 1] ** 2 + p[:, 2] ** 2) + + +@pytest.mark.parametrize("top_pt_max", [math.inf, 50.0], ids=["all", "top pt < 50"]) +def test_cut_handed_to_mapping_keeps_volume(top_pt_max): + """The volume passing the gluon cut is the same whether the mapping knows + the cut or it is applied afterwards.""" + _, _, _, p_free, _, det_free = sample([0.0, 0.0, 0.0], 1) + _, _, _, p_cut, _, det_cut = sample([0.0, 0.0, PT_MIN], 2) + # beams first: the top is momenta[2], the gluon momenta[4] + external = torch.where( + (pt(p_free[4]) >= PT_MIN) & (pt(p_free[2]) < top_pt_max), det_free, 0.0 + ) + piped = torch.where( + (pt(p_cut[4]) >= PT_MIN) & (pt(p_cut[2]) < top_pt_max), det_cut, 0.0 + ) + error = math.sqrt((external.var(correction=0) + piped.var(correction=0)) / N) + assert error > 0.0 + assert abs(piped.mean() - external.mean()).item() < 5.0 * error + + +def test_round_trip_massive_side_with_cut(): + mapping, _, conditions, momenta, r, det = sample([0.0, 0.0, PT_MIN], 3) + out = mapping.map_inverse(momenta, conditions) + r_inv = torch.stack(out[: mapping.random_dim()], dim=1) + ok = det > 0 + rt = torch.abs(det / 2.0 ** mapping.discrete_dim() * out[-1] - 1.0)[ok] + assert torch.mean((rt > 1e-2).double()) < 1e-4 + assert torch.quantile(rt, 0.999) < 1e-6 + r_err = torch.max(torch.abs(r_inv - r), dim=1).values[ok] + assert torch.mean((r_err > 1e-6).double()) < 1e-4 diff --git a/madspace/tests/test_event_stream.py b/madspace/tests/test_event_stream.py new file mode 100644 index 0000000000..4e75b35bee --- /dev/null +++ b/madspace/tests/test_event_stream.py @@ -0,0 +1,241 @@ +"""EventStream: in-memory event generation from channels with a fixed maximum +weight, as stored in a fix_max_weight gridpack. + +The channels of an e+ e- -> d d~ g process (flat matrix element, so the weights +only vary through the phase-space mapping) are surveyed and their maximum +weights fixed with EventGenerator.fix_max_weights(), then reloaded from their +JSON like a gridpack does. +""" + +import json +import os + +import numpy as np +import pytest + +import madspace as ms + +CM_ENERGY = 1000.0 +PIDS = [11, -11, 1, -1, 21] +SEED = 1234 + + +def build_topology(): + propagators = [ + ms.Propagator( + mass=0.0, width=0.0, integration_order=0, e_min=0.0, e_max=0.0, pdg_id=pid + ) + for pid in (22, -1) + ] + diagram = ms.Diagram( + [0.0, 0.0], + [0.0, 0.0, 0.0], + propagators, + [["i0", "i1", "p0"], ["o0", "p1", "p0"], ["o1", "o2", "p1"]], + ) + return ms.Topology.topologies(diagram)[0] + + +TOPOLOGY = build_topology() +PERMUTATIONS = [[0, 1, 2, 3, 4]] + + +def lhe_completer(): + subproc_args = ms.SubprocArgs( + process_id=0, + topologies=[TOPOLOGY], + permutations=[PERMUTATIONS], + diagram_indices=[[0]], + diagram_color_indices=[[[0]]], + color_flows=[[[0, 0], [0, 0], [501, 0], [0, 502], [502, 501]]], + pdg_color_types={11: 1, -11: 1, 1: 3, -1: -3, 22: 1, 21: 8}, + helicities=[[1, -1, 1, -1, 1]], + pdg_ids=[[PIDS]], + ) + return ms.LHECompleter([subproc_args], bw_cutoff=15.0) + + +def make_config(): + config = ms.GeneratorConfig() + config.cpu_batch_size = 1000 + config.freeze_max_weight_after = 1000 + config.max_overweight_truncation = 0.05 + return config + + +def make_context(run_dir, thread_count): + context = ms.Context(ms.cpu_device(), thread_count) + # RunningCoupling only reads these two keys of an LHAPDF .info file + alphas_file = os.path.join(run_dir, "alphas.info") + with open(alphas_file, "w") as f: + f.write("AlphaS_Qs: [1.0, 10.0, 100.0, 1000.0, 10000.0]\n") + f.write("AlphaS_Vals: [0.3, 0.2, 0.12, 0.1, 0.08]\n") + ms.AlphaSGrid(alphas_file).initialize_globals(context) + return context, alphas_file + + +def make_channel(context, alphas_file, run_dir, name, pt_min): + cuts = ms.Cuts( + [ + ms.CutItem(ms.Observable(PIDS, ms.Observable.obs_pt, [[i]]), min=pt_min) + for i in (2, 3, 4) + ] + ) + mapping = ms.PhaseSpaceMapping( + TOPOLOGY, CM_ENERGY, permutations=PERMUTATIONS, leptonic=True, cuts=cuts + ) + # 0xBADCAFE: the built-in flat test matrix element, no library needed + matrix_element = ms.MatrixElement( + 0xBADCAFE, + len(PIDS), + ms.Integrand.matrix_element_inputs, + ms.Integrand.matrix_element_outputs, + 1, + True, + ) + diff_xs = ms.DifferentialCrossSection( + matrix_element=matrix_element, + cm_energy=CM_ENERGY, + running_coupling=None, + energy_scale=ms.CachedScale(), + ) + integrand = ms.Integrand( + mapping, + [diff_xs], + running_coupling=ms.RunningCoupling(ms.AlphaSGrid(alphas_file)), + energy_scale=ms.EnergyScale(len(PIDS), 91.188), + channel_indices=[0], + ) + return ms.ChannelEventGenerator( + contexts=[context], + integrand=integrand, + event_file=os.path.join(run_dir, f"events.{name}.npy"), + weight_file=os.path.join(run_dir, f"weights.{name}.npy"), + config=make_config(), + subprocess_index=0, + name=name, + histograms=None, + ) + + +@pytest.fixture(scope="module") +def gridpack(tmp_path_factory): + """Channel JSONs with fixed max weights and the globals, like a gridpack.""" + run_dir = str(tmp_path_factory.mktemp("event_stream")) + context, alphas_file = make_context(run_dir, 4) + channels = [ + make_channel(context, alphas_file, run_dir, "a", 20.0), + make_channel(context, alphas_file, run_dir, "b", 60.0), + ] + generator = ms.EventGenerator([context], channels, seed=1, config=make_config()) + generator.survey() + generator.fix_max_weights() + globals_file = os.path.join(run_dir, "globals.tar") + context.save_globals(globals_file) + return { + "run_dir": run_dir, + "channels": [channel.to_json(True) for channel in channels], + "globals": globals_file, + "integral": generator.status().mean, + } + + +def make_stream(gridpack, thread_count=4, seed=SEED): + context = ms.Context(ms.cpu_device(), thread_count) + context.load_globals(gridpack["globals"]) + config = make_config() + channels = [ + ms.ChannelEventGenerator.load_json(channel, [context], config=config) + for channel in gridpack["channels"] + ] + return ms.EventStream(context, channels, seed, lhe_completer(), config) + + +def stream_batches(stream, batch_sizes): + batches = [stream.next_batch(size) for size in batch_sizes] + return {key: np.concatenate([b[key] for b in batches]) for key in batches[0]} + + +def test_reproducible_across_batch_sizes_and_threads(gridpack): + reference = stream_batches(make_stream(gridpack, 4), [3000]) + for thread_count, batch_sizes in [(1, [1000, 1000, 1000]), (3, [1, 999, 2000])]: + events = stream_batches(make_stream(gridpack, thread_count), batch_sizes) + for key, values in reference.items(): + np.testing.assert_array_equal(events[key], values, err_msg=key) + + +def test_seed_changes_events(gridpack): + events1 = make_stream(gridpack, seed=1).next_batch(100) + events2 = make_stream(gridpack, seed=2).next_batch(100) + assert not np.array_equal(events1["px"], events2["px"]) + + +def test_lhe_events_match_batches(gridpack): + lhe_events = make_stream(gridpack).next_events(50) + batch = make_stream(gridpack).next_batch(50) + for i, event in enumerate(lhe_events): + assert event.weight == batch["weight"][i] + assert [p.pdg_id for p in event.particles] == list(batch["pdg_id"][i]) + assert [p.px for p in event.particles] == list(batch["px"][i]) + + +def test_weights_and_particles(gridpack): + stream = make_stream(gridpack) + batch = stream.next_batch(5000) + weights = batch["weight"] + # weights are 1, except for overweight events + assert np.all(weights >= 1.0) + assert np.mean(weights == 1.0) > 0.9 + assert np.all(batch["particle_count"] == len(PIDS)) + assert set(map(tuple, batch["pdg_id"])) == {tuple(PIDS)} + assert batch["px"].shape == (5000, stream.max_particle_count()) + # momentum conservation + outgoing = batch["status_code"] == 1 + assert np.allclose(np.sum(batch["px"] * outgoing, axis=1), 0.0, atol=1e-6) + + +def test_status(gridpack): + stream = make_stream(gridpack) + stream.next_batch(2000) + status = stream.status() + assert stream.event_count() == 2000 + assert status.count_unweighted == 2000 + assert status.count > 0 + assert status.mean == pytest.approx(gridpack["integral"], rel=0.1) + channel_status = stream.channel_status() + assert sum(s.count_unweighted for s in channel_status) == 2000 + assert sum(stream.channel_probabilities()) == pytest.approx(1.0) + + +def test_channel_fractions_follow_probabilities(gridpack): + stream = make_stream(gridpack) + stream.next_batch(20000) + probs = np.array(stream.channel_probabilities()) + counts = np.array([s.count_unweighted for s in stream.channel_status()]) + fractions = counts / counts.sum() + assert np.allclose(fractions, probs, atol=0.02) + + +def test_no_files_written(gridpack, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + stream = make_stream(gridpack) + stream.next_batch(1000) + stream.close() + assert os.listdir(tmp_path) == [] + + +def test_requires_fixed_max_weight(gridpack): + context = ms.Context(ms.cpu_device(), 1) + context.load_globals(gridpack["globals"]) + config = make_config() + channels = [] + for channel in gridpack["channels"]: + data = json.loads(channel) + del data["max_weight"] + channels.append( + ms.ChannelEventGenerator.load_json( + json.dumps(data), [context], config=config + ) + ) + with pytest.raises(ValueError, match="fixed maximum weight"): + ms.EventStream(context, channels, SEED, lhe_completer(), config) diff --git a/madspace/tests/test_event_stream_reweight.py b/madspace/tests/test_event_stream_reweight.py new file mode 100644 index 0000000000..2babcce75b --- /dev/null +++ b/madspace/tests/test_event_stream_reweight.py @@ -0,0 +1,132 @@ +"""EventStream with a reweighting step after the first unweighting, like the +default [amplicol] mode = "reweight": the streamed events carry the ratio of the +step matrix element to the integrand one, not unit weights. Uses the setup of +test_weighting_steps. +""" + +import os + +import numpy as np +import pytest +from test_event_stream import CM_ENERGY, lhe_completer, make_config, make_context +from test_integrand_me_outputs import python_matrix_element +from test_weighting_steps import FULL_ME_OUTPUTS, generate, make_integrand, true_me + +import madspace as ms + +PT_MINS = (20.0, 60.0) + + +def biased_me(momenta): + """Does not average to the true one, so that the ratio changes the norm.""" + return true_me(momenta) * (0.6 + 0.8 * momenta[:, 2, 0] / CM_ENERGY) + + +def full_me_with_indices(momenta): + zeros = np.zeros(len(momenta), np.int32) + return true_me(momenta), zeros, zeros, zeros + + +def load_matrix_elements(context): + return ( + python_matrix_element(context, biased_me), + python_matrix_element(context, full_me_with_indices, FULL_ME_OUTPUTS), + ) + + +@pytest.fixture(scope="module") +def reweight_gridpack(tmp_path_factory): + """Two grouped channels with fixed maxima and a reweighting step.""" + config = make_config() + run_dir = str(tmp_path_factory.mktemp("reweight")) + context, alphas_file = make_context(run_dir, 2) + biased, full = load_matrix_elements(context) + groups = [ms.WeightingGroup([0, 1], [ms.WeightingStep([full], mode="reweight")])] + channels = [ + ms.ChannelEventGenerator( + contexts=[context], + integrand=make_integrand(alphas_file, biased, pt_min), + event_file=os.path.join(run_dir, f"events.{i}.npy"), + weight_file=os.path.join(run_dir, f"weights.{i}.npy"), + config=config, + subprocess_index=0, + name=str(i), + histograms=None, + ) + for i, pt_min in enumerate(PT_MINS) + ] + generator = ms.EventGenerator( + [context], channels, seed=1, config=config, groups=groups + ) + generator.survey() + generator.fix_max_weights() + globals_file = os.path.join(run_dir, "globals.tar") + context.save_globals(globals_file) + return { + "channels": [channel.to_json(True) for channel in channels], + "groups": ms.weighting_groups_to_json(groups), + "globals": globals_file, + } + + +def make_stream(gridpack, thread_count=4, seed=5): + context = ms.Context(ms.cpu_device(), thread_count) + context.load_globals(gridpack["globals"]) + load_matrix_elements(context) + config = make_config() + channels = [ + ms.ChannelEventGenerator.load_json( + channel_json, [context], event_file="", weight_file="", config=config + ) + for channel_json in gridpack["channels"] + ] + groups = ms.weighting_groups_from_json(gridpack["groups"]) + return ms.EventStream( + context, channels, seed, lhe_completer(), config, groups=groups + ) + + +def test_event_stream_reweight(reweight_gridpack, tmp_path): + stream = make_stream(reweight_gridpack) + batch = stream.next_batch(20000) + weights = batch["weight"] + # the ratio of the full to the biased matrix element, times the first-pass + # weight (1, or more for the overweights) + ratio = 1.0 / (0.6 + 0.8 * batch["energy"][:, 2] / CM_ENERGY) + assert np.mean(np.isclose(weights, ratio, rtol=1e-6)) > 0.9 + assert np.all(weights >= ratio * (1.0 - 1e-6)) + + # mean weight x first-pass cross section = full cross section + first_pass, _, _ = generate( + str(tmp_path / "biased"), + biased_me, + pt_mins=PT_MINS, + grouped=(0, 1), + target_count=1000, + ) + full, full_events, _ = generate( + str(tmp_path / "full"), + true_me, + pt_mins=PT_MINS, + grouped=(0, 1), + target_count=20000, + ) + status = stream.status() + mean_weight_error = np.std(weights) / np.sqrt(len(weights)) + xs_error = np.hypot( + mean_weight_error * first_pass.mean, np.mean(weights) * first_pass.error + ) + assert abs(np.mean(weights) * first_pass.mean - full.mean) < 4 * np.hypot( + xs_error, full.error + ) + assert abs(status.mean - full.mean) < 4 * np.hypot(status.error, full.error) + + # the weighted events follow the full matrix element + energies = [ + np.sum(w * e) / np.sum(w) + for w, e in ( + (weights, batch["energy"][:, 2]), + (full_events["weight"], full_events["part3_energy"]), + ) + ] + assert energies[0] == pytest.approx(energies[1], rel=0.01) diff --git a/madspace/tests/test_integrand_me_outputs.py b/madspace/tests/test_integrand_me_outputs.py new file mode 100644 index 0000000000..386efac57b --- /dev/null +++ b/madspace/tests/test_integrand_me_outputs.py @@ -0,0 +1,230 @@ +"""Integrand with matrix elements that only return |M|^2: the channel weights, +helicity, color and diagram index are optional outputs of the matrix element. + +Uses the e+ e- -> d d~ g setup of test_event_stream with PythonFunction matrix +elements, as a stand-in for surrogates that only predict |M|^2. +""" + +import os + +import numpy as np +import pytest +from test_event_stream import ( + CM_ENERGY, + PERMUTATIONS, + PIDS, + TOPOLOGY, + make_config, + make_context, +) + +import madspace as ms + + +def python_matrix_element(context, function, outputs=("matrix_element",), 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, inputs, list(outputs), particle_count=len(PIDS)) + ) + output_enums = { + "matrix_element": ms.MatrixElement.matrix_element_out, + "color_index": ms.MatrixElement.color_index_out, + "helicity_index": ms.MatrixElement.helicity_index_out, + "diagram_index": ms.MatrixElement.diagram_index_out, + "matrix_element_ratio": ms.MatrixElement.matrix_element_ratio_out, + } + return ms.MatrixElement( + api.index(), + len(PIDS), + [ms.MatrixElement.momenta_in] + + ([] if channel < 0 else [ms.MatrixElement.channel_in]), + [output_enums[output] for output in outputs], + channel=channel, + ) + + +def flat_matrix_element(): + return ms.MatrixElement( + 0xBADCAFE, + len(PIDS), + ms.Integrand.matrix_element_inputs, + ms.Integrand.matrix_element_outputs, + 1, + True, + ) + + +def make_integrand(alphas_file, matrix_element, return_matrix_element=False, **kwargs): + mapping = ms.PhaseSpaceMapping( + TOPOLOGY, + CM_ENERGY, + permutations=PERMUTATIONS, + leptonic=True, + cuts=ms.Cuts( + [ + ms.CutItem(ms.Observable(PIDS, ms.Observable.obs_pt, [[i]]), min=20.0) + for i in (2, 3, 4) + ] + ), + ) + diff_xs = ms.DifferentialCrossSection( + matrix_element=matrix_element, + cm_energy=CM_ENERGY, + running_coupling=None, + energy_scale=ms.CachedScale(), + return_matrix_element=return_matrix_element, + ) + return ms.Integrand( + mapping, + [diff_xs], + running_coupling=ms.RunningCoupling(ms.AlphaSGrid(alphas_file)), + energy_scale=ms.EnergyScale(len(PIDS), 91.188), + channel_indices=[0], + **kwargs, + ) + + +def generate(run_dir, context, integrand, name): + config = make_config() + config.target_count = 2000 + channel = ms.ChannelEventGenerator( + contexts=[context], + integrand=integrand, + event_file=os.path.join(run_dir, f"events.{name}.npy"), + weight_file=os.path.join(run_dir, f"weights.{name}.npy"), + config=config, + subprocess_index=0, + name=name, + histograms=None, + ) + generator = ms.EventGenerator([context], [channel], seed=7, config=config) + generator.survey() + generator.generate() + output_file = os.path.join(run_dir, f"combined.{name}.npy") + generator.combine_to_compact_npy(output_file) + return generator.status(), np.load(output_file) + + +def test_only_matrix_element_output(tmp_path): + """Without the index outputs the events carry -1, and |M|^2 = 1 reproduces + the integral of the built-in flat matrix element.""" + context, alphas_file = make_context(str(tmp_path), 2) + matrix_element = python_matrix_element( + context, lambda momenta: np.ones(len(momenta)) + ) + status, events = generate( + str(tmp_path), context, make_integrand(alphas_file, matrix_element), "py" + ) + flat_status, flat_events = generate( + str(tmp_path), + context, + make_integrand(alphas_file, flat_matrix_element()), + "flat", + ) + assert len(events) == len(flat_events) == 2000 + for field in ("helicity_index", "color_index", "diagram_index"): + assert (events[field] == -1).all() + assert (flat_events[field] >= 0).all() + error = np.hypot(status.error, flat_status.error) + assert abs(status.mean - flat_status.mean) < 5 * error + + +def test_provided_index_output(tmp_path): + """An index the matrix element returns is written to the events.""" + context, alphas_file = make_context(str(tmp_path), 1) + matrix_element = python_matrix_element( + context, + lambda momenta: (np.ones(len(momenta)), np.full(len(momenta), 1, np.int32)), + ("matrix_element", "helicity_index"), + ) + _, events = generate( + str(tmp_path), context, make_integrand(alphas_file, matrix_element), "hel" + ) + assert (events["helicity_index"] == 1).all() + assert (events["color_index"] == -1).all() + + +def test_raw_matrix_element_output(tmp_path): + """Returning the raw |M|^2 adds an integrand output without changing the + generated events.""" + context, alphas_file = make_context(str(tmp_path), 1) + matrix_element = python_matrix_element( + context, lambda momenta: 1.0 + momenta[:, 2, 0] / CM_ENERGY + ) + results = [ + generate( + str(tmp_path), + context, + make_integrand(alphas_file, matrix_element, raw), + f"raw{raw}", + ) + for raw in (False, True) + ] + (status_a, events_a), (status_b, events_b) = results + assert status_a.mean == status_b.mean + assert np.array_equal(events_a, events_b) + + +def test_diagram_channel_weights_need_amp2(tmp_path): + context, alphas_file = make_context(str(tmp_path), 1) + matrix_element = python_matrix_element( + context, lambda momenta: np.ones(len(momenta)) + ) + with pytest.raises(ValueError, match="diagram_amp2"): + make_integrand( + alphas_file, + matrix_element, + first_chan_weight_remap=[[0, 1]], + first_remapped_chan_count=2, + ) + + +@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_ordered=( + None if order is None else ms.ColorOrderedMapping.Options(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 diff --git a/madspace/tests/test_lhe.py b/madspace/tests/test_lhe.py index 63eb3ba93b..9f83841b0e 100644 --- a/madspace/tests/test_lhe.py +++ b/madspace/tests/test_lhe.py @@ -491,3 +491,105 @@ 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)) + + +def weighted_flavor_completer(pdg_id_weights): + """e+ e- > l+ l- with two flavors of two options each.""" + options = [ + [[11, -11, 13, -13], [11, -11, 15, -15]], + [[-11, 11, 13, -13], [-11, 11, 15, -15]], + ] + subproc_args = ms.SubprocArgs( + topologies=[chain_topology([0.0] * 4)], + permutations=[[]], + diagram_indices=[[]], + diagram_color_indices=[[]], + color_flows=[[[0, 0]] * 4], + pdg_color_types={pdg: 1 for pdg in (11, -11, 13, -13, 15, -15)}, + helicities=[[1, -1, 1, -1]], + pdg_ids=options, + pdg_id_weights=pdg_id_weights, + ) + return ms.LHECompleter([subproc_args], bw_cutoff=BW_CUTOFF) + + +def drawn_taus(lhe_completer, flavor_index, count, seed): + rand_gen = ms.MixMaxRandom(seed) + taus = 0 + for _ in range(count): + event = build_event(np.ones((4, 4)) * [100.0, 0.0, 0.0, 0.0]) + lhe_completer.complete_event_data(event, 0, -1, 0, flavor_index, 0, rand_gen) + taus += event.particles[2].pdg_id == 15 + return taus + + +def test_pdg_id_options_are_drawn_in_proportion_to_their_weights(): + """Options are written with probability proportional to their weight; + equal weights and flavors without weights draw them uniformly.""" + lhe_completer = weighted_flavor_completer([[1.0, 3.0], []]) + count = 4000 + assert drawn_taus(lhe_completer, 0, count, 5) / count == approx(0.75, abs=0.03) + assert drawn_taus(lhe_completer, 1, count, 5) / count == approx(0.5, abs=0.03) + + with tempfile.TemporaryDirectory() as tmpdir: + save_path = os.path.join(tmpdir, "lhe_completer.json") + lhe_completer.save(save_path) + loaded = ms.LHECompleter.load(save_path) + assert drawn_taus(loaded, 0, 200, 9) == drawn_taus(lhe_completer, 0, 200, 9) + + # equal weights consume the random numbers like the uniform draw + equal = weighted_flavor_completer([[2.0, 2.0], [2.0, 2.0]]) + uniform = weighted_flavor_completer([]) + assert drawn_taus(equal, 0, 200, 3) == drawn_taus(uniform, 0, 200, 3) + + +@pytest.mark.parametrize( + "pdg_id_weights", [[[1.0, 0.0], []], [[1.0, -1.0], []], [[1.0], []], [[1.0, 1.0]]] +) +def test_invalid_pdg_id_weights_raise(pdg_id_weights): + with pytest.raises(ValueError): + weighted_flavor_completer(pdg_id_weights) diff --git a/madspace/tests/test_mapping_channel_weights.py b/madspace/tests/test_mapping_channel_weights.py new file mode 100644 index 0000000000..eca5dfe122 --- /dev/null +++ b/madspace/tests/test_mapping_channel_weights.py @@ -0,0 +1,159 @@ +"""MappingChannelWeights: channel weights from the densities of the channel mappings. + +- The weights are the masked densities of the inverse mappings, normalised per + event, and the sampled diagram follows them. +- A color-ordered chain with a resonance among its children integrates to the + same as the plain chain, and can be inverted at the events of the plain one. +""" + +import pytest +import torch +from madspace_testing import default_context, generator, rand, randint + +import madspace as ms + +E_CM = 13000.0 +M_Z, W_Z = 91.1876, 2.4952 +PIDS = [2, -2, 21, 11, -11] + + +def cuts(): + O = ms.Observable + return ms.Cuts( + [ + ms.CutItem(O(PIDS, O.obs_pt, [O.jet_pids]), min=20.0), + ms.CutItem(O(PIDS, O.obs_pt, [[11, -11]]), min=10.0), + ] + ) + + +def chain(order, resonant): + """q q~ > g e- e+ as a chain, the lepton pair from a Z if resonant.""" + if resonant: + propagators = [ms.Propagator(), ms.Propagator(mass=M_Z, width=W_Z, pdg_id=23)] + vertices = [["i0", "p0", "o0"], ["i1", "p0", "p1"], ["p1", "o1", "o2"]] + else: + propagators = [ms.Propagator(), ms.Propagator()] + vertices = [["i0", "p0", "o0"], ["p0", "p1", "o1"], ["i1", "p1", "o2"]] + diagram = ms.Diagram([0.0, 0.0], [0.0] * 3, propagators, vertices) + return ms.PhaseSpaceMapping( + ms.Topology(diagram), + E_CM, + t_channel_mode=ms.PhaseSpaceMapping.color_ordered, + cuts=cuts(), + permutations=[list(range(5))], + color_ordered=ms.ColorOrderedMapping.Options(order), + ) + + +RESONANT = chain([0, 2, 3, 1], True) +PLAIN = chain([0, 2, 3, 4, 1], False) + + +def sample(mapping, n, seed): + g = generator(seed) + inputs = [rand((n, mapping.random_dim()), generator=g)] + if mapping.discrete_dim(): + inputs.append( + randint(0, 2, (n, mapping.discrete_dim()), generator=g, dtype=torch.int32) + ) + p, x1, x2, det = mapping.map_forward(inputs) + return p, x1, x2, det * 2 ** mapping.discrete_dim() + + +def density(mapping, p, x1, x2): + return mapping.map_inverse([p, x1, x2])[-1] / 2 ** mapping.discrete_dim() + + +def test_resonant_chain_integrates_like_the_plain_one(): + def bw(p): + q = p[:, 3] + p[:, 4] + m2 = q[:, 0] ** 2 - torch.sum(q[:, 1:] ** 2, dim=1) + return 1e6 / ((m2 - M_Z**2) ** 2 + (M_Z * W_Z) ** 2) + + n = 400_000 + results = [] + for mapping in (RESONANT, PLAIN): + p, _, _, w = sample(mapping, n, 1) + values = torch.where(w > 0, w * torch.nan_to_num(bw(p)), 0.0) + results.append((values.mean(), values.std() / n**0.5)) + (res, res_err), (plain, plain_err) = results + assert res_err < 0.02 * res + assert abs(res - plain) < 3 * (res_err**2 + plain_err**2) ** 0.5 + + +def test_resonant_chain_inverts_the_plain_events(): + p, x1, x2, w = sample(PLAIN, 100_000, 2) + ok = w > 0 + out = RESONANT.map_inverse([p[ok], x1[ok], x2[ok]]) + r = out[0] + assert torch.all((r >= 0.0) & (r <= 1.0)) + p_back, *_ = RESONANT.map_forward([r]) + assert torch.max(torch.abs(p_back - p[ok])) < 1e-8 * E_CM + + +@pytest.mark.parametrize("flavor", [0, 1]) +def test_weights_are_the_normalised_masked_densities(flavor): + masks = [[1.0, 1.0], [0.0, 1.0]] + # diagrams 0 and 1 of flavor 0 on channels 0 and 1, diagram 2 of flavor 1 + weights = ms.MappingChannelWeights( + [RESONANT, PLAIN], masks, [0, 1, 1], [[1.0, 1.0, 0.0], [0.0, 0.0, 1.0]] + ) + runtime = ms.FunctionRuntime(weights.function(), default_context()) + p, x1, x2, w = sample(PLAIN, 40_000, 3) + ok = w > 0 + p, x1, x2 = p[ok], x1[ok], x2[ok] + n = p.shape[0] + flavors = torch.full((n,), flavor, dtype=torch.int32) + alpha, diagram = (torch.from_dlpack(t) for t in runtime.call([p, x1, x2, flavors])) + expected = torch.stack( + [density(RESONANT, p, x1, x2), density(PLAIN, p, x1, x2)], dim=1 + ) * torch.tensor(masks[flavor]) + expected /= expected.sum(dim=1, keepdim=True) + torch.testing.assert_close(alpha, expected, rtol=1e-12, atol=1e-14) + counts = torch.bincount(diagram.long(), minlength=3).double() + if flavor == 1: + assert counts.tolist() == [0.0, 0.0, n] + else: + assert counts[2] == 0 + mean = alpha[:, 0].sum() + assert abs(counts[0] - mean) < 5 * mean**0.5 + + +def test_channels_without_the_flavor_get_no_weight(): + weights = ms.MappingChannelWeights([PLAIN, PLAIN], [[0.0, 1.0]]) + runtime = ms.FunctionRuntime(weights.function(), default_context()) + p, x1, x2, w = sample(PLAIN, 1000, 4) + ok = w > 0 + n = int(ok.sum()) + alpha = torch.from_dlpack( + runtime.call( + [ + p[ok], + x1[ok], + x2[ok], + torch.zeros(n, dtype=torch.int32), + ] + )[0] + ) + torch.testing.assert_close(alpha[:, 1], torch.ones(n, dtype=alpha.dtype)) + torch.testing.assert_close(alpha[:, 0], torch.zeros(n, dtype=alpha.dtype)) + + +def test_uniform_weights_without_valid_densities(): + """If no inverse gives a valid density, the weights are uniform over the + channels of the flavor: a function of the event alone.""" + weights = ms.MappingChannelWeights([RESONANT, PLAIN, PLAIN], [[1.0, 0.0, 1.0]]) + runtime = ms.FunctionRuntime(weights.function(), default_context()) + p, x1, x2, w = sample(PLAIN, 1000, 5) + ok = w > 0 + n = int(ok.sum()) + flavors = torch.zeros(n, dtype=torch.int32) + p, x1, x2 = p[ok], x1[ok], x2[ok] + alpha = torch.from_dlpack(runtime.call([p, x1, x2, flavors])[0]) + assert not torch.allclose(alpha[:, 0], alpha[:, 2]) + # inverse mappings that fail numerically + p = torch.full_like(p, float("nan")) + alpha = torch.from_dlpack(runtime.call([p, x1, x2, flavors])[0]) + expected = torch.tensor([0.5, 0.0, 0.5], dtype=alpha.dtype).expand(n, 3) + torch.testing.assert_close(alpha, expected) diff --git a/madspace/tests/test_python_function.py b/madspace/tests/test_python_function.py index acc89cdaca..527a8b6989 100644 --- a/madspace/tests/test_python_function.py +++ b/madspace/tests/test_python_function.py @@ -126,6 +126,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", [ diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py new file mode 100644 index 0000000000..ddf7a1a44f --- /dev/null +++ b/madspace/tests/test_weighting_steps.py @@ -0,0 +1,745 @@ +"""Re- and unweighting steps after the first unweighting (WeightingStep). + +The integrand uses a deformed "surrogate" matrix element, and a later step +re- or unweights the events with the "true" one. The results have to agree with +a run that uses the true matrix element in the integrand, and must not depend +on the number of threads. Uses the e+ e- -> d d~ g setup of test_event_stream +with PythonFunction matrix elements. +""" + +import json +import os + +import numpy as np +import pytest +from test_event_stream import ( + CM_ENERGY, + PERMUTATIONS, + PIDS, + TOPOLOGY, + lhe_completer, + make_config, + make_context, +) +from test_integrand_me_outputs import python_matrix_element + +import madspace as ms + + +def true_me(momenta): + return 1.0 + 3.0 * (momenta[:, 2, 0] / CM_ENERGY) ** 2 + + +def surrogate_me(momenta): + return true_me(momenta) * (1.0 + 0.3 * np.sin(momenta[:, 3, 1] / 50.0)) + + +def make_integrand(alphas_file, matrix_element, pt_min=20.0): + mapping = ms.PhaseSpaceMapping( + TOPOLOGY, + CM_ENERGY, + permutations=PERMUTATIONS, + leptonic=True, + cuts=ms.Cuts( + [ + ms.CutItem(ms.Observable(PIDS, ms.Observable.obs_pt, [[i]]), min=pt_min) + for i in (2, 3, 4) + ] + ), + ) + diff_xs = ms.DifferentialCrossSection( + matrix_element=matrix_element, + cm_energy=CM_ENERGY, + running_coupling=None, + energy_scale=ms.CachedScale(), + return_matrix_element=True, + ) + return ms.Integrand( + mapping, + [diff_xs], + running_coupling=ms.RunningCoupling(ms.AlphaSGrid(alphas_file)), + energy_scale=ms.EnergyScale(len(PIDS), 91.188), + channel_indices=[0], + ) + + +def generate( + run_dir, + integrand_me, + step_modes=(), + step_me=true_me, + target_count=5000, + threads=2, + pt_mins=(20.0,), + grouped=(0,), + seed=7, + step_outputs=("matrix_element",), + override_outputs=(), + integrand_outputs=("matrix_element",), + step_channel=-1, + step_kwargs=None, + config_kwargs=None, + survey=True, +): + """Survey and generate, returns the status, the combined events and the + channels. step_me and step_outputs can be lists, one entry per step, and + integrand_me one entry per channel.""" + os.makedirs(run_dir, exist_ok=True) + context, alphas_file = make_context(run_dir, threads) + if not isinstance(integrand_me, list): + integrand_me = [integrand_me] * len(pt_mins) + me0s = [ + python_matrix_element(context, me, integrand_outputs) for me in integrand_me + ] + if not isinstance(step_me, list): + step_me = [step_me] * len(step_modes) + step_outputs = [step_outputs] * len(step_modes) + step_mes = [ + python_matrix_element(context, me, outputs, step_channel) + for me, outputs in zip(step_me, step_outputs) + ] + config = make_config() + config.target_count = target_count + config.cpu_batch_size = 500 + for key, value in (config_kwargs or {}).items(): + setattr(config, key, value) + channels = [ + ms.ChannelEventGenerator( + contexts=[context], + integrand=make_integrand(alphas_file, me0, pt_min), + event_file=os.path.join(run_dir, f"events.{i}.npy"), + weight_file=os.path.join(run_dir, f"weights.{i}.npy"), + config=config, + subprocess_index=0, + name=str(i), + histograms=None, + ) + for i, (pt_min, me0) in enumerate(zip(pt_mins, me0s)) + ] + steps = [ + ms.WeightingStep( + [me], + mode=mode, + override_outputs=list(override_outputs), + **(step_kwargs or {}), + ) + for me, mode in zip(step_mes, step_modes) + ] + groups = [ms.WeightingGroup(list(grouped), steps)] if steps else [] + generator = ms.EventGenerator( + [context], channels, seed=seed, config=config, groups=groups + ) + if survey: + generator.survey() + generator.generate() + output_file = os.path.join(run_dir, "combined.npy") + generator.combine_to_compact_npy(output_file) + return generator.status(), np.load(output_file), channels + + +@pytest.fixture(scope="module") +def reference(tmp_path_factory): + """The true matrix element in the integrand.""" + return generate( + str(tmp_path_factory.mktemp("reference")), true_me, target_count=20000 + ) + + +def check_compatible(status, events, reference): + ref_status, ref_events, _ = reference + error = np.hypot(status.error, ref_status.error) + assert abs(status.mean - ref_status.mean) < 4 * error + # the events are normalized to the cross section + assert np.mean(events["weight"]) == pytest.approx(status.mean_abs, rel=1e-9) + # mean energy of the first quark, weighted + energies = [ + np.sum(ev["weight"] * ev["part3_energy"]) / np.sum(ev["weight"]) + for ev in (events, ref_events) + ] + assert energies[0] == pytest.approx(energies[1], rel=0.01) + + +def test_unweight_step(tmp_path, reference): + status, events, channels = generate( + str(tmp_path), surrogate_me, ["unweight"], target_count=20000 + ) + assert len(events) == 20000 + check_compatible(status, events, reference) + # unweighted: apart from the overweights, all events have the same weight + weights = events["weight"] + assert np.mean(np.isclose(weights, np.min(weights))) > 0.98 + assert channels[0].step_max_weight(1) > 0 + + +def test_reweight_step(tmp_path, reference): + status, events, _ = generate( + str(tmp_path), surrogate_me, ["reweight"], target_count=20000 + ) + check_compatible(status, events, reference) + # the ratios of the matrix elements stay in the weights + weights = events["weight"] + assert np.std(weights) / np.mean(weights) > 0.1 + + +def test_reweight_then_unweight(tmp_path, reference): + """The unweighting step absorbs the ratio of the reweighting step: the + second step has the same matrix element, so it only unweights.""" + status, events, _ = generate( + str(tmp_path), surrogate_me, ["reweight", "unweight"], target_count=20000 + ) + check_compatible(status, events, reference) + weights = events["weight"] + assert np.mean(np.isclose(weights, np.min(weights))) > 0.98 + + +@pytest.mark.parametrize("step_modes", [["unweight"], ["reweight", "unweight"]]) +def test_reproducible(tmp_path, step_modes): + """Two grouped channels and one without steps, which share the batches + differently with different thread counts.""" + results = [ + generate( + str(tmp_path / str(threads)), + surrogate_me, + step_modes, + threads=threads, + pt_mins=(20.0, 40.0, 60.0), + grouped=(0, 2), + ) + for threads in (1, 4) + ] + (status_a, events_a, _), (status_b, events_b, _) = results + assert len(events_a) == 5000 + assert np.array_equal(events_a, events_b) + assert status_a.mean == status_b.mean + + +def heavy_tailed_me(momenta): + return surrogate_me(momenta) * ( + 1.0 + 0.05 / (0.002 + np.abs(np.sin(momenta[:, 3, 1] / 7.0))) + ) + + +def test_reproducible_step_freeze(tmp_path): + """The step maximum freezes while the first one still rises, in an order + relative to the step decisions that depends on the threads.""" + results = [ + generate( + str(tmp_path / str(threads)), + heavy_tailed_me, + ["unweight"], + threads=threads, + pt_mins=(20.0, 40.0, 60.0), + grouped=(0, 1, 2), + step_kwargs={"freeze_max_weight_after": 300}, + config_kwargs={ + "freeze_max_weight_after": 100000, + "max_overweight_truncation": 0.002, + }, + ) + for threads in (1, 2, 4, 8) + ] + status_a, events_a, channels_a = results[0] + for status_b, events_b, channels_b in results[1:]: + assert [c.step_max_weight(1) for c in channels_a] == [ + c.step_max_weight(1) for c in channels_b + ] + assert np.array_equal(events_a, events_b) + assert status_a.mean == status_b.mean + + +def test_survey_skips_cleared_iterations(tmp_path): + """Only the middle channel iterates again in the survey, so the steps of its + earlier iteration are skipped, independently of the threads.""" + results = [ + generate( + str(tmp_path / str(threads)), + [surrogate_me, heavy_tailed_me, surrogate_me], + ["reweight", "unweight"], + threads=threads, + pt_mins=(20.0, 40.0, 60.0), + grouped=(0, 1, 2), + config_kwargs={"survey_max_iters": 6, "survey_target_precision": 0.02}, + ) + for threads in (1, 4) + ] + (status_a, events_a, channels_a), (status_b, events_b, _) = results + assert [c.status().iterations for c in channels_a] == [3, 4, 3] + assert np.array_equal(events_a, events_b) + assert status_a.mean == status_b.mean + + +def test_step_budget_of_smaller_target(tmp_path): + """A smaller target applies its truncation budget to the step maxima too, + and rescales the unweighted count accordingly. EventGenerator does this at + the end of each round, as generate() starts with a generous target.""" + small = lambda momenta: surrogate_me(momenta) * (momenta[:, 3, 1] > 300.0) + _, _, channels = generate( + str(tmp_path), + [surrogate_me, small], + ["unweight"], + step_me=heavy_tailed_me, + pt_mins=(20.0, 20.0), + grouped=(0, 1), + target_count=20000, + survey=False, + step_kwargs={"freeze_max_weight_after": 10**6}, + ) + channel = channels[1] + target = channel.status().count_target + assert target < 5000 + # the budget of the final target is already applied + max_weight, count = channel.step_max_weight(1), channel.status().count_unweighted + channel.apply_step_truncation_budgets() + assert channel.step_max_weight(1) == max_weight + channel.set_target_count(target // 20) + channel.apply_step_truncation_budgets() + new_max_weight = channel.step_max_weight(1) + assert new_max_weight > max_weight + assert channel.status().count_unweighted == pytest.approx( + count * max_weight / new_max_weight, rel=1e-12 + ) + + +def broken_me(momenta, invalid=(np.nan, np.inf)): + """true_me, with invalid values in two regions.""" + px = momenta[:, 3, 1] + me = np.where(px > 100.0, invalid[0], true_me(momenta)) + return np.where(px < -100.0, invalid[1], me) + + +@pytest.mark.parametrize("step_modes", [["unweight"], ["reweight"]]) +def test_invalid_step_values(tmp_path, capfd, step_modes): + """Non-finite ratios count as weight 0, like a step matrix element of 0, + and are counted and warned about.""" + results = [ + generate( + str(tmp_path / name), + surrogate_me, + step_modes, + step_me=step_me, + pt_mins=(20.0, 40.0), + grouped=(0, 1), + ) + for name, step_me in [ + ("broken", broken_me), + ("zero", lambda momenta: broken_me(momenta, (0.0, 0.0))), + ] + ] + (status_broken, events_broken, channels), (status_zero, events_zero, _) = results + assert np.array_equal(events_broken, events_zero) + assert status_broken.mean == status_zero.mean + assert status_zero.count_invalid_step_values == 0 + invalid_counts = [c.status().count_invalid_step_values for c in channels] + assert all(count > 0 for count in invalid_counts) + assert status_broken.count_invalid_step_values == sum(invalid_counts) + if step_modes == ["unweight"]: + assert all(np.isfinite(c.step_max_weight(1)) for c in channels) + # after the survey and after generate() + warnings = [line for line in capfd.readouterr().out.splitlines() if "WARN" in line] + assert all("got the weight 0 in the weighting steps" in w for w in warnings) + assert {w.split("'")[1] for w in warnings} == {"0", "1"} + + +def test_index_outputs(tmp_path): + """A step fills in indices the integrand matrix element does not return, + and replaces the others only with override_outputs.""" + + def with_helicity(momenta): + return true_me(momenta), np.full(len(momenta), 1, np.int32) + + outputs = ("matrix_element", "helicity_index") + _, events, _ = generate( + str(tmp_path / "fill"), + surrogate_me, + ["unweight"], + step_me=with_helicity, + step_outputs=outputs, + ) + assert (events["helicity_index"] == 1).all() + assert (events["color_index"] == -1).all() + + def surrogate_with_helicity(momenta): + return surrogate_me(momenta), np.zeros(len(momenta), np.int32) + + for override, helicity in [((), 0), ((ms.MatrixElement.helicity_index_out,), 1)]: + _, events, _ = generate( + str(tmp_path / f"override{helicity}"), + surrogate_with_helicity, + ["unweight"], + step_me=with_helicity, + step_outputs=outputs, + integrand_outputs=outputs, + override_outputs=override, + ) + assert (events["helicity_index"] == helicity).all() + + +def true_over_surrogate(momenta): + return true_me(momenta) / surrogate_me(momenta) + + +@pytest.mark.parametrize("step_modes", [["unweight"], ["reweight"]]) +def test_ratio_step(tmp_path, step_modes): + """A step that returns |M_true|^2 / |M_surrogate|^2 gives the same events + as a step that returns |M_true|^2: the pipeline computes the same ratio.""" + results = [ + generate( + str(tmp_path / name), + surrogate_me, + step_modes, + step_me=step_me, + step_outputs=outputs, + ) + for name, step_me, outputs in [ + ("me", true_me, ("matrix_element",)), + ("ratio", true_over_surrogate, ("matrix_element_ratio",)), + ] + ] + (status_me, events_me, _), (status_ratio, events_ratio, _) = results + assert np.array_equal(events_me, events_ratio) + assert status_me.mean == status_ratio.mean + + +def test_ratio_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.""" + + def intermediate_me(momenta): + return true_me(momenta) * (1.0 + 0.2 * np.cos(momenta[:, 2, 2] / 40.0)) + + status, events, _ = generate( + str(tmp_path), + surrogate_me, + ["reweight", "unweight"], + step_me=[lambda p: intermediate_me(p) / surrogate_me(p), true_me], + step_outputs=[("matrix_element_ratio",), ("matrix_element",)], + target_count=20000, + ) + check_compatible(status, events, reference) + weights = events["weight"] + assert np.mean(np.isclose(weights, np.min(weights))) > 0.98 + + +def test_ratio_matrix_element_rules(tmp_path): + context, _ = make_context(str(tmp_path), 1) + ratio_me = python_matrix_element( + context, true_over_surrogate, ("matrix_element_ratio",) + ) + assert ms.WeightingStep([ratio_me]).returns_ratio() + assert not ms.WeightingStep( + [python_matrix_element(context, true_me)] + ).returns_ratio() + with pytest.raises(ValueError, match="weighting step"): + ms.DifferentialCrossSection( + matrix_element=ratio_me, + cm_energy=CM_ENERGY, + running_coupling=None, + energy_scale=ms.CachedScale(), + ) + with pytest.raises(ValueError, match="not both"): + ms.MatrixElement( + ratio_me.matrix_element_index(), + len(PIDS), + [ms.MatrixElement.momenta_in], + [ + ms.MatrixElement.matrix_element_out, + ms.MatrixElement.matrix_element_ratio_out, + ], + ) + groups = [ms.WeightingGroup([0], [ms.WeightingStep([ratio_me])])] + text = ms.weighting_groups_to_json(groups) + loaded = ms.weighting_groups_from_json(text) + assert loaded[0].steps[0].returns_ratio() + assert ms.weighting_groups_to_json(loaded) == text + # a fixed channel (AmpliCol colour ordering) survives the round trip + channel_me = python_matrix_element( + context, true_over_surrogate, ("matrix_element_ratio",), 3 + ) + text = ms.weighting_groups_to_json( + [ms.WeightingGroup([0], [ms.WeightingStep([channel_me])])] + ) + loaded = ms.weighting_groups_from_json(text) + assert loaded[0].steps[0].matrix_elements[0].channel() == 3 + + +def test_first_pass_settings_of_group(tmp_path): + """The first unweighting of the channels of a group uses the settings of + the group, the other channels the global ones.""" + context, alphas_file = make_context(str(tmp_path), 2) + surrogate = python_matrix_element(context, surrogate_me) + full = python_matrix_element(context, true_me) + config = make_config() + channels = [ + ms.ChannelEventGenerator( + [context], + make_integrand(alphas_file, surrogate), + str(tmp_path / f"events.{name}.npy"), + str(tmp_path / f"weights.{name}.npy"), + config, + 0, + name, + None, + ) + for name in ("group", "plain") + ] + group = ms.WeightingGroup( + [0], + [ms.WeightingStep([full], freeze_max_weight_after=300)], + first_pass_max_overweight_truncation=0.2, + first_pass_freeze_max_weight_after=2000, + ) + generator = ms.EventGenerator( + [context], channels, seed=3, config=config, groups=[group] + ) + assert [c.max_overweight_truncation() for c in channels] == [0.2, 0.05] + assert [c.freeze_max_weight_after() for c in channels] == [2000, 1000] + generator.survey() + # each channel runs until its own maximum freezes + generator.fix_max_weights() + assert [c.status().count_target for c in channels] == [2000, 1000] + + # stored with the group, old JSON without them still loads + text = ms.weighting_groups_to_json([group]) + loaded = ms.weighting_groups_from_json(text)[0] + assert loaded.first_pass_max_overweight_truncation == 0.2 + assert loaded.first_pass_freeze_max_weight_after == 2000 + old = json.loads(text) + del old[0]["first_pass_max_overweight_truncation"] + del old[0]["first_pass_freeze_max_weight_after"] + loaded = ms.weighting_groups_from_json(json.dumps(old))[0] + assert loaded.first_pass_max_overweight_truncation is None + assert loaded.first_pass_freeze_max_weight_after is None + assert "first_pass" not in ms.weighting_groups_to_json([loaded]) + + +def test_invalid_groups(tmp_path): + context, alphas_file = make_context(str(tmp_path), 1) + me = python_matrix_element(context, true_me) + config = make_config() + + def channel(return_matrix_element=True): + integrand = make_integrand(alphas_file, me) + if not return_matrix_element: + diff_xs = ms.DifferentialCrossSection( + matrix_element=me, + cm_energy=CM_ENERGY, + running_coupling=None, + energy_scale=ms.CachedScale(), + ) + integrand = ms.Integrand( + ms.PhaseSpaceMapping( + TOPOLOGY, CM_ENERGY, permutations=PERMUTATIONS, leptonic=True + ), + [diff_xs], + running_coupling=ms.RunningCoupling(ms.AlphaSGrid(alphas_file)), + energy_scale=ms.EnergyScale(len(PIDS), 91.188), + channel_indices=[0], + ) + return ms.ChannelEventGenerator( + [context], integrand, "", "", config, 0, "c", None + ) + + step = ms.WeightingStep([me]) + with pytest.raises(ValueError, match="matrix element"): + ms.EventGenerator( + [context], [channel(False)], seed=1, groups=[ms.WeightingGroup([0], [step])] + ) + with pytest.raises(ValueError, match="one weighting group"): + ms.EventGenerator( + [context], + [channel()], + seed=1, + groups=[ms.WeightingGroup([0], [step]), ms.WeightingGroup([0], [step])], + ) + with pytest.raises(ValueError, match="sampled internally"): + sampled = ms.MatrixElement( + me.matrix_element_index(), + len(PIDS), + [ms.MatrixElement.momenta_in, ms.MatrixElement.random_helicity_in], + [ms.MatrixElement.matrix_element_out], + 1, + True, + ) + ms.EventGenerator( + [context], + [channel()], + seed=1, + groups=[ms.WeightingGroup([0], [ms.WeightingStep([sampled])])], + ) + + +def full_me_with_indices(momenta): + zeros = np.zeros(len(momenta), np.int32) + return true_me(momenta), zeros, zeros, zeros + + +FULL_ME_OUTPUTS = ("matrix_element", "color_index", "helicity_index", "diagram_index") + + +def load_matrix_elements(context): + """The matrix elements, always in the same order, so that their indices + match the ones stored in the channels and groups.""" + surrogate = python_matrix_element(context, surrogate_me) + full = python_matrix_element(context, full_me_with_indices, FULL_ME_OUTPUTS) + return surrogate, full + + +@pytest.fixture(scope="module") +def fixed_gridpack(tmp_path_factory): + """Two grouped channels with all maxima fixed by fix_max_weights(), stored + like in a gridpack.""" + config = make_config() + run_dir = str(tmp_path_factory.mktemp("fixed")) + context, alphas_file = make_context(run_dir, 2) + surrogate, full = load_matrix_elements(context) + step = ms.WeightingStep( + [full], freeze_max_weight_after=1000, max_overweight_truncation=0.05 + ) + groups = [ms.WeightingGroup([0, 1], [step])] + channels = [ + ms.ChannelEventGenerator( + contexts=[context], + integrand=make_integrand(alphas_file, surrogate, pt_min), + event_file=os.path.join(run_dir, f"events.{i}.npy"), + weight_file=os.path.join(run_dir, f"weights.{i}.npy"), + config=config, + subprocess_index=0, + name=str(i), + histograms=None, + ) + for i, pt_min in enumerate((20.0, 60.0)) + ] + generator = ms.EventGenerator( + [context], channels, seed=1, config=config, groups=groups + ) + generator.survey() + generator.fix_max_weights() + globals_file = os.path.join(run_dir, "globals.tar") + context.save_globals(globals_file) + return { + "channels": [channel.to_json(True) for channel in channels], + "groups": ms.weighting_groups_to_json(groups), + "globals": globals_file, + "max_weights": [ + (channel.max_weight(), channel.step_max_weight(1)) for channel in channels + ], + "status": generator.status(), + } + + +def load_gridpack(gridpack, thread_count, run_dir=None): + context = ms.Context(ms.cpu_device(), thread_count) + context.load_globals(gridpack["globals"]) + load_matrix_elements(context) + config = make_config() + channels = [ + ms.ChannelEventGenerator.load_json( + channel_json, + [context], + event_file=os.path.join(run_dir, f"events.{i}.npy") if run_dir else "", + weight_file=os.path.join(run_dir, f"weights.{i}.npy") if run_dir else "", + config=config, + ) + for i, channel_json in enumerate(gridpack["channels"]) + ] + groups = ms.weighting_groups_from_json(gridpack["groups"]) + return context, channels, groups, config + + +def test_fixed_max_weights(fixed_gridpack, tmp_path): + """fix_max_weights() freezes the maxima of all steps. Reloaded like in a + gridpack, with the groups from JSON, the channels keep them.""" + max_weights = fixed_gridpack["max_weights"] + assert all(m0 > 0 and m1 > 0 for m0, m1 in max_weights) + for channel_json, (_, m1) in zip(fixed_gridpack["channels"], max_weights): + assert json.loads(channel_json)["step_max_weights"] == [m1] + + context, channels, groups, config = load_gridpack(fixed_gridpack, 2, str(tmp_path)) + assert ms.weighting_groups_to_json(groups) == fixed_gridpack["groups"] + generator = ms.EventGenerator( + [context], channels, seed=2, config=config, groups=groups + ) + generator.generate() + assert [ + (channel.max_weight(), channel.step_max_weight(1)) for channel in channels + ] == max_weights + status, new_status = fixed_gridpack["status"], generator.status() + assert new_status.count_unweighted == config.target_count + assert abs(status.mean - new_status.mean) < 4 * np.hypot( + status.error, new_status.error + ) + + +def make_stream(fixed_gridpack, thread_count=4, seed=5): + context, channels, groups, config = load_gridpack(fixed_gridpack, thread_count) + return ms.EventStream( + context, channels, seed, lhe_completer(), config, groups=groups + ) + + +def stream_batches(stream, batch_sizes): + batches = [stream.next_batch(size) for size in batch_sizes] + return {key: np.concatenate([b[key] for b in batches]) for key in batches[0]} + + +def test_event_stream_reproducible(fixed_gridpack): + reference = stream_batches(make_stream(fixed_gridpack, 4), [3000]) + for thread_count, batch_sizes in [(1, [1000, 2000]), (3, [1, 999, 2000])]: + events = stream_batches(make_stream(fixed_gridpack, thread_count), batch_sizes) + for key, values in reference.items(): + np.testing.assert_array_equal(events[key], values, err_msg=key) + + +def test_event_stream_cross_section(fixed_gridpack): + """The streamed events have unit weights apart from overweights, and the + integral after the step matches the generator's.""" + stream = make_stream(fixed_gridpack) + batch = stream.next_batch(5000) + weights = batch["weight"] + assert np.all(weights >= 1.0) + assert np.mean(weights == 1.0) > 0.9 + status, stream_status = fixed_gridpack["status"], stream.status() + assert abs(status.mean - stream_status.mean) < 4 * np.hypot( + status.error, stream_status.error + ) + + +def test_event_stream_status_without_inactive_channels(fixed_gridpack): + """A channel without fixed step maxima returns no events, so it does not + count in the cross section of the stream.""" + gridpack = dict(fixed_gridpack) + channel_json = json.loads(gridpack["channels"][1]) + del channel_json["step_max_weights"] + gridpack["channels"] = [gridpack["channels"][0], json.dumps(channel_json)] + stream = make_stream(gridpack) + integral = json.loads(gridpack["channels"][0])["integral"] + status = stream.status() + assert status.mean == pytest.approx(integral["mean"], rel=1e-12) + assert status.error == pytest.approx(integral["error"], rel=1e-12) + assert stream.channel_probabilities()[1] == 0.0 diff --git a/pyamplicol/docs/development/CONFIG_CONTRACT.md b/pyamplicol/docs/development/CONFIG_CONTRACT.md index c568afb0ee..58db22bbd3 100644 --- a/pyamplicol/docs/development/CONFIG_CONTRACT.md +++ b/pyamplicol/docs/development/CONFIG_CONTRACT.md @@ -39,6 +39,7 @@ overrides in order, then recorded license/resource clamping. - `selected_color_sector_ids: list[int] = []` - `selected_source_helicities: table[int, int] = {}` - `flavor_blind: bool = false` +- `flavor_blind_strict: bool = false` - `consider_charge_conjugation: bool = true` - `only_unit_multiplier: bool = false` - `analyze_process_groups: none | numerical = none` @@ -52,9 +53,10 @@ The explicit color-sector, reference-order, and source-helicity controls are developer-facing generation constraints. Ordinary users should leave them at their defaults and use the runtime selectors below. -The four process-group fields decide what `process-groups.json` records: +The five process-group fields decide what `process-groups.json` records: `flavor_blind` generates one representative per flavour-relabelling class and -records the others as unverified members; `analyze_process_groups = "numerical"` +records the others as unverified members; `flavor_blind_strict` implies it but +routes only members whose relabelling is a symmetry of the allowed vertices; `analyze_process_groups = "numerical"` merges generated processes whose matrix elements agree numerically up to a rational factor, restricted to factor 1 by `only_unit_multiplier`. `consider_charge_conjugation` also admits charge-conjugate relations in both diff --git a/pyamplicol/docs/user/command-line-interface.md b/pyamplicol/docs/user/command-line-interface.md index 4c1193950e..e708a37629 100644 --- a/pyamplicol/docs/user/command-line-interface.md +++ b/pyamplicol/docs/user/command-line-interface.md @@ -108,7 +108,9 @@ Useful generation options include: - `--name NAME` to name a request; - `--multiparticle 'p=d,d~,g'` to define a multiparticle label; - `--flavor-scheme N` and `--max-quark-lines N` to constrain expansion; -- `--flavor-blind` to generate one subprocess per flavour-relabelling class +- `--flavor-blind` to generate one subprocess per flavour-relabelling class, + or `--flavor-blind-strict` to route only the relations that are symmetries + of the allowed vertices (see [Group equivalent subprocesses](#group-equivalent-subprocesses)); - `--analyze-process-groups {none,numerical}`, `--only-unit-multiplier`, and `--[no-]consider-charge-conjugation` to search the generated subprocesses diff --git a/pyamplicol/docs/user/configuration.md b/pyamplicol/docs/user/configuration.md index 12ade6e57b..2c9d9a757e 100644 --- a/pyamplicol/docs/user/configuration.md +++ b/pyamplicol/docs/user/configuration.md @@ -166,12 +166,13 @@ assign arbitrary names fixed QCD or electroweak meanings. ### Process groups -Four `process` fields control which subprocesses are recorded as sharing one +Five `process` fields control which subprocesses are recorded as sharing one matrix element in the artifact's `process-groups.json`: ```toml [process] flavor_blind = false # generate one subprocess per flavour class +flavor_blind_strict = false # same, routing only vertex-symmetric relations consider_charge_conjugation = true # also relate charge-conjugate subprocesses only_unit_multiplier = false # accept numerical relations only with factor 1 analyze_process_groups = "none" # or "numerical", run after generation @@ -179,6 +180,7 @@ analyze_process_groups = "none" # or "numerical", run after generation `flavor_blind` is an assumption, not a check: members whose flavour relabelling is not a symmetry of the model couplings are reported with a warning. +`flavor_blind_strict` generates those members instead of routing them. `consider_charge_conjugation` matters only with `flavor_blind` or numerical analysis. See [Group equivalent subprocesses](command-line-interface.md#group-equivalent-subprocesses) diff --git a/pyamplicol/docs/user/process-grouping.md b/pyamplicol/docs/user/process-grouping.md index 89917daedd..7e5c2c189a 100644 --- a/pyamplicol/docs/user/process-grouping.md +++ b/pyamplicol/docs/user/process-grouping.md @@ -272,6 +272,20 @@ one fixed way, so it can warn about a relation that holds. In the same run, `d <-> s` alone moves the allowed `W` vertex `[d w+ u]`, although also exchanging `u <-> c` would restore it. +### Strict flavour-blind generation + +`--flavor-blind-strict` (`process.flavor_blind_strict = true`, which implies +`flavor_blind`) routes only the relations that pass the check above. Within a +flavour class, each member tries every flavour pairing, conjugated ones last, +against each representative kept so far, and joins the first one without a +warning. A member that matches none is generated and becomes a representative +of its own, which later members may join. No member of the written map has a +warning. For `p p > j j` with explicit electroweak orders, `u u~ > g g` is then +generated next to `d d~ > g g`, while `s c > s c` still joins `d u > d u`. + +The check is local and conservative, so strict generation may generate a +member whose relation holds; it never routes a member the check rejects. + ### When to use it Equal matrix elements within a class are **your responsibility**. The flag is diff --git a/pyamplicol/docs/user/umami-api.md b/pyamplicol/docs/user/umami-api.md index 1b6811c32a..79473c4b82 100644 --- a/pyamplicol/docs/user/umami-api.md +++ b/pyamplicol/docs/user/umami-api.md @@ -209,6 +209,23 @@ A provider's metadata contains: optional αs parameter. - `channels[].processes`: the channel/flavour selectors consumed by UMAMI, with physical PDGs and integration weights referenced by `matrix_elements`. +- `channels[].phasespace_channels`: the colour-ordered phase-space channels + integrating the selectors of a channel, each a cyclic `phasespace_order` of + the legs and the s-channel `resonances` (`pdg`, decay `legs`) it maps out. + A selector lists its phase-space channels in `multichannels`; they differ in + where its colour singlets attach to the colour lines (after the antiquark + ends, in every order and distribution, as in AmpliCol), in the non-cyclic + order of three or more lines, and in which compatible (nested or disjoint) + set of colour-singlet resonances of the process they follow, with a + non-resonant one always included. The candidates are every resonance whose + mass and width are not zero in the model: drop those that cannot go on shell + with the parameter-card values, and merge the channels that then agree. A + UMAMI channel collects the selectors linked by shared phase-space channels. + Weight the channels of a selector by a partition of unity, for example their + normalised mapping densities, rather than adding them. +- `channels[].colored_resonances`: the coloured resonances (`pdg`, decay + `legs`, `color`) of its selectors, which the phase-space channels do not + follow; reject the parameters if one of them can go on shell. - `pdg_ids`, `color_orders`, `color_flows` and `helicities`: shared physical configuration tables. Indices are zero-based and colour-flow references can be absent for contracted colour. @@ -219,8 +236,8 @@ A provider's metadata contains: Physical incoming PDGs remain oriented for PDF evaluation. Aliases name the same computation; they are not automatically extra integrable contributions. -Each exported contribution uses one phase-space ordering, rather than several -alternative maps whose naive sum would double count it. +An exported contribution integrated over several phase-space channels must be +weighted between them, since their naive sum would double count it. ## Automatic symmetry grouping diff --git a/pyamplicol/examples/all_options.toml b/pyamplicol/examples/all_options.toml index c0eb70f3a1..0d8c4dc69b 100644 --- a/pyamplicol/examples/all_options.toml +++ b/pyamplicol/examples/all_options.toml @@ -51,6 +51,10 @@ selected_source_helicities = { "1" = -1, "2" = 1 } # class and record the others as unverified members in process-groups.json. # Default false. flavor_blind = false +# process.flavor_blind_strict: Like flavor_blind, but route only the relations +# that are symmetries of the allowed model vertices and generate the members +# that would draw a warning. Implies flavor_blind. Default false. +flavor_blind_strict = false # process.consider_charge_conjugation: Also relate processes by charge # conjugation for flavor_blind and numerical grouping. Default true. consider_charge_conjugation = true @@ -102,6 +106,8 @@ mode = "error" workers = "auto" # generation.emit_api_bundle: Emit Python/Rust/C/C++/Fortran/UMAMI API bundle. Default true. emit_api_bundle = true +# generation.pack_payloads: Store per-process payloads and recurrence schedules in one payloads.pacbin. Default false. +pack_payloads = false [generation.validation] # generation.validation.enabled: Run generation validation. Default true. enabled = true diff --git a/pyamplicol/rust/crates/rusticol-core/src/artifact.rs b/pyamplicol/rust/crates/rusticol-core/src/artifact.rs index d7b6d61a29..2f773c1492 100644 --- a/pyamplicol/rust/crates/rusticol-core/src/artifact.rs +++ b/pyamplicol/rust/crates/rusticol-core/src/artifact.rs @@ -1,6 +1,6 @@ // SPDX-License-Identifier: 0BSD -use crate::pacbin::{PacbinMemberKind, PacbinReader}; +use crate::pacbin::{PacbinFileReader, PacbinMemberKind, PacbinReader}; use crate::{ ARTIFACT_MANIFEST_FILE, C_ABI_VERSION, COMPILED_MODEL_SCHEMA_VERSION, PROCESS_ARTIFACT_SCHEMA_VERSION, PYTHON_API_VERSION, RUNTIME_PHYSICS_SCHEMA_VERSION, @@ -40,6 +40,8 @@ const EVALUATOR_PAYLOAD_CONTAINER_KIND: &str = "pyamplicol-evaluator-payload-con const EVALUATOR_PAYLOAD_CONTAINER_STORAGE_ABI: &str = "pacbin-v1"; #[cfg(any(feature = "f64-compiled", feature = "f64-symjit"))] const RECURRENCE_BOOTSTRAP_FILENAME: &str = "recurrence-bootstrap.bin"; +/// Holds declared payloads that are not stored as loose files. +pub const PAYLOAD_CONTAINER_PATH: &str = "payloads.pacbin"; const SUPPORTED_ARTIFACT_TARGETS: [&str; 3] = [ "aarch64-apple-darwin", "x86_64-apple-darwin", @@ -209,6 +211,7 @@ pub enum PayloadRole { ApiBuildFile, SdkMetadata, ProcessGroups, + PayloadContainer, } #[derive(Clone, Debug, Decode, Deserialize, Eq, PartialEq, Serialize)] @@ -317,6 +320,7 @@ pub struct VerifiedArtifact { manifest: ArtifactManifest, payloads: Arc>, evaluator_payload_container: Option>, + payload_container: Option>, #[cfg(feature = "f64-compiled")] native_library_cache: Arc>>>, } @@ -363,6 +367,11 @@ pub(crate) enum EvaluatorPayloadSource { container: Arc, logical_path: String, }, + /// A declared payload stored in the artifact's payload container. + PackedDeclared { + container: Arc, + payload: Payload, + }, } /// A dynamic library loaded only from process-private snapshot bytes. @@ -633,6 +642,9 @@ impl EvaluatorPayloadSource { container, logical_path, } => container.member_bytes(logical_path).map(Cow::Borrowed), + Self::PackedDeclared { container, payload } => { + read_packed_payload(container, payload).map(Cow::Owned) + } } } @@ -643,6 +655,9 @@ impl EvaluatorPayloadSource { root.join(&payload.path).display().to_string() } Self::Packed { logical_path, .. } => format!("evaluators.pacbin:{logical_path}"), + Self::PackedDeclared { payload, .. } => { + format!("{PAYLOAD_CONTAINER_PATH}:{}", payload.path) + } } } } @@ -654,6 +669,7 @@ pub(crate) struct EvaluatorPayloadStore { relative_root: PathBuf, container: Option>, payloads: Option>>, + payload_container: Option>, #[cfg(feature = "f64-compiled")] native_library_cache: Arc>>>, } @@ -665,6 +681,7 @@ impl EvaluatorPayloadStore { relative_root: root.to_path_buf(), container: None, payloads: None, + payload_container: None, #[cfg(feature = "f64-compiled")] native_library_cache: Arc::new(Mutex::new(BTreeMap::new())), } @@ -698,6 +715,14 @@ impl EvaluatorPayloadStore { payload.role ))); } + if let Some(container) = &self.payload_container + && container.contains(&logical_path) + { + return Ok(EvaluatorPayloadSource::PackedDeclared { + container: container.clone(), + payload: payload.clone(), + }); + } return Ok(EvaluatorPayloadSource::AuthenticatedFile { root: self.artifact_root.clone(), payload: payload.clone(), @@ -914,6 +939,7 @@ impl VerifiedArtifact { &decoded.evaluator_payload_container, &payloads, )?; + let payload_container = load_bootstrap_payload_container(&root, &payloads)?; let artifact = Self { root, manifest_path, @@ -923,6 +949,7 @@ impl VerifiedArtifact { // leaves. Avoid opening and retaining the model-wide evaluator // PACBIN for source/intrinsic-only recurrence processes. evaluator_payload_container: None, + payload_container, #[cfg(feature = "f64-compiled")] native_library_cache: Arc::new(Mutex::new(BTreeMap::new())), }; @@ -1004,6 +1031,7 @@ impl VerifiedArtifact { validate_references(&manifest, &payloads)?; let evaluator_payload_container = load_evaluator_payload_container(&root, &manifest, &payloads)?; + let payload_container = load_payload_container(&root, &payloads)?; let payloads = Arc::new(payloads); Ok(Self { root, @@ -1011,6 +1039,7 @@ impl VerifiedArtifact { manifest, payloads, evaluator_payload_container, + payload_container, #[cfg(feature = "f64-compiled")] native_library_cache: Arc::new(Mutex::new(BTreeMap::new())), }) @@ -1050,7 +1079,7 @@ impl VerifiedArtifact { "artifact declares more than one process-groups payload", )); } - let bytes = read_declared_payload(&self.root, payload)?; + let bytes = self.read_declared(payload)?; if format!("{:x}", Sha256::digest(&bytes)) != payload.sha256 { return Err(RusticolError::integrity(format!( "process-groups payload {:?} does not match its declared sha256", @@ -1068,7 +1097,42 @@ impl VerifiedArtifact { pub fn read_payload(&self, path: &str) -> RusticolResult> { let payload = self.payload(path)?; - read_declared_payload(&self.root, payload) + self.read_declared(payload) + } + + fn read_declared(&self, payload: &Payload) -> RusticolResult> { + match &self.payload_container { + Some(container) if container.contains(&payload.path) => { + read_packed_payload(container, payload) + } + _ => read_declared_payload(&self.root, payload), + } + } + + fn is_packed(&self, path: &str) -> bool { + self.payload_container + .as_ref() + .is_some_and(|container| container.contains(path)) + } + + /// Open a declared payload that is itself a pacbin container, loose or packed, + /// and authenticate it against `expected_sha256`. + pub(crate) fn open_payload_pacbin( + &self, + path: &str, + expected_sha256: &[u8; 32], + ) -> RusticolResult { + if self.is_packed(path) { + let bytes = self.read_payload(path)?; + if <[u8; 32]>::from(Sha256::digest(&bytes)) != *expected_sha256 { + return Err(RusticolError::integrity(format!( + "pacbin container {path:?} payload digest mismatch" + ))); + } + return PacbinReader::from_bytes(bytes); + } + let file = self.open_payload_file(path)?; + PacbinReader::open_file_with_sha256(file, &self.root.join(path), expected_sha256) } /// Open one declared loose payload and pin the checked file description. @@ -1078,12 +1142,20 @@ impl VerifiedArtifact { /// the bytes between declaration checks and use. pub(crate) fn open_payload_file(&self, path: &str) -> RusticolResult { let payload = self.payload(path)?; + if self.is_packed(path) { + return Err(RusticolError::compatibility(format!( + "payload {path:?} is stored in {PAYLOAD_CONTAINER_PATH} and cannot be opened as a file" + ))); + } open_checked_payload(&self.root, payload) } + /// The location of a declared payload; packed payloads have no file there. pub(crate) fn payload_path(&self, path: &str) -> RusticolResult { let payload = self.payload(path)?; - drop(open_checked_payload(&self.root, payload)?); + if !self.is_packed(path) { + drop(open_checked_payload(&self.root, payload)?); + } Ok(self.root.join(path)) } @@ -1101,6 +1173,7 @@ impl VerifiedArtifact { relative_root: relative_root.to_path_buf(), container: self.evaluator_payload_container.clone(), payloads: Some(self.payloads.clone()), + payload_container: self.payload_container.clone(), #[cfg(feature = "f64-compiled")] native_library_cache: self.native_library_cache.clone(), }) @@ -1524,6 +1597,7 @@ fn payload_role_name(role: PayloadRole) -> &'static str { PayloadRole::ApiBuildFile => "api-build-file", PayloadRole::SdkMetadata => "sdk-metadata", PayloadRole::ProcessGroups => "process-groups", + PayloadRole::PayloadContainer => "payload-container", } } @@ -1764,18 +1838,118 @@ pub(crate) fn open_staged_recurrence_bootstrap_artifact( } } validate_evaluator_payload_container_declaration(&evaluator_payload_container, &payloads)?; + let payload_container = load_payload_container(&root, &payloads)?; Ok(VerifiedArtifact { manifest_path: root.join(ARTIFACT_MANIFEST_FILE), root, manifest, payloads: Arc::new(payloads), - // Process-ready lowering reads only loose plan/binding/color payloads. + // Process-ready lowering reads only plan/binding/color payloads. evaluator_payload_container: None, + payload_container, #[cfg(feature = "f64-compiled")] native_library_cache: Arc::new(Mutex::new(BTreeMap::new())), }) } +/// Open the payload container, checking that every member is a declared payload +/// with the declared size and digest. +fn load_payload_container( + root: &Path, + payloads: &BTreeMap, +) -> RusticolResult>> { + let mut declared = payloads + .values() + .filter(|payload| payload.role == PayloadRole::PayloadContainer); + let Some(payload) = declared.next() else { + return Ok(None); + }; + if declared.next().is_some() || payload.path != PAYLOAD_CONTAINER_PATH { + return Err(RusticolError::artifact(format!( + "an artifact declares at most one payload container, at {PAYLOAD_CONTAINER_PATH:?}" + ))); + } + open_payload_container(root, payload, payloads, true).map(Some) +} + +/// The payload container next to a recurrence bootstrap. Its inventory predates +/// the container, so only the members it declares are checked and used. +fn load_bootstrap_payload_container( + root: &Path, + payloads: &BTreeMap, +) -> RusticolResult>> { + let metadata = match fs::symlink_metadata(root.join(PAYLOAD_CONTAINER_PATH)) { + Ok(metadata) => metadata, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => { + return Err(RusticolError::security(format!( + "could not inspect {PAYLOAD_CONTAINER_PATH:?}: {error}" + ))); + } + }; + let container = Payload { + path: PAYLOAD_CONTAINER_PATH.to_string(), + role: PayloadRole::PayloadContainer, + media_type: "application/octet-stream".to_string(), + size_bytes: metadata.len(), + sha256: String::new(), + executable: false, + target: None, + process_id: None, + }; + open_payload_container(root, &container, payloads, false).map(Some) +} + +fn open_payload_container( + root: &Path, + payload: &Payload, + payloads: &BTreeMap, + require_declared: bool, +) -> RusticolResult> { + let file = open_checked_payload(root, payload)?; + let reader = PacbinFileReader::open_file(file, &root.join(&payload.path))?; + if reader.index().file_size() != payload.size_bytes { + return Err(RusticolError::integrity( + "payload container size disagrees with its declaration", + )); + } + for member in reader.index().members() { + let Some(packed) = payloads.get(member.logical_path()) else { + if require_declared { + return Err(RusticolError::integrity(format!( + "payload container member {:?} is not a declared payload", + member.logical_path() + ))); + } + continue; + }; + if member.kind() != PacbinMemberKind::DeclaredPayload + || packed.role == PayloadRole::PayloadContainer + || packed.executable + || member.length() != packed.size_bytes + || hex_digest(member.sha256()) != packed.sha256 + { + return Err(RusticolError::integrity(format!( + "invalid payload container member {:?}", + member.logical_path() + ))); + } + } + Ok(Arc::new(reader)) +} + +fn read_packed_payload(container: &PacbinFileReader, payload: &Payload) -> RusticolResult> { + validate_payload_declaration(payload)?; + let bytes = container.read_member(&payload.path)?; + if u64::try_from(bytes.len()).ok() != Some(payload.size_bytes) { + return Err(RusticolError::integrity(format!( + "packed payload {:?} disagrees with its declared size", + payload.path + ))); + } + Ok(bytes) +} + fn load_evaluator_payload_container( root: &Path, manifest: &ArtifactManifest, diff --git a/pyamplicol/rust/crates/rusticol-core/src/engine/recurrence_load.rs b/pyamplicol/rust/crates/rusticol-core/src/engine/recurrence_load.rs index 943c4722c7..2982579969 100644 --- a/pyamplicol/rust/crates/rusticol-core/src/engine/recurrence_load.rs +++ b/pyamplicol/rust/crates/rusticol-core/src/engine/recurrence_load.rs @@ -27,7 +27,7 @@ use super::recurrence_process_pack::{ decode_process_executor_descriptors, read_string, semantic_digest_from_bytes, }; use super::*; -use crate::pacbin::{PacbinMemberKind, PacbinReader}; +use crate::pacbin::PacbinMemberKind; #[cfg(feature = "on-the-fly-test-support")] use crate::recurrence::direct_backend::{ DirectCurrentObservation, begin_direct_current_observation, take_direct_current_observation, @@ -1248,7 +1248,6 @@ fn load_plan_reference( )?; validate_ready_plan_reference(&spec)?; let container = spec.plan.runtime_schedule; - let path = artifact.root().join(&container.path); let payload = artifact.payload(&container.path)?; if payload.role != PayloadRole::EvaluatorState || payload.media_type != "application/octet-stream" @@ -1265,8 +1264,7 @@ fn load_plan_reference( let process_binding = validate_process_binding(artifact, evaluator_root, &spec)?; let expected_file_sha = decode_sha256(&container.sha256)?; - let container_file = artifact.open_payload_file(&container.path)?; - let reader = PacbinReader::open_file_with_sha256(container_file, &path, &expected_file_sha)?; + let reader = artifact.open_payload_pacbin(&container.path, &expected_file_sha)?; let index = reader.index(); if index.file_size() != container.size_bytes || reader.container_size() as u64 != container.size_bytes diff --git a/pyamplicol/rust/crates/rusticol-core/src/pacbin.rs b/pyamplicol/rust/crates/rusticol-core/src/pacbin.rs index cf72999a47..7a7ebe7080 100644 --- a/pyamplicol/rust/crates/rusticol-core/src/pacbin.rs +++ b/pyamplicol/rust/crates/rusticol-core/src/pacbin.rs @@ -53,6 +53,8 @@ pub enum PacbinMemberKind { RecurrenceColorProjectionCertificate = 8, OnTheFlyProcessSeed = 9, ColorContraction = 10, + /// A payload declared by the artifact manifest, stored in its payload container. + DeclaredPayload = 11, } impl PacbinMemberKind { @@ -67,6 +69,7 @@ impl PacbinMemberKind { 8 => Ok(Self::RecurrenceColorProjectionCertificate), 9 => Ok(Self::OnTheFlyProcessSeed), 10 => Ok(Self::ColorContraction), + 11 => Ok(Self::DeclaredPayload), _ => Err(compatibility(format!( "unknown pacbin member kind: {value}" ))), @@ -590,6 +593,130 @@ impl PacbinReader { } } +/// A container read member by member from its open file. +/// +/// Only the frame and index are read when opening, so a large container costs +/// neither memory nor a full digest. Each member is copied out and checked +/// against its indexed SHA-256 before use. +#[derive(Debug)] +pub struct PacbinFileReader { + file: std::sync::Mutex, + display_path: PathBuf, + index: PacbinIndex, +} + +impl PacbinFileReader { + pub(crate) fn open_file(mut file: File, display_path: &Path) -> RusticolResult { + let file_size = file + .metadata() + .map_err(|error| { + RusticolError::artifact(format!( + "could not inspect pacbin container {}: {error}", + display_path.display() + )) + })? + .len(); + let minimum_size = (HEADER_SIZE + INDEX_HEADER_SIZE + FOOTER_SIZE) as u64; + if file_size < minimum_size { + return Err(integrity("truncated pacbin container")); + } + let header = read_file_range(&mut file, display_path, 0, HEADER_SIZE as u64)?; + let footer = read_file_range( + &mut file, + display_path, + file_size - FOOTER_SIZE as u64, + FOOTER_SIZE as u64, + )?; + let frame = parse_frame(&header, &footer, file_size)?; + let index_bytes = read_file_range( + &mut file, + display_path, + frame.index_offset, + frame.index_length, + )?; + let (members, index_sha256) = parse_index(&index_bytes, &frame)?; + // Padding is never handed out, so it is not read back here. + validate_canonical_payload_layout(&members, frame.index_offset, |_, _| Ok(()))?; + Ok(Self { + file: std::sync::Mutex::new(file), + display_path: display_path.to_path_buf(), + index: PacbinIndex { + version: frame.version, + index_offset: frame.index_offset, + index_sha256, + file_size, + members, + }, + }) + } + + pub fn index(&self) -> &PacbinIndex { + &self.index + } + + pub fn member(&self, logical_path: &str) -> RusticolResult<&PacbinMember> { + let normalized = normalize_logical_path(logical_path)?; + self.index + .members + .binary_search_by(|member| member.logical_path.as_bytes().cmp(normalized.as_bytes())) + .map(|index| &self.index.members[index]) + .map_err(|_| { + RusticolError::invalid_argument(format!("unknown pacbin member: {normalized}")) + }) + } + + pub fn contains(&self, logical_path: &str) -> bool { + self.member(logical_path).is_ok() + } + + /// Copy one member out of the container and authenticate the copy. + pub fn read_member(&self, logical_path: &str) -> RusticolResult> { + let member = self.member(logical_path)?; + let bytes = { + let mut file = self + .file + .lock() + .map_err(|_| RusticolError::artifact("pacbin container file handle is poisoned"))?; + read_file_range(&mut file, &self.display_path, member.offset, member.length)? + }; + let digest: [u8; 32] = Sha256::digest(&bytes).into(); + if digest != member.sha256 { + return Err(integrity(format!( + "pacbin member digest mismatch: {}", + member.logical_path + ))); + } + Ok(bytes) + } +} + +fn read_file_range( + file: &mut File, + display_path: &Path, + offset: u64, + length: u64, +) -> RusticolResult> { + let length = usize::try_from(length) + .map_err(|_| integrity("pacbin read length exceeds platform bounds"))?; + let mut bytes = Vec::new(); + bytes.try_reserve_exact(length).map_err(|error| { + RusticolError::artifact(format!( + "could not allocate {length} bytes for pacbin container {}: {error}", + display_path.display() + )) + })?; + bytes.resize(length, 0); + file.seek(SeekFrom::Start(offset)) + .and_then(|_| file.read_exact(&mut bytes)) + .map_err(|error| { + RusticolError::artifact(format!( + "could not read pacbin container {}: {error}", + display_path.display() + )) + })?; + Ok(bytes) +} + /// Atomically stream and publish one deterministic `pacbin-v1` container. /// /// Member metadata is retained in memory, while payload bytes are copied and @@ -1040,8 +1167,48 @@ fn parse_and_validate(bytes: &[u8], verify_payloads: bool) -> RusticolResult RusticolResult { if &header[0..8] != HEADER_MAGIC { return Err(integrity("invalid pacbin header magic")); } @@ -1066,17 +1233,13 @@ fn parse_and_validate(bytes: &[u8], verify_payloads: bool) -> RusticolResult= footer_offset_u64 { return Err(integrity("pacbin index offset is out of bounds")); } - let footer = checked_slice(bytes, footer_offset, FOOTER_SIZE, "pacbin footer")?; if &footer[0..8] != FOOTER_MAGIC { return Err(integrity("invalid pacbin footer magic")); } @@ -1104,17 +1267,23 @@ fn parse_and_validate(bytes: &[u8], verify_payloads: bool) -> RusticolResult RusticolResult<(Vec, [u8; 32])> { + let member_count = frame.member_count; let actual_index_digest: [u8; 32] = Sha256::digest(index_bytes).into(); let index_header = checked_slice(index_bytes, 0, INDEX_HEADER_SIZE, "pacbin index header")?; @@ -1228,21 +1397,11 @@ fn parse_and_validate(bytes: &[u8], verify_payloads: bool) -> RusticolResult Rusti } fn validate_canonical_payload_layout( - bytes: &[u8], members: &[PacbinMember], index_offset: u64, + mut validate_padding: impl FnMut(u64, u64) -> RusticolResult<()>, ) -> RusticolResult<()> { let mut expected_offset = HEADER_SIZE as u64; for member in members { @@ -1323,7 +1482,7 @@ fn validate_canonical_payload_layout( member.logical_path ))); } - validate_zero_region(bytes, unaligned_offset, padding, "pacbin payload padding")?; + validate_padding(unaligned_offset, padding)?; expected_offset = end; } @@ -1339,12 +1498,7 @@ fn validate_canonical_payload_layout( "pacbin payload region has a non-canonical trailing gap", )); } - validate_zero_region( - bytes, - expected_offset, - trailing_padding, - "pacbin payload padding", - )?; + validate_padding(expected_offset, trailing_padding)?; Ok(()) } diff --git a/pyamplicol/rust/crates/rusticol-core/src/pacbin_tests.rs b/pyamplicol/rust/crates/rusticol-core/src/pacbin_tests.rs index 4afba1b5c4..3f07319f20 100644 --- a/pyamplicol/rust/crates/rusticol-core/src/pacbin_tests.rs +++ b/pyamplicol/rust/crates/rusticol-core/src/pacbin_tests.rs @@ -198,6 +198,36 @@ fn open_keeps_one_immutable_snapshot_after_atomic_source_replacement() { fs::remove_file(path).unwrap(); } +#[test] +fn file_reader_authenticates_each_member_it_reads() { + static NEXT_PATH: AtomicU64 = AtomicU64::new(0); + let path = std::env::temp_dir().join(format!( + "rusticol-pacbin-file-reader-{}-{}.pacbin", + std::process::id(), + NEXT_PATH.fetch_add(1, Ordering::Relaxed) + )); + let mut bytes = decode_hex(PYTHON_GOLDEN_HEX); + fs::write(&path, &bytes).unwrap(); + let reader = PacbinFileReader::open_file(File::open(&path).unwrap(), &path).unwrap(); + assert_eq!(reader.index().members().len(), 2); + assert!(reader.contains("a/jit.symjit") && !reader.contains("a/missing")); + assert_eq!(reader.read_member("a/jit.symjit").unwrap(), b"jit-v1"); + assert_eq!(reader.read_member("z/exact.bin").unwrap(), b"exact-state"); + + let offset = reader.member("a/jit.symjit").unwrap().offset() as usize; + bytes[offset] ^= 1; + fs::write(&path, &bytes).unwrap(); + let error = reader.read_member("a/jit.symjit").unwrap_err(); + assert_eq!(error.kind(), RusticolErrorKind::Integrity); + assert!(error.to_string().contains("member digest mismatch")); + + let truncated = &decode_hex(PYTHON_GOLDEN_HEX)[..100]; + fs::write(&path, truncated).unwrap(); + let error = PacbinFileReader::open_file(File::open(&path).unwrap(), &path).unwrap_err(); + assert_eq!(error.kind(), RusticolErrorKind::Integrity); + fs::remove_file(path).unwrap(); +} + #[test] fn expected_container_digest_authenticates_an_immutable_snapshot() { static NEXT_PATH: AtomicU64 = AtomicU64::new(0); diff --git a/pyamplicol/src/pyamplicol/artifacts/inspection.py b/pyamplicol/src/pyamplicol/artifacts/inspection.py index 242348fb64..cc5a5e1a3c 100644 --- a/pyamplicol/src/pyamplicol/artifacts/inspection.py +++ b/pyamplicol/src/pyamplicol/artifacts/inspection.py @@ -30,7 +30,7 @@ from pyamplicol.api.errors import ArtifactError from pyamplicol.generation.process_grouping import GroupMember, ProcessGroupMap -from .manifest import MANIFEST_NAME, ArtifactManifest, load_manifest +from .manifest import MANIFEST_NAME, ArtifactManifest, ArtifactPayloads, load_manifest from .process_groups import initial_process_groups, read_process_groups from .security import confined_path, normalize_relative_path @@ -321,12 +321,30 @@ def _string(value: object, context: str) -> str: return value -def _json_mapping(path: Path, context: str) -> Mapping[str, object]: +# the payload reader of the last inspected manifest, which it keeps alive +_PAYLOAD_READER: list[tuple[ArtifactManifest, ArtifactPayloads]] = [] + + +def _payloads(manifest: ArtifactManifest) -> ArtifactPayloads: + if not _PAYLOAD_READER or _PAYLOAD_READER[0][0] is not manifest: + _PAYLOAD_READER[:] = [(manifest, ArtifactPayloads(manifest))] + return _PAYLOAD_READER[0][1] + + +def _read_payload(manifest: ArtifactManifest, path: Path) -> bytes: + """Read a payload at ``path`` below the root, loose or packed.""" + + return _payloads(manifest).read_bytes(path.relative_to(manifest.root).as_posix()) + + +def _json_mapping( + manifest: ArtifactManifest, path: Path, context: str +) -> Mapping[str, object]: try: - value = json.loads(path.read_text(encoding="utf-8")) + value = json.loads(_read_payload(manifest, path)) except OSError as exc: raise ArtifactError(f"cannot read {context}: {exc}") from exc - except json.JSONDecodeError as exc: + except (UnicodeDecodeError, json.JSONDecodeError) as exc: raise ArtifactError(f"invalid {context}: {exc}") from exc return _mapping(value, context) @@ -342,6 +360,8 @@ def _artifact_path( raise except (TypeError, ValueError) as exc: raise ArtifactError(f"{context} is not a valid artifact path") from exc + if _payloads(manifest).is_packed(normalized): + return manifest.root / normalized return confined_path(manifest.root, normalized) @@ -355,7 +375,9 @@ def _execution_paths(manifest: ArtifactManifest) -> Mapping[str, Path]: evaluator_relative, "runtime.evaluator_manifest_path", ) - evaluator_set = _json_mapping(evaluator_path, "runtime evaluator manifest") + evaluator_set = _json_mapping( + manifest, evaluator_path, "runtime evaluator manifest" + ) records = _sequence( evaluator_set.get("processes"), "runtime evaluator manifest.processes", @@ -504,7 +526,7 @@ def _table_rows( f"{context}.path", ) try: - payload = path.read_bytes() + payload = _read_payload(manifest, path) except OSError as exc: raise ArtifactError(f"cannot read {context} payload: {exc}") from exc expected_size = count * row_size @@ -905,7 +927,7 @@ def _eager_execution_inspection( ), "eager execution.kernel_pack.manifest_path", ) - pack = _json_mapping(pack_path, "prepared eager kernel pack") + pack = _json_mapping(manifest, pack_path, "prepared eager kernel pack") kernels = _sequence(pack.get("kernels"), "prepared eager kernel pack.kernels") kernel_ids = { _integer( @@ -1000,7 +1022,7 @@ def _eager_v3_execution_inspection( ), "eager execution.kernel_pack.manifest_path", ) - pack = _json_mapping(pack_path, "prepared eager kernel pack") + pack = _json_mapping(manifest, pack_path, "prepared eager kernel pack") kernels = _sequence(pack.get("kernels"), "prepared eager kernel pack.kernels") invocation_count = _integer( summary.get("invocation_count"), @@ -1202,7 +1224,7 @@ def _recurrence_execution_inspection( ), "recurrence execution.kernel_pack.manifest_path", ) - pack = _json_mapping(pack_path, "prepared recurrence kernel pack") + pack = _json_mapping(manifest, pack_path, "prepared recurrence kernel pack") kernels = _sequence( pack.get("kernels"), "prepared recurrence kernel pack.kernels", @@ -1949,7 +1971,9 @@ def _execution_inspection( selected_color_capability == ON_THE_FLY_CONTRACTED_COLOR_RUNTIME_CAPABILITY ) expected_color_accuracies = {"nlc", "full"} if contracted_color else {"lc"} - execution = _json_mapping(execution_path, "process execution manifest") + execution = _json_mapping( + manifest, execution_path, "process execution manifest" + ) expected_root_fields = { "schema_version", "kind", @@ -2268,7 +2292,9 @@ def _execution_inspection( "recurrence direct-arena execution is missing required " f"capabilities: {missing}" ) - execution = _json_mapping(execution_path, "process execution manifest") + execution = _json_mapping( + manifest, execution_path, "process execution manifest" + ) kind = _string(execution.get("kind"), "process execution manifest.kind") if kind != _RECURRENCE_RUNTIME_KIND: raise ArtifactError( @@ -2284,7 +2310,7 @@ def _execution_inspection( execution_mode="compiled", native_profile_phases=_COMPILED_PROFILE_PHASES, ) - execution = _json_mapping(execution_path, "process execution manifest") + execution = _json_mapping(manifest, execution_path, "process execution manifest") kind = _string(execution.get("kind"), "process execution manifest.kind") if kind == _EAGER_RUNTIME_KIND: return _eager_execution_inspection(manifest, execution, execution_path) @@ -2365,14 +2391,13 @@ def _physics_inspection( execution: _ExecutionInspection, ) -> _PhysicsInspection: relative = str(process["physics_path"]) - path = manifest.root / relative try: - payload = json.loads(path.read_text(encoding="utf-8")) + payload = json.loads(_payloads(manifest).read_bytes(relative)) except OSError as exc: raise ArtifactError( f"cannot read runtime physics metadata {relative}: {exc}" ) from exc - except json.JSONDecodeError as exc: + except (UnicodeDecodeError, json.JSONDecodeError) as exc: raise ArtifactError( f"invalid runtime physics metadata {relative}: {exc}" ) from exc diff --git a/pyamplicol/src/pyamplicol/artifacts/manifest.py b/pyamplicol/src/pyamplicol/artifacts/manifest.py index 6b430c636a..a64266f0f2 100644 --- a/pyamplicol/src/pyamplicol/artifacts/manifest.py +++ b/pyamplicol/src/pyamplicol/artifacts/manifest.py @@ -3,6 +3,7 @@ from __future__ import annotations +import hashlib import importlib import json import re @@ -23,6 +24,12 @@ verify_native_module, ) from pyamplicol.api.errors import ArtifactError, CompatibilityError +from pyamplicol.generation.evaluator_container import ( + PacbinError, + PacbinMember, + PacbinMemberKind, + PacbinReader, +) from .security import ( confined_path, @@ -37,6 +44,8 @@ ARTIFACT_IDENTITY_KIND = "pyamplicol-runtime-payload-identity" ARTIFACT_IDENTITY_SCHEMA_VERSION = 1 PORTABLE_64LE_TARGET = "portable-64le" +PAYLOAD_CONTAINER_PATH = "payloads.pacbin" +PAYLOAD_CONTAINER_ROLE = "payload-container" _SHA256 = re.compile(r"^[a-f0-9]{64}$") _GIT_REVISION = re.compile(r"^[a-f0-9]{40}$") _PUBLIC_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:+,~-]{0,254}$") @@ -70,6 +79,7 @@ "api-build-file", "sdk-metadata", "process-groups", + "payload-container", } _RUNTIME_IDENTITY_PAYLOAD_ROLES = frozenset( { @@ -905,6 +915,97 @@ def _validate_target_compatibility( ) +class ArtifactPayloads: + """Read declared payloads stored as files or packed in payloads.pacbin. + + The container index is checked against the payload declarations once; + each packed read opens its own stream and verifies the member digest. + """ + + def __init__(self, manifest: ArtifactManifest) -> None: + self.root = manifest.root + self._records = {record.path: record for record in manifest.payloads} + self.container_path: Path | None = None + self._members: Mapping[str, PacbinMember] = {} + containers = [ + record + for record in manifest.payloads + if record.role == PAYLOAD_CONTAINER_ROLE + ] + if not containers: + return + if len(containers) != 1 or containers[0].path != PAYLOAD_CONTAINER_PATH: + raise ArtifactError( + "an artifact declares at most one payload container, at " + f"{PAYLOAD_CONTAINER_PATH!r}" + ) + container = containers[0] + path = confined_path(self.root, container.path) + try: + with PacbinReader.open(path, verify_payloads=False) as reader: + index = reader.index + except (OSError, PacbinError) as exc: + raise ArtifactError( + f"could not open payload container {container.path}: {exc}" + ) from exc + if index.file_size != container.size_bytes: + raise ArtifactError("payload container size disagrees with its declaration") + for member in index.members: + record = self._records.get(member.logical_path) + if ( + record is None + or member.kind is not PacbinMemberKind.DECLARED_PAYLOAD + or record.role == PAYLOAD_CONTAINER_ROLE + or record.executable + or member.length != record.size_bytes + or member.sha256 != record.sha256 + ): + raise ArtifactError( + f"invalid payload container member {member.logical_path!r}" + ) + self.container_path = path + self._members = {member.logical_path: member for member in index.members} + + def is_packed(self, relative: str) -> bool: + """Whether ``relative`` is stored in the payload container.""" + + return normalize_relative_path(relative) in self._members + + def read_bytes(self, relative: str) -> bytes: + """The bytes of one payload, wherever it is stored.""" + + normalized = normalize_relative_path(relative) + member = self._members.get(normalized) + if member is None or self.container_path is None: + return confined_path(self.root, normalized).read_bytes() + with self.container_path.open("rb") as stream: + stream.seek(member.offset) + payload = stream.read(member.length) + if len(payload) != member.length: + raise ArtifactError(f"packed payload is truncated: {normalized}") + if hashlib.sha256(payload).hexdigest() != member.sha256: + raise ArtifactError(f"packed payload digest mismatch: {normalized}") + return payload + + def read_text(self, relative: str) -> str: + """The UTF-8 text of one payload, wherever it is stored.""" + + return self.read_bytes(relative).decode("utf-8") + + def verify_packed(self) -> None: + """Stream and verify the digest of every packed payload.""" + + if self.container_path is None: + return + try: + with PacbinReader.open(self.container_path, verify_payloads=True): + pass + except (OSError, PacbinError) as exc: + raise ArtifactError( + f"payload container verification failed: {exc}" + ) from exc + + def validate_payloads( manifest: ArtifactManifest, *, @@ -914,7 +1015,15 @@ def validate_payloads( raise ArtifactError("artifact manifest identity digest mismatch") _validate_target_compatibility(manifest, expected_target=expected_target) seen = {record.path for record in manifest.payloads} + payloads = ArtifactPayloads(manifest) + payloads.verify_packed() for record in manifest.payloads: + if payloads.is_packed(record.path): + if confined_path(manifest.root, record.path, must_exist=False).exists(): + raise ArtifactError( + f"packed artifact payload is also stored as a file: {record.path}" + ) + continue path = confined_path(manifest.root, record.path) if path.stat().st_size != record.size_bytes: raise ArtifactError(f"artifact payload size mismatch: {record.path}") @@ -999,9 +1108,12 @@ def load_manifest( "ARTIFACT_IDENTITY_KIND", "ARTIFACT_IDENTITY_SCHEMA_VERSION", "MANIFEST_NAME", + "PAYLOAD_CONTAINER_PATH", + "PAYLOAD_CONTAINER_ROLE", "PORTABLE_64LE_TARGET", "SCHEMA_VERSION", "ArtifactManifest", + "ArtifactPayloads", "PayloadRecord", "artifact_identity_extension", "canonical_manifest_bytes", diff --git a/pyamplicol/src/pyamplicol/artifacts/umami.py b/pyamplicol/src/pyamplicol/artifacts/umami.py index b609e38fbc..7300b9d854 100644 --- a/pyamplicol/src/pyamplicol/artifacts/umami.py +++ b/pyamplicol/src/pyamplicol/artifacts/umami.py @@ -13,10 +13,12 @@ from __future__ import annotations import json +import logging import math from collections.abc import Mapping, Sequence from contextlib import suppress from dataclasses import dataclass +from itertools import combinations_with_replacement, islice, permutations from typing import TYPE_CHECKING, Any from .api_bundle import ApiBundlePayload @@ -24,6 +26,8 @@ if TYPE_CHECKING: from ..generation.process_grouping import GroupMember +_LOGGER = logging.getLogger(__name__) + FLAVOR_BLIND_ASSUMPTION = ( "flavor-blind matrix elements: ungenerated process-group members are " "evaluated through their relabelled representative (assumed, not proven)" @@ -347,6 +351,159 @@ def _relabelled_runtime(runtime, relabel): } +# Phase-space partners of one colour order beyond this many are dropped: any +# subset gives correct channel weights, as each selector lists its own. +_MAX_PARTNERS = 64 +# Resonance configurations of one process beyond this many are not enumerated. +_MAX_CONFIGURATIONS = 1024 + + +def _roles(pdgs, incoming, particles): + """The colour representation of each leg in the all-outgoing convention.""" + roles = [] + for slot, pdg in enumerate(pdgs): + representation = int(particles[int(pdg)]["color"]) + if slot < incoming and abs(representation) == 3: + representation = -representation + roles.append(representation) + return roles + + +def _line_blocks(word_slots, roles): + """The complete open colour lines of an LC word in its order, or None for a + trace or a word without fundamental lines.""" + blocks, current = [], None + for slot in word_slots: + if roles[slot] == 3: + current = [slot] + elif current is not None: + current.append(slot) + if roles[slot] == -3: + blocks.append(current) + current = None + if not blocks or current is not None or sum(map(len, blocks)) != len(word_slots): + return None + return blocks + + +def _classify_resonances(resonances, roles): + """Split the candidates ``(pdg, legs, color)`` of a process into the colour + singlets a phase-space channel can follow, ``(pdg, legs)``, the coloured + ones and the singlets decaying to coloured and colourless particles.""" + mappable, coloured, mixed = [], [], [] + for pdg, legs, color in resonances: + if color != 1: + coloured.append((pdg, legs, color)) + elif len({roles[slot] == 1 for slot in legs}) > 1: + mixed.append((pdg, legs)) + else: + mappable.append((pdg, legs)) + return mappable, coloured, mixed + + +def _configurations(resonances, limit=_MAX_CONFIGURATIONS): + """Every set of resonances that can be mapped together, pairwise nested or + disjoint: the non-resonant one first, then the larger ones first. Returns + ``(configurations, complete)``, enumerating at most ``limit`` of them.""" + + def compatible(a, b): + return a < b or b < a or not a & b + + found, truncated = [], False + + def extend(chosen, start): + nonlocal truncated + found.append(tuple(chosen)) + for index in range(start, len(resonances)): + candidate = resonances[index] + if all(compatible(candidate[1], r[1]) for r in chosen): + if len(found) >= limit: + truncated = True + return + extend([*chosen, candidate], index + 1) + + extend([], 0) + found[1:] = sorted(found[1:], key=len, reverse=True) + return found, not truncated + + +def _resonance_legs(resonance, configuration): + """The legs of a resonance with its nested resonances kept contiguous.""" + inner = [ + other + for other in configuration + if other[1] < resonance[1] + and not any(other[1] < r[1] < resonance[1] for r in configuration) + ] + covered = set().union(*(other[1] for other in inner)) + blocks = [_resonance_legs(other, configuration) for other in inner] + blocks += [[slot] for slot in resonance[1] - covered] + return [slot for block in sorted(blocks, key=min) for slot in block] + + +def _placements(word_slots, blocks, units): + """Phase-space orders of one colour order, generated lazily: as in AmpliCol, + the colour singlet units follow the antiquark ends of the lines, in every + order and every distribution over the lines, and with three or more lines + every non-cyclic order of the lines is a partner too (the same colour + tensor).""" + if blocks is None: + yield list(word_slots) + [slot for unit in units for slot in unit] + return + for rest in permutations(blocks[1:]): + lines = [blocks[0], *rest] + for unit_order in permutations(units): + for cuts in combinations_with_replacement( + range(len(units) + 1), len(lines) - 1 + ): + bounds = (0, *cuts, len(units)) + order = [] + for index, line in enumerate(lines): + order += line + for unit in unit_order[bounds[index] : bounds[index + 1]]: + order += unit + yield order + + +def _phasespace_partners(word_slots, singlet_slots, roles, configurations): + """The phase-space orders and resonance configurations integrating one + colour order together, the first one its own: ``[(order, configuration)]`` + with ``configuration`` one of the process's ``configurations`` of + ``(pdg, legs)``. A resonance decaying to coloured particles applies only if + they are a complete colour line of this order. At most ``_MAX_PARTNERS`` + configurations are kept, in their order, and as many placements of each + as fit.""" + blocks = _line_blocks(word_slots, roles) if word_slots else None + lines = {frozenset(block) for block in blocks or ()} + applicable = [ + configuration + for configuration in configurations + if all( + legs in lines or all(roles[slot] == 1 for slot in legs) + for _, legs in configuration + ) + ][:_MAX_PARTNERS] + per_configuration = max(1, _MAX_PARTNERS // len(applicable)) + partners = [] + for configuration in applicable: + top = [ + r + for r in configuration + if not any(r[1] < other[1] for other in configuration) + ] + units = [ + _resonance_legs(r, configuration) + for r in top + if all(roles[slot] == 1 for slot in r[1]) + ] + covered = {slot for unit in units for slot in unit} + units += [[slot] for slot in singlet_slots if slot not in covered] + units.sort(key=min) + orders = islice(_placements(word_slots, blocks, units), per_configuration) + partners += [(order, configuration) for order in orders] + return partners + + class _InternTable(list): """A list with a hash index of its entries, so interning is O(1).""" @@ -442,6 +599,45 @@ def _add_terms(entry, terms): existing["factor"] += term["factor"] +def _group_phasespace_channels(data, selectors, phasespace_channels): + """One UMAMI channel per set of contributions linked by shared phase-space + channels: its contributions are integrated together, each over the + phase-space channels listed in its ``multichannels``. The coloured + resonances of its contributions are listed for the launch to check.""" + parent = {key: key for key in phasespace_channels} + + def root(key): + while parent[key] != key: + parent[key] = parent[parent[key]] + key = parent[key] + return key + + for _, keys, _ in selectors: + for key in keys[1:]: + parent[root(key)] = root(keys[0]) + groups: dict[tuple, dict] = {} + for entry, keys, colored in selectors: + group = groups.setdefault( + root(keys[0]), {"lookup": {}, "processes": [], "colored": []} + ) + entry["multichannels"] = [ + group["lookup"].setdefault(key, len(group["lookup"])) for key in keys + ] + group["processes"].append(entry) + group["colored"] += [r for r in colored if r not in group["colored"]] + data["channels"] = [] + for group in groups.values(): + channels = [phasespace_channels[key] for key in group["lookup"]] + data["channels"].append( + { + "phasespace_order": channels[0]["phasespace_order"], + "phasespace_channels": channels, + "colored_resonances": group["colored"], + "processes": group["processes"], + } + ) + + def _merge_flavour_orbits(data): """Merge different flavours with a common proven computation. @@ -450,9 +646,9 @@ def _merge_flavour_orbits(data): A merged flavour is evaluated by the kept selector, so the complete runtime evaluation (permutation, factor and colour tags included) must agree. """ - candidates = {} incoming = data["provider"]["incoming_count"] for channel in data["channels"]: + candidates = {} kept = [] for entry in channel["processes"]: runtime = entry["runtime"] @@ -464,6 +660,7 @@ def _merge_flavour_orbits(data): tuple(runtime["permutation"]), runtime["factor"], tuple(data["pdg_ids"][next(iter(pdgs))][:incoming]), + tuple(entry["multichannels"]), ) previous = next( ( @@ -481,9 +678,6 @@ def _merge_flavour_orbits(data): previous["members"].extend(entry["members"]) channel["processes"] = kept data["channels"] = [c for c in data["channels"] if c["processes"]] - for index, channel in enumerate(data["channels"]): - for entry in channel["processes"]: - entry["multichannels"] = [index] def _export_helicities(data, runtime_helicities): @@ -634,7 +828,7 @@ def _build_provider( assumptions.append(FLAVOR_BLIND_ASSUMPTION) data: dict[str, Any] = { "kind": "pyamplicol-umami-metadata", - "schema_version": 1, + "schema_version": 2, "provider": { "id": provider_id, "library": "libumami.so" if single else f"libumami_{provider_id}.so", @@ -658,8 +852,11 @@ def _build_provider( } structural_representatives: dict[str, dict] = {} orbit_entries: dict[tuple, tuple[dict, list[int]]] = {} - channel_lookup = {} channel_patterns: dict[tuple, list[int]] = {} + # (order, resonances) of every phase-space channel, and the exported + # contributions with the channels integrating them and coloured resonances + phasespace_channels: dict[tuple, dict] = {} + selectors: list[tuple[dict, list[tuple], list[dict]]] = [] physical_count = 0 runtime_helicities: list[list[list[int]]] = [] runtime_indices: dict[str, int] = {} @@ -726,6 +923,33 @@ def _build_provider( ) replay = _safe_replay(process) sectors_by_word = {_word(sector): sector for sector in process.color_sectors} + roles = _roles(pdgs, incoming, particles) + resonances = [ + ( + int(r["pdg"]), + frozenset(labels[int(label)] for label in r["labels"]), + int(r["color"]), + ) + for r in _records( + physics.get("extensions", {}).get("s_channel_resonances", ()) + ) + ] + if relation is not None: + resonances = [ + (pdg, frozenset(relation[slot] for slot in legs), color) + for pdg, legs, color in resonances + ] + mappable, coloured, unmapped = _classify_resonances(resonances, roles) + configurations, complete = _configurations(mappable) + if not complete and process.color_accuracy == "lc": + _LOGGER.warning( + "%s: only the first %d s-channel resonance configurations are " + "followed by the [amplicol] phase space", + source_id, + len(configurations), + ) + # the legs of a complete colour line in some colour order of the process + line_legs = set() colors = sorted( _records(physics.get("color_components")), key=lambda color: tuple(color.get("word", ())), @@ -761,6 +985,10 @@ def _build_provider( [int(external[slot]["label"]) for slot in word_slots] ) color_index = None if flow is None else _intern(data["color_flows"], flow) + if flow is not None: + coloured_slots = [slot for slot in order if roles[slot] != 1] + blocks = _line_blocks(coloured_slots, roles) if coloured_slots else None + line_legs.update(frozenset(block) for block in blocks or ()) rep_color, perm = replay.get(color_id, (color_id, tuple(identity))) runtime = { "process_index": process_index, @@ -860,17 +1088,91 @@ def _build_provider( "members": [member], "multichannels": [], } - # A single complete map per channel pattern. Maps are shared only if - # their fixed external masses agree (already a provider invariant). - if pattern not in channel_lookup: - channel_lookup[pattern] = len(data["channels"]) - data["channels"].append( - {"phasespace_order": channel_order, "processes": []} + # The phase-space channels of the contribution, in the relabelled + # legs. Maps are shared only if their fixed external masses agree + # (already a provider invariant) and they follow the same resonances. + partners = ( + [(order, ())] + if flow is None + else _phasespace_partners( + [slot for slot in order if roles[slot] != 1], + [slot for slot in order if roles[slot] == 1], + roles, + configurations, + ) + ) + # Contributions share a phase-space channel only with the same kind + # of particle in every leg, so that the cuts select the same legs + kinds = [None] * count + for slot, pdg in enumerate(pdgs): + particle = particles[int(pdg)] + kinds[relabel[slot]] = ( + abs(int(particle["color"])), + abs(float(particle.get("charge", 0.0))), + particle.get("spin"), + ) + keys = [] + for partner_order, configuration in partners: + _, partner_rotated = _channel_pattern( + [relabel[slot] for slot in partner_order], + incoming, + data["provider"]["masses"], + ) + partner_resonances = sorted( + (sorted(relabel[slot] for slot in legs), pdg) + for pdg, legs in configuration ) - channel_index = channel_lookup[pattern] - entry["multichannels"] = [channel_index] - data["channels"][channel_index]["processes"].append(entry) + key = ( + tuple(partner_rotated), + tuple((tuple(legs), abs(pdg)) for legs, pdg in partner_resonances), + tuple(kinds), + ) + if key not in keys: + keys.append(key) + phasespace_channels.setdefault( + key, + { + "phasespace_order": partner_rotated, + "resonances": [ + {"pdg": pdg, "legs": legs} + for legs, pdg in partner_resonances + ], + }, + ) + # not followed by the channels: the launch rejects those that go on shell + colored_resonances = ( + [] + if flow is None + else [ + { + "pdg": pdg, + "legs": sorted(relabel[slot] for slot in legs), + "color": color, + } + for pdg, legs, color in coloured + ] + ) + selectors.append((entry, keys, colored_resonances)) orbit_entries[orbit] = (entry, aligned_order, runtime, relabel) + unmapped += [ + (pdg, legs) + for pdg, legs in mappable + if legs not in line_legs and any(roles[slot] != 1 for slot in legs) + ] + if unmapped and process.color_accuracy == "lc": + _LOGGER.warning( + "%s: the [amplicol] phase space does not follow the s-channel " + "resonances %s, which decay to coloured and colourless particles " + "or not to one complete colour line", + source_id, + ", ".join( + f"{pdg} > {' '.join(str(pdgs[slot]) for slot in sorted(legs))}" + for pdg, legs in sorted( + unmapped, key=lambda r: (r[0], sorted(r[1])) + ) + ), + ) + _group_phasespace_channels(data, selectors, phasespace_channels) _merge_flavour_orbits(data) _export_helicities(data, runtime_helicities) alpha_names = {p["alpha_s_parameter"] for p in data["runtime_processes"]} diff --git a/pyamplicol/src/pyamplicol/artifacts/writer.py b/pyamplicol/src/pyamplicol/artifacts/writer.py index e7dc464711..26aa455d5a 100644 --- a/pyamplicol/src/pyamplicol/artifacts/writer.py +++ b/pyamplicol/src/pyamplicol/artifacts/writer.py @@ -3,11 +3,12 @@ from __future__ import annotations +import hashlib import io import json import os import tempfile -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from contextlib import suppress from datetime import UTC, datetime from pathlib import Path @@ -16,6 +17,8 @@ from .manifest import ( ARTIFACT_IDENTITY_EXTENSION, MANIFEST_NAME, + PAYLOAD_CONTAINER_PATH, + PAYLOAD_CONTAINER_ROLE, PayloadRecord, artifact_identity_extension, canonical_manifest_bytes, @@ -31,6 +34,9 @@ ) from .transaction import ArtifactTransaction, ArtifactWriteMode +# packable payloads are appended here while staging, then moved into the container +_PAYLOAD_SPOOL_NAME = ".payloads.spool" + class ArtifactBuilder: """Build one artifact in a private staging directory.""" @@ -41,13 +47,24 @@ def __init__( *, mode: ArtifactWriteMode = "error", expected_artifact_id: str | None = None, + packs: Callable[[str], bool] | None = None, ) -> None: + """``packs`` selects the payloads stored in payloads.pacbin; they are + spooled into one staging file instead of being written as files.""" if expected_artifact_id is not None and mode != "append": raise ValueError("expected_artifact_id is valid only for append mode") + if packs is not None and mode == "append": + raise ValueError("packed payloads cannot be appended") self._transaction = ArtifactTransaction(destination, mode=mode) self._expected_artifact_id = expected_artifact_id self.root: Path | None = None self._payloads: dict[str, PayloadRecord] = {} + self._packs = packs + self._spool: BinaryIO | None = None + # path -> (offset, length) in the spool + self._spooled: dict[str, tuple[int, int]] = {} + self._sealed = False + self._packed: frozenset[str] = frozenset() def __enter__(self) -> ArtifactBuilder: self.root = self._transaction.__enter__() @@ -114,13 +131,36 @@ def register_staged_file( if executable is not None: path.chmod(0o755 if executable else 0o644) fsync_file(path) - record = PayloadRecord( - path=path_value, + return self._record_payload( + path_value, role=role, media_type=media_type, size_bytes=path.stat().st_size, sha256=sha256_file(path), executable=executable_bit(path), + target=target, + process_id=process_id, + ) + + def _record_payload( + self, + path_value: str, + *, + role: str, + media_type: str, + size_bytes: int, + sha256: str, + executable: bool, + target: Mapping[str, object] | None, + process_id: str | None, + ) -> PayloadRecord: + record = PayloadRecord( + path=path_value, + role=role, + media_type=media_type, + size_bytes=size_bytes, + sha256=sha256, + executable=executable, target=dict(target) if target is not None else None, process_id=process_id, ) @@ -172,12 +212,25 @@ def add_stream( if not callable(read): raise TypeError("artifact stream source must provide read()") path_value = normalize_relative_path(relative) + if self._packs is not None and not executable and self._packs(path_value): + return self._spool_stream( + path_value, + source, + role=role, + media_type=media_type, + target=target, + process_id=process_id, + chunk_size=chunk_size, + ) path = self.staged_path(path_value, create_parent=True) descriptor, temporary_name = tempfile.mkstemp( prefix=f".{path.name}.", suffix=".tmp", dir=path.parent ) temporary = Path(temporary_name) descriptor_open = True + # hashed while written: the staged file is not read back + digest = hashlib.sha256() + size_bytes = 0 try: with os.fdopen(descriptor, "wb") as stream: descriptor_open = False @@ -192,6 +245,8 @@ def add_stream( "artifact stream source returned more bytes than requested" ) view = memoryview(chunk) + digest.update(view) + size_bytes += len(view) written = 0 while written < len(view): count = stream.write(view[written:]) @@ -207,10 +262,12 @@ def add_stream( os.close(descriptor) with suppress(FileNotFoundError): temporary.unlink() - return self.register_staged_file( + return self._record_payload( path_value, role=role, media_type=media_type, + size_bytes=size_bytes, + sha256=digest.hexdigest(), executable=executable, target=target, process_id=process_id, @@ -288,6 +345,9 @@ def discard_payloads(self, relative: str, *, recursive: bool = False) -> None: if path == path_value or (recursive and path.startswith(prefix)) ) for owned_path in owned: + if self._spooled.pop(owned_path, None) is not None: + self._payloads.pop(owned_path, None) + continue path = confined_path(self._root(), owned_path, must_exist=False) if path.is_symlink() or path.is_dir(): raise ValueError( @@ -312,6 +372,124 @@ def discard_payloads(self, relative: str, *, recursive: bool = False) -> None: with suppress(OSError): directory.rmdir() + def _spool_stream( + self, + path_value: str, + source: BinaryIO, + *, + role: str, + media_type: str, + target: Mapping[str, object] | None, + process_id: str | None, + chunk_size: int, + ) -> PayloadRecord: + if self._sealed: + raise RuntimeError( + f"payload {path_value!r} arrived after the payload container was sealed" + ) + if path_value == MANIFEST_NAME: + raise ValueError(f"{MANIFEST_NAME} is reserved for the artifact manifest") + if self._spool is None: + self._spool = (self._root() / _PAYLOAD_SPOOL_NAME).open("w+b") + spool = self._spool + offset = spool.seek(0, os.SEEK_END) + digest = hashlib.sha256() + while True: + chunk = source.read(chunk_size) + if not isinstance(chunk, bytes | bytearray | memoryview): + raise TypeError("artifact stream source must return bytes") + if not chunk: + break + digest.update(chunk) + spool.write(chunk) + length = spool.tell() - offset + self._spooled[path_value] = (offset, length) + return self._record_payload( + path_value, + role=role, + media_type=media_type, + size_bytes=length, + sha256=digest.hexdigest(), + executable=False, + target=target, + process_id=process_id, + ) + + def is_spooled(self, relative: str) -> bool: + """Whether a staged payload is held in the spool rather than a file.""" + + return normalize_relative_path(relative) in self._spooled + + def has_payload(self, relative: str) -> bool: + """Whether ``relative`` is a payload staged so far.""" + + return normalize_relative_path(relative) in self._payloads + + def read_staged(self, relative: str) -> bytes: + """The bytes of a staged payload, wherever it is stored.""" + + path_value = normalize_relative_path(relative) + spooled = self._spooled.get(path_value) + if spooled is not None and self._spool is not None: + offset, length = spooled + self._spool.flush() + data = os.pread(self._spool.fileno(), length, offset) + if len(data) != length: + raise OSError(f"short read of spooled payload {path_value!r}") + return data + if self._sealed and path_value in self._packed: + from pyamplicol.generation.evaluator_container import PacbinReader + + record = self._payloads[path_value] + with PacbinReader.open( + self._root() / PAYLOAD_CONTAINER_PATH, verify_payloads=False + ) as reader: + return reader.read_member(path_value, length=record.size_bytes) + return confined_path(self._root(), path_value).read_bytes() + + def seal_payload_container(self) -> PayloadRecord | None: + """Move the spooled payloads into payloads.pacbin; later packable + payloads are rejected. Their records are kept.""" + + if self._sealed: + raise RuntimeError("the payload container is already sealed") + self._sealed = True + self._packed = frozenset(self._spooled) + if self._spool is None: + return None + from pyamplicol.generation.evaluator_container import ( + PacbinMemberKind, + PacbinMemberSource, + write_pacbin_atomic, + ) + + spool = self._spool + spool.flush() + sources = [ + PacbinMemberSource( + path, + PacbinMemberKind.DECLARED_PAYLOAD, + _SpoolSlice(spool.fileno(), offset, length), + ) + for path, (offset, length) in self._spooled.items() + ] + index = write_pacbin_atomic(self.staged_path(PAYLOAD_CONTAINER_PATH), sources) + for member in index.members: + record = self._payloads[member.logical_path] + if member.length != record.size_bytes or member.sha256 != record.sha256: + raise ValueError( + f"packed payload {member.logical_path!r} changed while packing" + ) + spool.close() + self._spool = None + self._spooled = {} + (self._root() / _PAYLOAD_SPOOL_NAME).unlink() + return self.register_staged_file( + PAYLOAD_CONTAINER_PATH, + role=PAYLOAD_CONTAINER_ROLE, + media_type="application/octet-stream", + ) + def payload_records(self) -> tuple[PayloadRecord, ...]: """Return the immutable, path-ordered staged payload inventory. @@ -338,8 +516,11 @@ def finalize( default_process_id: str | None = None, extensions: Mapping[str, object] | None = None, ) -> Path: + if not self._sealed and self._packs is not None: + self.seal_payload_container() if not self._payloads: raise ValueError("an artifact must contain at least one payload") + self._remove_empty_directories() manifest_extensions = dict(extensions or {}) supplied_identity = manifest_extensions.get(ARTIFACT_IDENTITY_EXTENSION) expected_identity = artifact_identity_extension() @@ -378,10 +559,43 @@ def finalize( fsync_file(path) return path + def _remove_empty_directories(self) -> None: + """Remove the directories left empty by payloads moved into containers.""" + + root = self._root() + for directory, _, _ in os.walk(root, topdown=False): + if Path(directory) != root: + with suppress(OSError): + os.rmdir(directory) + def __exit__(self, *exception: object) -> bool: + if self._spool is not None: + self._spool.close() + self._spool = None if exception[0] is None and not (self._root() / MANIFEST_NAME).is_file(): raise RuntimeError("artifact builder exited without finalize()") return self._transaction.__exit__(*exception) # type: ignore[arg-type] -__all__ = ["ArtifactBuilder"] +class _SpoolSlice: + """One spooled payload, read with positioned reads from the open spool.""" + + def __init__(self, descriptor: int, offset: int, length: int) -> None: + self._descriptor = descriptor + self._position = offset + self._end = offset + length + + def read(self, size: int = -1) -> bytes: + remaining = self._end - self._position + if size < 0 or size > remaining: + size = remaining + if size == 0: + return b"" + data = os.pread(self._descriptor, size, self._position) + if not data: + raise OSError("short read of a spooled payload") + self._position += len(data) + return data + + +__all__ = ["PAYLOAD_CONTAINER_PATH", "PAYLOAD_CONTAINER_ROLE", "ArtifactBuilder"] diff --git a/pyamplicol/src/pyamplicol/assets/api_templates/umami/umami.h b/pyamplicol/src/pyamplicol/assets/api_templates/umami/umami.h deleted file mode 100644 index 4eb2dba380..0000000000 --- a/pyamplicol/src/pyamplicol/assets/api_templates/umami/umami.h +++ /dev/null @@ -1,324 +0,0 @@ -/* - * _ - * (_) - * _ _ _ __ ___ __ _ _ __ ___ _ - * | | | | '_ ` _ \ / _` | '_ ` _ \| | - * | |_| | | | | | | (_| | | | | | | | - * \__,_|_| |_| |_|\__,_|_| |_| |_|_| - * - * Unified MAtrix eleMent Interface - * - * - */ - -#ifndef UMAMI_HEADER -#define UMAMI_HEADER 1 - -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif - -/** - * Major version number of the UMAMI interface. If the major version is the same - * between caller and implementation, binary compatibility is ensured. - */ -#define UMAMI_MAJOR_VERSION 1 -/** - * Minor version number of the UMAMI interface. Between minor versions, new keys for - * errors, devices, metadata, inputs and outputs can be added. - */ -#define UMAMI_MINOR_VERSION 2 - -typedef enum { - /** operation was executed successfully */ - UMAMI_SUCCESS, - /** an unspecified error */ - UMAMI_ERROR, - /** operation not implemented */ - UMAMI_ERROR_NOT_IMPLEMENTED, - /** the provided input key is not supported by this matrix element implementation */ - UMAMI_ERROR_UNSUPPORTED_INPUT, - /** the provided output key is not supported by this matrix element implementation - */ - UMAMI_ERROR_UNSUPPORTED_OUTPUT, - /** the meta_key was not yet initialized */ - UMAMI_ERROR_UNINITIALIZED_META, - /** the provided metadata key is not supported by this matrix element implementation - */ - UMAMI_ERROR_UNSUPPORTED_META, - /** a mandatory matrix element input was not provided */ - UMAMI_ERROR_MISSING_INPUT, -} UmamiStatus; - -typedef enum { - UMAMI_DEVICE_CPU, - UMAMI_DEVICE_CUDA, - UMAMI_DEVICE_HIP, -} UmamiDevice; - -typedef enum { - /** `UmamiDevice` specifying the type of device */ - UMAMI_META_DEVICE, - /** `int` specifying the number of external particles */ - UMAMI_META_PARTICLE_COUNT, - /** `int` specifying the number of Feynman diagrams */ - UMAMI_META_DIAGRAM_COUNT, - /** `int` specifying the number of helicities */ - UMAMI_META_HELICITY_COUNT, - /** `int` specifying the number of colors */ - UMAMI_META_COLOR_COUNT, - /** `double` array of length particle count: the external-leg masses (GeV) */ - UMAMI_META_MASSES, -} UmamiMetaKey; - -typedef enum { - /** momenta of the external legs, type: `double`, shape: `(particle count, 4)` */ - UMAMI_IN_MOMENTA, - /** value for the strong coupling, type: `double`, shape: `()` */ - UMAMI_IN_ALPHA_S, - /** flavor index, type: `int`, shape: `()` */ - UMAMI_IN_FLAVOR_INDEX, - /** random number for color selection, type: `double`, shape: `()` */ - UMAMI_IN_RANDOM_COLOR, - /** random number for helicity selection, type: `double`, shape: `()` */ - UMAMI_IN_RANDOM_HELICITY, - /** random number for diagram selection, type: `double`, shape: `()` */ - UMAMI_IN_RANDOM_DIAGRAM, - /** externally selected helicity index, type: `int`, shape: `()` */ - UMAMI_IN_HELICITY_INDEX, - /** externally selected diagram index, type: `unsigned int`, shape: `()` */ - UMAMI_IN_DIAGRAM_INDEX, - /** externally selected channel index, type: `unsigned int`, shape: `()` */ - UMAMI_IN_CHANNEL_INDEX, - /** number of external legs, type: `int`, shape: `()` */ - UMAMI_IN_PARTICLE_COUNT, - /** PDG ids of the external legs, type: `int`, shape: `(particle count)` */ - UMAMI_IN_PDG_IDS, - /** momentum fraction of an incoming parton, type: `double`, shape: `()` */ - UMAMI_IN_MOMENTUM_FRACTION, - /** energy scale, e.g. the factorization scale, type: `double`, shape: `()` */ - UMAMI_IN_SCALE, - /** PDG id of a single parton, type: `int`, shape: `()` */ - UMAMI_IN_PDG_ID, - /** index of the beam (0 or 1), type: `int`, shape: `()` */ - UMAMI_IN_BEAM_INDEX, -} UmamiInputKey; - -/** Number of values in `UmamiInputKey` */ -#define UMAMI_INPUT_KEY_COUNT (UMAMI_IN_BEAM_INDEX + 1) - -typedef enum { - /** value of the matrix element, type: `double`, shape: `()` */ - UMAMI_OUT_MATRIX_ELEMENT, - /** squared amplitudes of the diagrams, type: `double`, shape: `(diagram count)` */ - UMAMI_OUT_DIAGRAM_AMP2, - /** selected color index, type: `int`, shape: `()` */ - UMAMI_OUT_COLOR_INDEX, - /** selected helicity index, type: `int`, shape: `()` */ - UMAMI_OUT_HELICITY_INDEX, - /** selected diagram index, type: `int`, shape: `()` */ - UMAMI_OUT_DIAGRAM_INDEX, - /** CUDA or HIP stream for asynchronous execution. Listed as an output as it is a - * mutable pointer */ - UMAMI_OUT_GPU_STREAM, - /** value of a custom cut, 0 rejects the event, type: `double`, shape: `()` */ - UMAMI_OUT_CUT, - /** value of a custom energy scale, type: `double`, shape: `()` */ - UMAMI_OUT_SCALE, - /** value x f(x, Q) of a custom parton density, type: `double`, shape: `()` */ - UMAMI_OUT_PDF, - /** value of a custom event bias, type: `double`, shape: `()` */ - UMAMI_OUT_BIAS, - /** ratio of the matrix element to a reference matrix element defined by the - * caller, type: `double`, shape: `()` */ - UMAMI_OUT_MATRIX_ELEMENT_RATIO, -} UmamiOutputKey; - -/** Number of values in `UmamiOutputKey` */ -#define UMAMI_OUTPUT_KEY_COUNT (UMAMI_OUT_MATRIX_ELEMENT_RATIO + 1) - -/** Implementation-defined pointer to a matrix element instance */ -typedef void* UmamiHandle; - -/** - * Retrieve metadata about the implemented matrix element - * - * @param meta_key - * key specifying the type of metadata to be retrieved - * @param result - * pointer to store the result. It's type depends on the metadata key - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_get_meta(UmamiMetaKey meta_key, void* result); - -/** - * Reports which input keys can be passed to this matrix element implementation, to - * be written by whoever implements the UMAMI interface. Implementing this function - * is optional; a caller that fails to resolve it should treat it as unavailable. - * - * @param supported - * pointer set on return to an implementation-owned array of booleans, where - * element `i` indicates whether the input key with numeric value `i` - * (see `UmamiInputKey`) is accepted by this implementation. - * @param count - * pointer set on return to the number of elements in `*supported`, reflecting - * the number of `UmamiInputKey` values known to this implementation. A caller - * built against a newer header should treat any key index beyond `*count` as - * not supported. - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_supported_inputs(bool const** supported, int* count); - -/** - * Reports which input keys must be passed to this matrix element implementation for - * `umami_matrix_element` to succeed, to be written by whoever implements the UMAMI - * interface. This is a subset of the keys reported by `umami_supported_inputs`. - * Implementing this function is optional; a caller that fails to resolve it should - * treat it as unavailable. - * - * @param required - * pointer set on return to an implementation-owned array of booleans, where - * element `i` indicates whether the input key with numeric value `i` - * (see `UmamiInputKey`) is mandatory for this implementation. - * @param count - * pointer set on return to the number of elements in `*required`. A caller - * built against a newer header should treat any key index beyond `*count` as - * not required. - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_required_inputs(bool const** required, int* count); - -/** - * Reports which output keys can be requested from this matrix element - * implementation, to be written by whoever implements the UMAMI interface. - * Implementing this function is optional; a caller that fails to resolve it should - * treat it as unavailable. - * - * @param supported - * pointer set on return to an implementation-owned array of booleans, where - * element `i` indicates whether the output key with numeric value `i` - * (see `UmamiOutputKey`) can be produced by this implementation. - * @param count - * pointer set on return to the number of elements in `*supported`. A caller - * built against a newer header should treat any key index beyond `*count` as - * not supported. - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_supported_outputs(bool const** supported, int* count); - -/** - * Creates an instance of the matrix element. Each instance is independent, so thread - * safety can be achieved by creating a separate one for every thread. - * - * @param handle - * pointer to an instance of the subprocess. Has to be cleaned up by - * the caller with `free_subprocess`. - * @param param_card_path - * path to the parameter file - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_initialize(UmamiHandle* handle, char const* param_card_path); - -/** - * Sets the value of a model parameter - * - * @param handle - * handle of a matrix element instance - * @param name - * name of the parameter - * @param parameter_real - * real part of the parameter value - * @param parameter_imag - * imaginary part of the parameter value. Ignored for real valued parameters. - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_set_parameter( - UmamiHandle handle, char const* name, double parameter_real, double parameter_imag -); - -/** - * Retrieves the value of a model parameter - * - * @param handle - * handle of a matrix element instance - * @param name - * name of the parameter - * @param parameter_real - * pointer to double to return real part of the parameter value - * @param parameter_imag - * pointer to double to return imaginary part of the parameter value. Ignored - * for real-valued parameters (i.e. you may pass a null pointer) - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_get_parameter( - UmamiHandle handle, char const* name, double* parameter_real, double* parameter_imag -); - -/** - * Evaluates the matrix element as a function of the given inputs, filling the - * requested outputs. Unless otherwise specified, all inputs and outputs have a - * column-major memory layout and have a batch dimension that is contiguous in memory. - * - * @param handle - * handle of a matrix element instance - * @param count - * number of events to evaluate the matrix element for - * @param stride - * stride of the batch dimension of the input and output arrays to simplify - * parallel execution on CPUs, see memory layout - * @param offset - * offset of the event index - * @param input_count - * number of inputs to the matrix element - * @param input_keys - * pointer to an array of input keys, length `input_count` - * @param inputs - * pointer to an array of void pointers to the inputs. The type of the inputs - * depends on the input key - * @param output_count - * number of outputs to the matrix element - * @param output_keys - * pointer to an array of output keys, length `output_count` - * @param outputs - * pointer to an array of void pointers to the outputs. The type of the outputs - * depends on the output key. The caller is responsible for allocating memory for - * the outputs. - * @return - * UMAMI_SUCCESS on success, error code otherwise - */ -UmamiStatus umami_matrix_element( - UmamiHandle handle, - size_t count, - size_t stride, - size_t offset, - size_t input_count, - UmamiInputKey const* input_keys, - void const* const* inputs, - size_t output_count, - UmamiOutputKey const* output_keys, - void* const* outputs -); - -/** - * Frees matrix element instance - * - * @param handle - * handle of a matrix element instance - */ -UmamiStatus umami_free(UmamiHandle handle); - -#ifdef __cplusplus -} -#endif - -#endif // UMAMI_HEADER diff --git a/pyamplicol/src/pyamplicol/assets/api_templates/umami/umami.h b/pyamplicol/src/pyamplicol/assets/api_templates/umami/umami.h new file mode 120000 index 0000000000..dfa21b9ab2 --- /dev/null +++ b/pyamplicol/src/pyamplicol/assets/api_templates/umami/umami.h @@ -0,0 +1 @@ +../../../../../../madspace/include/madspace/umami.h \ No newline at end of file diff --git a/pyamplicol/src/pyamplicol/cli/parser.py b/pyamplicol/src/pyamplicol/cli/parser.py index 7727e07801..14931856b1 100644 --- a/pyamplicol/src/pyamplicol/cli/parser.py +++ b/pyamplicol/src/pyamplicol/cli/parser.py @@ -274,6 +274,16 @@ def _add_process_grouping_options( "record the others as its group members (assumed, not verified)" ), ) + parser.add_argument( + "--flavor-blind-strict", + dest="process.flavor_blind_strict", + action="store_true", + default=argparse.SUPPRESS, + help=( + "like --flavor-blind, but route only the relations that are " + "symmetries of the allowed model vertices and generate the others" + ), + ) parser.add_argument( "--consider-charge-conjugation", dest="process.consider_charge_conjugation", @@ -364,6 +374,12 @@ def _add_generation_options(parser: argparse.ArgumentParser) -> None: action=argparse.BooleanOptionalAction, default=argparse.SUPPRESS, ) + parser.add_argument( + "--pack-payloads", + dest="generation.pack_payloads", + action=argparse.BooleanOptionalAction, + default=argparse.SUPPRESS, + ) parser.add_argument( "--validation", dest="generation.validation.enabled", diff --git a/pyamplicol/src/pyamplicol/config/models.py b/pyamplicol/src/pyamplicol/config/models.py index 40bc9719da..eb2e23aad9 100644 --- a/pyamplicol/src/pyamplicol/config/models.py +++ b/pyamplicol/src/pyamplicol/config/models.py @@ -345,6 +345,7 @@ class ProcessConfig: metadata=_setting("map_int", dynamic_kind="int"), ) flavor_blind: bool = field(default=False, metadata=_setting("bool")) + flavor_blind_strict: bool = field(default=False, metadata=_setting("bool")) consider_charge_conjugation: bool = field(default=True, metadata=_setting("bool")) only_unit_multiplier: bool = field(default=False, metadata=_setting("bool")) analyze_process_groups: ProcessGroupAnalysis = field( @@ -431,6 +432,7 @@ def __post_init__(self) -> None: ) for name in ( "flavor_blind", + "flavor_blind_strict", "consider_charge_conjugation", "only_unit_multiplier", ): @@ -655,6 +657,8 @@ class GenerationConfig: ) workers: AutoInt = field(default="auto", metadata=_setting("auto_int")) emit_api_bundle: bool = field(default=True, metadata=_setting("bool")) + # store per-process payloads and schedules in one payloads.pacbin + pack_payloads: bool = field(default=False, metadata=_setting("bool")) validation: GenerationValidationConfig = field( default_factory=GenerationValidationConfig, metadata=_section() ) @@ -673,6 +677,8 @@ def __post_init__(self) -> None: ) if not isinstance(self.emit_api_bundle, bool): raise ConfigurationError("generation.emit_api_bundle must be a boolean") + if not isinstance(self.pack_payloads, bool): + raise ConfigurationError("generation.pack_payloads must be a boolean") if not isinstance(self.validation, GenerationValidationConfig): raise ConfigurationError( "generation.validation must be a GenerationValidationConfig" diff --git a/pyamplicol/src/pyamplicol/generation/artifact_writer.py b/pyamplicol/src/pyamplicol/generation/artifact_writer.py index be38636aff..3179ebb6d6 100644 --- a/pyamplicol/src/pyamplicol/generation/artifact_writer.py +++ b/pyamplicol/src/pyamplicol/generation/artifact_writer.py @@ -6,6 +6,7 @@ import hashlib import importlib import importlib.metadata +import io import json import math import os @@ -33,7 +34,7 @@ process_groups_payload, validate_process_groups, ) -from pyamplicol.artifacts.security import sha256_file +from pyamplicol.artifacts.writer import PAYLOAD_CONTAINER_PATH, PAYLOAD_CONTAINER_ROLE from pyamplicol.config import ( ConfigClamp, ConfigResolution, @@ -118,7 +119,10 @@ from .structural_source_proof import ( ROLE as STRUCTURAL_SOURCE_PROOF_ROLE, ) -from .structural_source_proof import build_generation_structural_proof +from .structural_source_proof import ( + StructuralProofCache, + build_generation_structural_proof, +) from .validation import ValidationPointRecord, validation_point_map if TYPE_CHECKING: @@ -146,6 +150,14 @@ _EVALUATOR_PAYLOAD_CONTAINER_EXTENSION = "evaluator_payload_container" _RECURRENCE_BOOTSTRAP_FILENAME = "recurrence-bootstrap.bin" _EVALUATOR_PAYLOAD_CONTAINER_PATH = "evaluators.pacbin" +_PAYLOAD_CONTAINER_PATH = PAYLOAD_CONTAINER_PATH +_PAYLOAD_CONTAINER_ROLE = PAYLOAD_CONTAINER_ROLE +# trees whose payloads generation.pack_payloads moves into payloads.pacbin +_PACKED_PAYLOAD_PREFIXES = ( + "processes/", + "recurrence/schedules/", + "model/eager-kernels/", +) _COMPILED_COLOR_CONTRACTION_MEMBER_PATH = "compiled-color.pacrclr3" _EVALUATOR_PAYLOAD_CONTAINER_KIND = "pyamplicol-evaluator-payload-container" _EVALUATOR_PAYLOAD_CONTAINER_SCHEMA_VERSION = 1 @@ -601,8 +613,12 @@ def _register_staged_source( relative: str, kind: PacbinMemberKind, ) -> None: - path = self._builder.staged_path(relative) - self._new_sources[relative] = PacbinMemberSource(relative, kind, path) + source: Path | BinaryIO = ( + io.BytesIO(self._builder.read_staged(relative)) + if self._builder.is_spooled(relative) + else self._builder.staged_path(relative) + ) + self._new_sources[relative] = PacbinMemberSource(relative, kind, source) self._staged_loose_paths.add(relative) def _is_discarded(self, relative: str) -> bool: @@ -848,11 +864,7 @@ def _write_umami_bundle( physics = ( process.umami_physics if process is not None and process.umami_physics is not None - else json.loads( - builder.staged_path(str(record["physics_path"])).read_text( - encoding="utf-8" - ) - ) + else json.loads(builder.read_staged(str(record["physics_path"]))) ) umami_input = UmamiProcessInput( process_id=identifier, @@ -942,6 +954,11 @@ def write_schema_v3_artifact( existing = load_manifest(output) if mode == "append" else None if existing is not None: _reject_legacy_eager_append(existing) + if any(record.role == _PAYLOAD_CONTAINER_ROLE for record in existing.payloads): + raise ValueError( + "cannot append to an artifact with packed payloads; regenerate it " + "with mode='replace'" + ) hook = api_bundle_hook or _default_api_bundle_hook() requested_config = _config_payload(configuration.requested) effective_config = _config_payload(configuration.effective) @@ -1029,6 +1046,7 @@ def write_schema_v3_artifact( output, mode=write_mode, expected_artifact_id=(existing.artifact_id if existing is not None else None), + packs=_packs_payload if _payloads_packed(configuration.effective) else None, ) as builder: evaluator_payloads = _EvaluatorPayloadCollector( builder, @@ -1151,13 +1169,12 @@ def write_schema_v3_artifact( staged_root = builder.root if staged_root is None: # pragma: no cover - builder invariant raise RuntimeError("artifact builder is not active") + proof_cache = StructuralProofCache() for process_record in process_records: process_id = str(process_record["id"]) execution_path = f"processes/{process_id}/execution.json" - execution_file = builder.staged_path(execution_path) - execution_payload = json.loads( - execution_file.read_text(encoding="utf-8") - ) + execution_bytes = builder.read_staged(execution_path) + execution_payload = json.loads(execution_bytes.decode("utf-8")) if not isinstance(execution_payload, Mapping): raise ValueError( f"execution payload for {process_id!r} is not an object" @@ -1170,10 +1187,12 @@ def write_schema_v3_artifact( source_revision=source_revision, native_build_inputs_sha256=native_build_inputs_sha256, execution_path=execution_path, - execution_sha256=sha256_file(execution_file), + execution_sha256=hashlib.sha256(execution_bytes).hexdigest(), execution=execution_payload, evaluator_container_path=evaluator_container_path, evaluator_container_index_sha256=(evaluator_container_index_sha256), + cache=proof_cache, + payloads=builder, ) builder.add_json( f"processes/{process_id}/structural-source-proof.json", @@ -1182,6 +1201,9 @@ def write_schema_v3_artifact( process_id=process_id, compact=True, ) + if _payloads_packed(configuration.effective): + # the native recurrence bootstrap reads the packed payloads + builder.seal_payload_container() _write_recurrence_bootstraps( builder, processes=recurrence_processes, @@ -1264,9 +1286,13 @@ def write_schema_v3_artifact( ) manifest = load_manifest(output) - _validate_artifact_references(manifest) + packed = _validate_artifact_references(manifest) files = tuple( - [output / record.path for record in manifest.payloads] + [ + output / record.path + for record in manifest.payloads + if record.path not in packed + ] + [output / "artifact.json"] ) return ArtifactWriteResult( @@ -2143,8 +2169,8 @@ def _write_recurrence_bootstraps( separators=(",", ":"), allow_nan=False, ).encode("utf-8"), - builder.staged_path(execution_path).read_bytes(), - builder.staged_path(physics_path).read_bytes(), + builder.read_staged(execution_path), + builder.read_staged(physics_path), ) if not isinstance(image, bytes) or not image: raise RuntimeError( @@ -5361,8 +5387,11 @@ def _call_api_bundle_hook( return "API" -def _validate_artifact_references(manifest: ArtifactManifest) -> None: +def _validate_artifact_references(manifest: ArtifactManifest) -> frozenset[str]: + """Check the payload references and files; return the packed payload paths.""" + declared = {record.path for record in manifest.payloads} + packed = _validate_payload_container(manifest) _validate_evaluator_payload_container(manifest) _validate_recurrence_schedule_sharing(manifest) required = { @@ -5395,6 +5424,63 @@ def _validate_artifact_references(manifest: ArtifactManifest) -> None: raise ValueError( "artifact contains undeclared files: " + ", ".join(sorted(undeclared)) ) + if actual & packed: + raise ValueError( + "packed payloads are also stored as files: " + + ", ".join(sorted(actual & packed)) + ) + return packed + + +def _payloads_packed(config: GenerationConfig | RunConfig) -> bool: + generation = config if isinstance(config, GenerationConfig) else config.generation + return bool(generation.pack_payloads) + + +def _packs_payload(path: str) -> bool: + """Whether generation.pack_payloads stores ``path`` in payloads.pacbin.""" + + return ( + path.startswith(_PACKED_PAYLOAD_PREFIXES) + # nested containers below processes/ are opened as files by their loaders + and not (path.startswith("processes/") and path.endswith(".pacbin")) + # a loose bootstrap sidecar selects the fast recurrence loading path + and not path.endswith("/recurrence-bootstrap.bin") + ) + + +def _validate_payload_container(manifest: ArtifactManifest) -> frozenset[str]: + """Check that the payload container holds exactly declared payloads.""" + + containers = [ + record for record in manifest.payloads if record.role == _PAYLOAD_CONTAINER_ROLE + ] + if not containers: + return frozenset() + if len(containers) != 1 or containers[0].path != _PAYLOAD_CONTAINER_PATH: + raise ValueError( + f"an artifact declares at most one payload container, at " + f"{_PAYLOAD_CONTAINER_PATH!r}" + ) + records = {record.path: record for record in manifest.payloads} + with PacbinReader.open( + manifest.root / _PAYLOAD_CONTAINER_PATH, verify_payloads=False + ) as reader: + members = reader.members + for member in members: + record = records.get(member.logical_path) + if ( + record is None + or member.kind is not PacbinMemberKind.DECLARED_PAYLOAD + or record.role == _PAYLOAD_CONTAINER_ROLE + or record.executable + or member.length != record.size_bytes + or member.sha256 != record.sha256 + ): + raise ValueError( + f"invalid payload container member {member.logical_path!r}" + ) + return frozenset(member.logical_path for member in members) def _validate_recurrence_schedule_sharing(manifest: ArtifactManifest) -> None: diff --git a/pyamplicol/src/pyamplicol/generation/evaluator_container.py b/pyamplicol/src/pyamplicol/generation/evaluator_container.py index e8a5a0fffe..e708b304ce 100644 --- a/pyamplicol/src/pyamplicol/generation/evaluator_container.py +++ b/pyamplicol/src/pyamplicol/generation/evaluator_container.py @@ -66,6 +66,8 @@ class PacbinMemberKind(IntEnum): RECURRENCE_COLOR_PROJECTION_CERTIFICATE = 8 ON_THE_FLY_PROCESS_SEED = 9 COLOR_CONTRACTION = 10 + # a payload declared by the artifact manifest, stored in its payload container + DECLARED_PAYLOAD = 11 @dataclass(frozen=True, slots=True) diff --git a/pyamplicol/src/pyamplicol/generation/recurrence_schedule_sharing.py b/pyamplicol/src/pyamplicol/generation/recurrence_schedule_sharing.py index 5ea6aed790..dd82dfe8f1 100644 --- a/pyamplicol/src/pyamplicol/generation/recurrence_schedule_sharing.py +++ b/pyamplicol/src/pyamplicol/generation/recurrence_schedule_sharing.py @@ -650,21 +650,52 @@ def recurrence_schedule_semantic_digest( return _canonical_digest(payload) -def build_recurrence_process_executor_pack( - *, - direct_catalog: object, - kernel_pack: object, - required_executor_ids: Sequence[int], - runtime_layout_digest: str, -) -> RecurrenceProcessExecutorPack: - """Project one validated model catalog to its exact fixed-plan executor set.""" +@dataclass(frozen=True, slots=True) +class _ExecutorCatalog: + """The validated model-wide part of every process executor pack.""" + + templates: Mapping[int, Mapping[str, object]] + kernels: Mapping[int, Mapping[str, object]] + backend: str + target_triple: str + portable: bool + cpu_features: tuple[str, ...] + compiled_model_digest: str + recurrence_template_catalog_digest: str + prepared_kernel_pack_digest: str + direct_template_catalog_digest: str + + +# by object identity; the entries keep their keys alive so ids are not reused +_EXECUTOR_CATALOGS: dict[tuple[int, int], tuple[object, object, _ExecutorCatalog]] = {} +_EXECUTOR_CATALOGS_LOCK = Lock() +_EXECUTOR_CATALOGS_LIMIT = 4 + +def _executor_catalog(direct_catalog: object, kernel_pack: object) -> _ExecutorCatalog: + key = (id(direct_catalog), id(kernel_pack)) + with _EXECUTOR_CATALOGS_LOCK: + entry = _EXECUTOR_CATALOGS.get(key) + if entry is None: + entry = ( + direct_catalog, + kernel_pack, + _build_executor_catalog(direct_catalog, kernel_pack), + ) + if len(_EXECUTOR_CATALOGS) >= _EXECUTOR_CATALOGS_LIMIT: + _EXECUTOR_CATALOGS.pop(next(iter(_EXECUTOR_CATALOGS))) + _EXECUTOR_CATALOGS[key] = entry + return entry[2] + + +def _build_executor_catalog( + direct_catalog: object, kernel_pack: object +) -> _ExecutorCatalog: direct_mapping = _object_mapping(direct_catalog, "direct-template catalog") pack_mapping = _object_mapping(kernel_pack, "prepared-kernel pack") raw_templates = _object_sequence( direct_mapping.get("templates"), "direct-template catalog templates" ) - catalog_count = len(raw_templates) templates: dict[int, Mapping[str, object]] = {} for expected_id, raw_template in enumerate(raw_templates): template = _require_mapping(raw_template, "direct-template descriptor") @@ -677,19 +708,6 @@ def build_recurrence_process_executor_pack( ) templates[executor_id] = template - selected_ids = tuple( - _wire_u32(value, "required direct-executor ID") - for value in required_executor_ids - ) - if selected_ids != tuple(sorted(set(selected_ids))): - raise RecurrenceScheduleSharingError( - "required direct-executor IDs must be sorted and unique" - ) - if any(executor_id not in templates for executor_id in selected_ids): - raise RecurrenceScheduleSharingError( - "required direct-executor ID is absent from its complete catalog" - ) - raw_kernels = _object_sequence( pack_mapping.get("kernels"), "prepared-kernel pack kernels" ) @@ -731,16 +749,13 @@ def build_recurrence_process_executor_pack( raise RecurrenceScheduleSharingError( "target CPU features must be sorted and unique" ) - - descriptors = tuple( - _encode_process_executor_descriptor( - templates[executor_id], - backend=backend, - kernels=kernels, - ) - for executor_id in selected_ids - ) - return RecurrenceProcessExecutorPack( + return _ExecutorCatalog( + templates=templates, + kernels=kernels, + backend=backend, + target_triple=target_triple, + portable=portable, + cpu_features=cpu_features, compiled_model_digest=_wire_sha256( direct_mapping.get("compiled_model_digest"), "compiled-model digest" ), @@ -755,14 +770,54 @@ def build_recurrence_process_executor_pack( direct_template_catalog_digest=_wire_sha256( direct_mapping.get("catalog_digest"), "direct-template catalog digest" ), + ) + + +def build_recurrence_process_executor_pack( + *, + direct_catalog: object, + kernel_pack: object, + required_executor_ids: Sequence[int], + runtime_layout_digest: str, +) -> RecurrenceProcessExecutorPack: + """Project one validated model catalog to its exact fixed-plan executor set.""" + + catalog = _executor_catalog(direct_catalog, kernel_pack) + templates = catalog.templates + selected_ids = tuple( + _wire_u32(value, "required direct-executor ID") + for value in required_executor_ids + ) + if selected_ids != tuple(sorted(set(selected_ids))): + raise RecurrenceScheduleSharingError( + "required direct-executor IDs must be sorted and unique" + ) + if any(executor_id not in templates for executor_id in selected_ids): + raise RecurrenceScheduleSharingError( + "required direct-executor ID is absent from its complete catalog" + ) + + descriptors = tuple( + _encode_process_executor_descriptor( + templates[executor_id], + backend=catalog.backend, + kernels=catalog.kernels, + ) + for executor_id in selected_ids + ) + return RecurrenceProcessExecutorPack( + compiled_model_digest=catalog.compiled_model_digest, + recurrence_template_catalog_digest=catalog.recurrence_template_catalog_digest, + prepared_kernel_pack_digest=catalog.prepared_kernel_pack_digest, + direct_template_catalog_digest=catalog.direct_template_catalog_digest, runtime_layout_digest=_wire_sha256( runtime_layout_digest, "runtime-layout digest" ), - backend=backend, - target_triple=target_triple, - portable=portable, - cpu_features=cpu_features, - catalog_executor_count=catalog_count, + backend=catalog.backend, + target_triple=catalog.target_triple, + portable=catalog.portable, + cpu_features=catalog.cpu_features, + catalog_executor_count=len(templates), executor_ids=selected_ids, descriptor_payloads=descriptors, ) diff --git a/pyamplicol/src/pyamplicol/generation/s_channel_resonances.py b/pyamplicol/src/pyamplicol/generation/s_channel_resonances.py new file mode 100644 index 0000000000..7f52d380bf --- /dev/null +++ b/pyamplicol/src/pyamplicol/generation/s_channel_resonances.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: 0BSD +"""The unstable s-channel propagators a process may put on shell. + +A phase-space mapping has to follow these resonances, or it undersamples their +peaks. This is a species-level reachability sweep over the subsets of external +legs, like the DAG prepass: it ignores helicity and colour, keeps the coupling +orders, and may overestimate. Masses and widths are only known at launch, so +every candidate whose mass and width are not zero in the model is kept. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +from ..models.base import CouplingOrders, Model, Particle +from ..processes.ir import CanonicalProcessIR +from .dag_reachability import ( + _combine_coupling_order_tuples, + _coupling_orders_within_limits, + _masks_by_size, + _normalize_coupling_order_limits, + _ordered_splits, +) + + +def s_channel_resonances( + process: CanonicalProcessIR, + model: Model, + max_coupling_orders: Mapping[str, int] | None = None, +) -> list[dict[str, object]]: + """``[{"pdg", "labels", "color"}]`` of every possibly unstable particle (see + ``_possibly_unstable``) that a tree current of two or more final-state legs + can carry, with a complementary current to close the amplitude with. The + pdg carries the summed charge of the outgoing legs it decays to.""" + + if not any(_possibly_unstable(model, p) for p in model.particles.values()): + return [] + limits = _normalize_coupling_order_limits(max_coupling_orders) + legs = process.legs + every = final = 0 + species: dict[int, dict[int, set[CouplingOrders]]] = {} + for leg in legs: + bit = 1 << (leg.label - 1) + every |= bit + if leg.is_final: + final |= bit + if leg.outgoing_pdg is not None: + species.setdefault(bit, {}).setdefault(int(leg.outgoing_pdg), set()).add(()) + for mask in _masks_by_size(every): + if mask & (mask - 1) == 0 or mask == every: + continue + current = species.setdefault(mask, {}) + for left_mask, right_mask in _ordered_splits(mask): + left, right = species.get(left_mask), species.get(right_mask) + if not left or not right: + continue + for left_pdg, left_orders in left.items(): + for right_pdg, right_orders in right.items(): + for vertex in model.vertices_accepting( + left_pdg, right_pdg, color_accuracy=process.color_accuracy + ): + if model.skip_duplicate_vertex_orientation(vertex): + continue + vertex_orders = model.vertex_coupling_orders(vertex) + for a in left_orders: + for b in right_orders: + orders = _combine_coupling_order_tuples( + a, b, vertex_orders + ) + if _coupling_orders_within_limits(orders, limits): + current.setdefault(vertex.particles[2], set()).add( + orders + ) + found: dict[tuple[int, tuple[int, ...]], dict[str, object]] = {} + for mask, particles in species.items(): + if mask & ~final or mask.bit_count() < 2: + continue + decay_legs = [leg for leg in legs if mask >> (leg.label - 1) & 1] + rest = species.get(every & ~mask, {}) + for pdg, orders in particles.items(): + if not _possibly_unstable(model, model.particle(pdg)): + continue + closing = rest.get(model.anti_particle(pdg), ()) + if not any( + _coupling_orders_within_limits( + _combine_coupling_order_tuples(a, b), limits + ) + for a in orders + for b in closing + ): + continue + charge = sum(model.charge(int(leg.pdg)) for leg in decay_legs) + if abs(model.charge(pdg) - charge) > 1e-9: + pdg = model.anti_particle(pdg) + labels = tuple(leg.label for leg in decay_legs) + found.setdefault( + (pdg, labels), + { + "pdg": pdg, + "labels": list(labels), + "color": int(model.color_rep(pdg)), + }, + ) + return [found[key] for key in sorted(found)] + + +def _possibly_unstable(model: Model, particle: Particle) -> bool: + """Whether neither the mass nor the width of a particle is zero in the + model: a zero value without a parameter (ZERO in UFO) is final.""" + + def zero(kind: str) -> bool: + if getattr(particle, kind) != 0.0: + return False + provider = getattr(model, f"runtime_{kind}_parameter_name", None) + if not callable(provider): + return True + try: + return provider(particle.pdg) is None + except KeyError: + return False + + return not zero("mass") and not zero("width") diff --git a/pyamplicol/src/pyamplicol/generation/service.py b/pyamplicol/src/pyamplicol/generation/service.py index 55dc5f2729..65ce7ae73a 100644 --- a/pyamplicol/src/pyamplicol/generation/service.py +++ b/pyamplicol/src/pyamplicol/generation/service.py @@ -162,6 +162,7 @@ flavor_blind_classes, flavor_blind_mapping, particle_table, + profile_mappings, ) from .progress import GenerationPhaseReporter, PhaseHandle from .recurrence_color import ( @@ -227,6 +228,7 @@ build_runtime_expression_schema, build_runtime_schema_layout, ) +from .s_channel_resonances import s_channel_resonances from .stage_compiler import ( build_and_write_generic_stage_evaluator_artifacts, write_model_parameter_evaluator_artifact, @@ -4488,6 +4490,7 @@ def _expand_process_set( table, consider_charge_conjugation=self._consider_charge_conjugation, relation_warning=self._flavor_blind_relation_warning(resolved, table), + strict=self._flavor_blind_strict, ) return tuple(result) @@ -6145,6 +6148,9 @@ def report_selector(details: Mapping[str, object]) -> None: ), color_plan=prepared.complete_color_plan, ) + physics["extensions"]["s_channel_resonances"] = s_channel_resonances( + expanded.process_ir, model, self._coupling_order_limits or None + ) profile["physics-metadata"] = time.perf_counter() - physics_started run = self._run_config assert run is not None @@ -6742,6 +6748,9 @@ def _construct_eager_plan_v3_artifact( runtime_model_parameter_records_from_eager_input(lowering_input) ), ) + physics["extensions"]["s_channel_resonances"] = s_channel_resonances( + ir, model, self._coupling_order_limits or None + ) return EagerPlanV3ProcessArtifact( process_id=process_name, expression=ir.process, @@ -7561,7 +7570,14 @@ def _coupling_order_limits(self) -> dict[str, int]: @property def _flavor_blind(self) -> bool: run = self._run_config - return run is not None and run.process.flavor_blind + return run is not None and ( + run.process.flavor_blind or run.process.flavor_blind_strict + ) + + @property + def _flavor_blind_strict(self) -> bool: + run = self._run_config + return run is not None and run.process.flavor_blind_strict @property def _consider_charge_conjugation(self) -> bool: @@ -8655,42 +8671,86 @@ def _grouping_process(entry: _ExpandedProcess) -> GroupingProcess: ) +def _flavor_blind_relation( + process: GroupingProcess, + representative: GroupingProcess, + table: Mapping[int, GroupingParticle], + *, + consider_charge_conjugation: bool, + relation_warning: _RelationWarning, + representative_ir: CanonicalProcessIR, + strict: bool, +) -> tuple[LegMapping, str | None]: + """The assumed mapping and its warning; ``strict`` tries every flavour + pairing, conjugated ones last, for one that draws no warning.""" + + assumed = flavor_blind_mapping( + process, + representative, + table, + consider_charge_conjugation=consider_charge_conjugation, + ) + warning = relation_warning(process, representative, assumed, representative_ir) + if warning is None or not strict: + return assumed, warning + candidates = profile_mappings(process, representative, table) + if consider_charge_conjugation: + candidates += profile_mappings(process, representative, table, conjugated=True) + for candidate in candidates: + mapping = replace(candidate, method="flavor-blind") + if ( + mapping != assumed + and relation_warning(process, representative, mapping, representative_ir) + is None + ): + return mapping, None + return assumed, warning + + def _collapse_flavor_blind_classes( entries: Sequence[_ExpandedProcess], table: Mapping[int, GroupingParticle], *, consider_charge_conjugation: bool, relation_warning: _RelationWarning, + strict: bool = False, ) -> list[_ExpandedProcess]: """Keep the representative of each flavour class and record the others on it. - Processes with explicit aliases keep their own artifact entry. + Processes with explicit aliases keep their own artifact entry. With + ``strict``, a member joins the first representative of its class whose + relation draws no warning, and becomes a generated representative otherwise. """ by_name = {entry.request.name: entry for entry in entries} members: dict[str, list[GroupMember]] = {} warnings: list[str] = [] + split = 0 for flavour_class in flavor_blind_classes( [_grouping_process(entry) for entry in entries], table, consider_charge_conjugation=consider_charge_conjugation, ): - representative = flavour_class[0] + representatives = [flavour_class[0]] for process in flavour_class[1:]: if by_name[process.process_id].aliases: continue - mapping = flavor_blind_mapping( - process, - representative, - table, - consider_charge_conjugation=consider_charge_conjugation, - ) - warning = relation_warning( - process, - representative, - mapping, - by_name[representative.process_id].process_ir, - ) + for representative in representatives: + mapping, warning = _flavor_blind_relation( + process, + representative, + table, + consider_charge_conjugation=consider_charge_conjugation, + relation_warning=relation_warning, + representative_ir=by_name[representative.process_id].process_ir, + strict=strict, + ) + if warning is None or not strict: + break + else: + representatives.append(process) + split += 1 + continue if warning is not None: warnings.append(warning) members.setdefault(representative.process_id, []).append( @@ -8727,6 +8787,13 @@ def _collapse_flavor_blind_classes( len(warnings), "" if len(warnings) == 1 else "s", ) + if strict and split: + _LOGGER.info( + "Strict flavor-blind grouping: %d subprocess%s generated because " + "their relabelling is not a symmetry of the allowed vertices", + split, + "" if split == 1 else "es", + ) return [ replace(entry, group_members=tuple(members.get(entry.request.name, ()))) for entry in entries diff --git a/pyamplicol/src/pyamplicol/generation/structural_source_proof.py b/pyamplicol/src/pyamplicol/generation/structural_source_proof.py index ee865f980a..7c6c5fe5ce 100644 --- a/pyamplicol/src/pyamplicol/generation/structural_source_proof.py +++ b/pyamplicol/src/pyamplicol/generation/structural_source_proof.py @@ -11,8 +11,9 @@ import hashlib import json from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field from pathlib import Path -from typing import Any +from typing import Any, Protocol from pyamplicol.artifacts.security import sha256_file @@ -30,6 +31,24 @@ _DIAGNOSTIC_PATH_FIELDS = frozenset({"compiler_path"}) +@dataclass +class StructuralProofCache: + """Payload records shared by the proofs of one staged artifact. + + Only valid while the referenced files do not change; existing paths are + cached, missing ones are checked again. + """ + + # loose path -> (object record, nested payload paths) + loose: dict[str, tuple[dict[str, object], tuple[str, ...]]] = field( + default_factory=dict + ) + # (container path, index sha256) -> (member records, container record) + containers: dict[ + tuple[str, str], tuple[dict[str, dict[str, object]], dict[str, object]] + ] = field(default_factory=dict) + + def canonical_sha256(domain: str, value: object) -> str: """Hash canonical JSON under a distinct semantic domain.""" @@ -219,15 +238,59 @@ def add(record: object, path: str) -> None: return tuple(lanes) +class StagedPayloads(Protocol): + """The payloads of a staged artifact, wherever they are stored.""" + + def has_payload(self, relative: str) -> bool: ... + + def read_staged(self, relative: str) -> bytes: ... + + +class _DirectoryPayloads: + """Payloads stored as files below ``root``.""" + + def __init__(self, root: Path) -> None: + self._root = root + + def has_payload(self, relative: str) -> bool: + path = self._root / relative + return path.is_file() and not path.is_symlink() + + def read_staged(self, relative: str) -> bytes: + return (self._root / relative).read_bytes() + + +def _loose_object( + payloads: StagedPayloads, + loose: str, +) -> tuple[dict[str, object], tuple[str, ...]]: + """The inventory record of a loose payload and the paths it references.""" + + data = payloads.read_staged(loose) + record: dict[str, object] = { + "path": loose, + "sha256": hashlib.sha256(data).hexdigest(), + "size_bytes": len(data), + } + if not loose.endswith(".json"): + return record, () + try: + nested = json.loads(data.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise ValueError( + f"referenced lane manifest is not valid JSON: {loose}" + ) from error + return record, _payload_paths(nested) + + def _resolve_loose_payload( - root: Path, + payloads: StagedPayloads, process_id: str, relative: str, ) -> str | None: candidates = (relative, f"processes/{process_id}/{relative}") for candidate in candidates: - path = root / candidate - if path.is_file() and not path.is_symlink(): + if payloads.has_payload(candidate): return candidate return None @@ -243,13 +306,30 @@ def build_generation_structural_proof( execution: Mapping[str, object], evaluator_container_path: str | None, evaluator_container_index_sha256: str | None, + cache: StructuralProofCache | None = None, + payloads: StagedPayloads | None = None, ) -> dict[str, object]: - """Build exact source/lane evidence from the private artifact staging tree.""" + """Build exact source/lane evidence from the private artifact staging tree. + + ``cache`` shares the records of payloads common to the processes of one + staged artifact. ``payloads`` reads the staged payloads, by default the + files below ``artifact_root``. + """ + + if payloads is None: + payloads = _DirectoryPayloads(artifact_root) lanes = execution_lanes(execution) member_records: dict[str, dict[str, object]] = {} container_record: dict[str, object] | None = None - if evaluator_container_path is not None: + container_key = ( + None + if evaluator_container_path is None or evaluator_container_index_sha256 is None + else (evaluator_container_path, evaluator_container_index_sha256) + ) + if cache is not None and container_key in cache.containers: + member_records, container_record = cache.containers[container_key] + elif evaluator_container_path is not None: container = artifact_root / evaluator_container_path if not container.is_file() or container.is_symlink(): raise ValueError("evaluator container is absent during structural proof") @@ -275,15 +355,19 @@ def build_generation_structural_proof( "index_sha256": reader.index.index_sha256, "member_count": len(reader.members), } + if cache is not None and container_key is not None: + cache.containers[container_key] = (member_records, container_record) loose_objects: dict[str, dict[str, object]] = { execution_path: { "path": execution_path, "sha256": execution_sha256, - "size_bytes": (artifact_root / execution_path).stat().st_size, + "size_bytes": len(payloads.read_staged(execution_path)), } } used_members: set[str] = set() + # the staged files do not change while this proof is built + resolved: dict[str, str | None] = {} lane_records: list[dict[str, object]] = [] roles: list[dict[str, str]] = [] for lane_id, lane in lanes: @@ -299,7 +383,14 @@ def include( lane_members.add(relative) used_members.add(relative) return - loose = _resolve_loose_payload(artifact_root, process_id, relative) + if relative in resolved: + loose = resolved[relative] + elif cache is not None and relative in cache.loose: + loose = resolved[relative] = relative + else: + loose = resolved[relative] = _resolve_loose_payload( + payloads, process_id, relative + ) if loose is None: # Some manifest paths are logical roots interpreted by the # runtime rather than physical files. Preserve them in the @@ -308,24 +399,15 @@ def include( if loose in lane_loose: return lane_loose.add(loose) - path = artifact_root / loose - loose_objects.setdefault( - loose, - { - "path": loose, - "sha256": sha256_file(path), - "size_bytes": path.stat().st_size, - }, - ) - if path.suffix == ".json": - try: - nested = json.loads(path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError) as error: - raise ValueError( - f"referenced lane manifest is not valid JSON: {loose}" - ) from error - for nested_relative in _payload_paths(nested): - include(nested_relative, lane_loose, lane_members) + cached = None if cache is None else cache.loose.get(loose) + if cached is None: + cached = _loose_object(payloads, loose) + if cache is not None: + cache.loose[loose] = cached + record, nested_paths = cached + loose_objects.setdefault(loose, dict(record)) + for nested_relative in nested_paths: + include(nested_relative, lane_loose, lane_members) for relative in _payload_paths(lane): include(relative, referenced_loose, referenced_members) @@ -641,6 +723,8 @@ def validate_generation_structural_proof( "ROLE", "SCHEMA", "SEMANTIC_MAP_DOMAINS", + "StagedPayloads", + "StructuralProofCache", "build_generation_structural_proof", "canonical_sha256", "execution_lanes", diff --git a/pyamplicol/src/pyamplicol/models/prepared.py b/pyamplicol/src/pyamplicol/models/prepared.py index 5e17fd9790..9fc0a9b52f 100644 --- a/pyamplicol/src/pyamplicol/models/prepared.py +++ b/pyamplicol/src/pyamplicol/models/prepared.py @@ -1151,6 +1151,12 @@ class PreparedKernelPack: default=None, ) ) + _recurrence_template_catalog: RecurrenceTemplateCatalog | None = field( + init=False, + repr=False, + compare=False, + default=None, + ) def __post_init__(self) -> None: if self.backend not in _BACKENDS: @@ -1412,12 +1418,18 @@ def recurrence_template_catalog(self) -> RecurrenceTemplateCatalog | None: recurrence_template = self.recurrence_template if recurrence_template is None: return None + # The catalog is immutable and model-wide, but rebuilding it is costly + cached = self._recurrence_template_catalog + if cached is not None: + return cached from .recurrence_template import RecurrenceTemplateCatalog thawed = _thaw_json(recurrence_template) - return RecurrenceTemplateCatalog.from_dict( + catalog = RecurrenceTemplateCatalog.from_dict( _mapping(thawed, "kernel_pack.recurrence_template") ) + object.__setattr__(self, "_recurrence_template_catalog", catalog) + return catalog @property def recurrence_direct_template_catalog( diff --git a/pyamplicol/src/pyamplicol/runtime/_evaluator_payloads.py b/pyamplicol/src/pyamplicol/runtime/_evaluator_payloads.py index a0270b1fd4..788c14c37f 100644 --- a/pyamplicol/src/pyamplicol/runtime/_evaluator_payloads.py +++ b/pyamplicol/src/pyamplicol/runtime/_evaluator_payloads.py @@ -8,7 +8,11 @@ from pathlib import Path from pyamplicol.api.errors import ArtifactError, CompatibilityError -from pyamplicol.artifacts.manifest import ArtifactManifest, PayloadRecord +from pyamplicol.artifacts.manifest import ( + ArtifactManifest, + ArtifactPayloads, + PayloadRecord, +) from pyamplicol.artifacts.security import confined_path, normalize_relative_path from pyamplicol.generation.evaluator_container import ( PacbinError, @@ -35,6 +39,8 @@ class ExactEvaluatorPayloadResolver: def __init__(self, manifest: ArtifactManifest) -> None: self._root = manifest.root self._loose = {record.path: record for record in manifest.payloads} + # declared payloads, loose or packed in payloads.pacbin + self.payloads = ArtifactPayloads(manifest) self._container_path: Path | None = None self._members: Mapping[str, PacbinMember] = {} @@ -152,7 +158,7 @@ def read_exact_state( record = self._loose.get(normalized) if record is not None: try: - return confined_path(self._root, normalized).read_bytes() + return self.payloads.read_bytes(normalized) except OSError as exc: raise CompatibilityError( f"could not read retained Symbolica evaluator state " diff --git a/pyamplicol/src/pyamplicol/runtime/backend.py b/pyamplicol/src/pyamplicol/runtime/backend.py index f5bae308f3..d5b8fdf0f8 100644 --- a/pyamplicol/src/pyamplicol/runtime/backend.py +++ b/pyamplicol/src/pyamplicol/runtime/backend.py @@ -45,6 +45,7 @@ WarmUpResult, ) from pyamplicol.artifacts import MANIFEST_NAME, ArtifactManifest, load_manifest +from pyamplicol.artifacts.manifest import ArtifactPayloads from pyamplicol.reporting import ( ProgressEnd, ProgressSink, @@ -405,12 +406,12 @@ def _has_reusable_runtime_selector_contract( ) -> bool: relative = str(process["physics_path"]) try: - payload = json.loads((manifest.root / relative).read_text(encoding="utf-8")) + payload = json.loads(ArtifactPayloads(manifest).read_bytes(relative)) except OSError as exc: raise ArtifactError( f"could not read runtime selector metadata {relative}: {exc}" ) from exc - except json.JSONDecodeError as exc: + except (UnicodeDecodeError, json.JSONDecodeError) as exc: raise ArtifactError( f"invalid runtime selector metadata {relative}: {exc}" ) from exc @@ -1622,8 +1623,8 @@ def validation_momenta(self) -> Momenta | None: f"process {process_id!r} must declare exactly one validation point" ) try: - raw = json.loads((manifest.root / payloads[0].path).read_text("utf-8")) - except (OSError, json.JSONDecodeError) as exc: + raw = json.loads(ArtifactPayloads(manifest).read_bytes(payloads[0].path)) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: raise ArtifactError( f"could not read validation point for process {process_id!r}: {exc}" ) from exc diff --git a/pyamplicol/src/pyamplicol/runtime/eager_exact/_contracts.py b/pyamplicol/src/pyamplicol/runtime/eager_exact/_contracts.py index b39c42f549..b3c4c0d6c0 100644 --- a/pyamplicol/src/pyamplicol/runtime/eager_exact/_contracts.py +++ b/pyamplicol/src/pyamplicol/runtime/eager_exact/_contracts.py @@ -15,8 +15,12 @@ PROCESS_ARTIFACT_SCHEMA_VERSION, ) from pyamplicol.api.errors import ArtifactError, CompatibilityError, EvaluationError -from pyamplicol.artifacts.manifest import ArtifactManifest, PayloadRecord -from pyamplicol.artifacts.security import confined_path, normalize_relative_path +from pyamplicol.artifacts.manifest import ( + ArtifactManifest, + ArtifactPayloads, + PayloadRecord, +) +from pyamplicol.artifacts.security import normalize_relative_path from pyamplicol.generation.eager_tables import ( EAGER_PLAN_ABI, EAGER_RUNTIME_CAPABILITY, @@ -180,10 +184,12 @@ def _integer(value: object, context: str, *, minimum: int = 0) -> int: return value -def _read_json(path: Path, context: str) -> Mapping[str, object]: +def _read_payload_json( + payloads: ArtifactPayloads, path: str, context: str +) -> Mapping[str, object]: try: - value = json.loads(path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError) as exc: + value = json.loads(payloads.read_bytes(path)) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: raise ArtifactError(f"could not load {context}: {exc}") from exc return _mapping(value, context) @@ -199,10 +205,19 @@ def _joined_payload_path(prefix: str, member: str) -> str: @dataclass(frozen=True, slots=True) class _PayloadIndex: records: Mapping[str, PayloadRecord] + files: ArtifactPayloads @classmethod def from_manifest(cls, manifest: ArtifactManifest) -> _PayloadIndex: - return cls({record.path: record for record in manifest.payloads}) + return cls( + {record.path: record for record in manifest.payloads}, + ArtifactPayloads(manifest), + ) + + def read_json(self, path: str, context: str) -> Mapping[str, object]: + """Load one declared JSON payload, loose or packed.""" + + return _read_payload_json(self.files, path, context) def require( self, @@ -469,9 +484,9 @@ def _load_table( ) global_path = _joined_payload_path(process_prefix, path) payloads.require(global_path, role="evaluator-state", process_id=process_id) - table_path = confined_path(process_root, path) + del process_root try: - content = table_path.read_bytes() + content = payloads.files.read_bytes(global_path) except OSError as exc: raise ArtifactError(f"could not read {context}: {exc}") from exc if len(content) != count * expected_size: diff --git a/pyamplicol/src/pyamplicol/runtime/eager_exact/_executor.py b/pyamplicol/src/pyamplicol/runtime/eager_exact/_executor.py index b7a0b58119..4bcf324354 100644 --- a/pyamplicol/src/pyamplicol/runtime/eager_exact/_executor.py +++ b/pyamplicol/src/pyamplicol/runtime/eager_exact/_executor.py @@ -12,7 +12,6 @@ from pyamplicol.api.protocols import Momenta from pyamplicol.api.results import ResolvedEvaluation from pyamplicol.artifacts import load_manifest -from pyamplicol.artifacts.security import confined_path from pyamplicol.runtime._color_topology_exact import ( apply_exact_color_replay_input_mapping, parse_exact_color_topology_replay, @@ -29,7 +28,6 @@ _KernelLoader, _mapping, _PayloadIndex, - _read_json, ) from pyamplicol.runtime.eager_exact._execution import _evaluate_point from pyamplicol.runtime.eager_exact._plan import _EagerExactPlan @@ -86,12 +84,11 @@ def __init__( payloads.require( physics_path, role="runtime-physics", process_id=representative_id ) - execution = _read_json( - confined_path(self._artifact, execution_records[0].path), - "eager execution metadata", + execution = payloads.read_json( + execution_records[0].path, "eager execution metadata" ) - representative_physics = _read_json( - confined_path(self._artifact, physics_path), "runtime physics metadata" + representative_physics = payloads.read_json( + physics_path, "runtime physics metadata" ) axes = native_physics_axes(native_runtime, representative_physics) physics = axes.public_physics diff --git a/pyamplicol/src/pyamplicol/runtime/eager_exact/_plan.py b/pyamplicol/src/pyamplicol/runtime/eager_exact/_plan.py index 853fe79a80..719b958715 100644 --- a/pyamplicol/src/pyamplicol/runtime/eager_exact/_plan.py +++ b/pyamplicol/src/pyamplicol/runtime/eager_exact/_plan.py @@ -11,7 +11,7 @@ from pyamplicol.api.errors import ArtifactError, CompatibilityError, EvaluationError from pyamplicol.artifacts.manifest import ArtifactManifest -from pyamplicol.artifacts.security import confined_path, normalize_relative_path +from pyamplicol.artifacts.security import normalize_relative_path from pyamplicol.generation.eager_tables import ( EAGER_KERNEL_ABI, EAGER_OUTPUT_FACTOR_NONE, @@ -48,7 +48,6 @@ _load_table, _mapping, _PayloadIndex, - _read_json, _sequence, _validate_execution_header, ) @@ -409,9 +408,7 @@ def load( pack_path = normalize_relative_path(pack_path) payload_root_name = normalize_relative_path(payload_root_name) payloads.require(pack_path, role="evaluator-manifest", process_id=None) - pack_payload = _read_json( - confined_path(artifact_root, pack_path), "eager kernel pack" - ) + pack_payload = payloads.read_json(pack_path, "eager kernel pack") raw_backend = pack_payload.get("backend") if raw_backend not in _SUPPORTED_PREPARED_BACKENDS: raise CompatibilityError( @@ -1109,9 +1106,7 @@ def _load_exact_kernel_pack( pack_path = normalize_relative_path(pack_path) payload_root_name = normalize_relative_path(payload_root_name) payloads.require(pack_path, role="evaluator-manifest", process_id=None) - pack_payload = _read_json( - confined_path(artifact_root, pack_path), "eager kernel pack" - ) + pack_payload = payloads.read_json(pack_path, "eager kernel pack") raw_backend = pack_payload.get("backend") if raw_backend not in _SUPPORTED_PREPARED_BACKENDS: raise CompatibilityError(f"unsupported prepared eager backend {raw_backend!r}") diff --git a/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_color.py b/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_color.py index 588fed2cb1..7b9e5d2724 100644 --- a/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_color.py +++ b/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_color.py @@ -12,8 +12,12 @@ from pathlib import Path, PurePosixPath from pyamplicol.api.errors import ArtifactError, CompatibilityError -from pyamplicol.artifacts.manifest import ArtifactManifest, PayloadRecord -from pyamplicol.artifacts.security import confined_path, normalize_relative_path +from pyamplicol.artifacts.manifest import ( + ArtifactManifest, + ArtifactPayloads, + PayloadRecord, +) +from pyamplicol.artifacts.security import normalize_relative_path from pyamplicol.color.symmetric_group import ( _inverse_permutation, _lexicographic_permutation_rank, @@ -252,7 +256,7 @@ def _load_recurrence_color_contraction( "recurrence color-contraction payload disagrees with execution metadata" ) try: - payload = confined_path(artifact_root, logical_path).read_bytes() + payload = ArtifactPayloads(manifest).read_bytes(logical_path) except OSError as exc: raise ArtifactError( f"could not read recurrence color-contraction payload: {exc}" diff --git a/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_executor.py b/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_executor.py index 8bde998036..930ea0c5d6 100644 --- a/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_executor.py +++ b/pyamplicol/src/pyamplicol/runtime/recurrence_exact/_executor.py @@ -12,7 +12,6 @@ from pyamplicol.api.protocols import Momenta from pyamplicol.api.results import ResolvedEvaluation from pyamplicol.artifacts import load_manifest -from pyamplicol.artifacts.security import confined_path from pyamplicol.runtime._evaluator_payloads import ExactEvaluatorPayloadResolver from pyamplicol.runtime._native_selection import ( native_physics_axes, @@ -24,7 +23,6 @@ _KernelLoader, _mapping, _PayloadIndex, - _read_json, _sequence, ) from pyamplicol.runtime.symbolica_exact import ( @@ -176,13 +174,11 @@ def __init__( role="runtime-physics", process_id=representative_id, ) - execution = _read_json( - confined_path(self._artifact, execution_records[0].path), - "recurrence execution metadata", + execution = payloads.read_json( + execution_records[0].path, "recurrence execution metadata" ) - representative_physics = _read_json( - confined_path(self._artifact, physics_path), - "runtime physics metadata", + representative_physics = payloads.read_json( + physics_path, "runtime physics metadata" ) axes = native_physics_axes(native_runtime, representative_physics) if execution.get("key") != selection.representative_process_key: diff --git a/pyamplicol/src/pyamplicol/runtime/symbolica_exact.py b/pyamplicol/src/pyamplicol/runtime/symbolica_exact.py index 30f357a2bc..1aef0fcbe1 100644 --- a/pyamplicol/src/pyamplicol/runtime/symbolica_exact.py +++ b/pyamplicol/src/pyamplicol/runtime/symbolica_exact.py @@ -27,7 +27,8 @@ from pyamplicol.api.protocols import Momenta from pyamplicol.api.results import ResolvedEvaluation from pyamplicol.artifacts import load_manifest -from pyamplicol.artifacts.security import confined_path, normalize_relative_path +from pyamplicol.artifacts.manifest import ArtifactPayloads +from pyamplicol.artifacts.security import normalize_relative_path from pyamplicol.runtime._color_topology_exact import ( apply_exact_color_replay_input_mapping, parse_exact_color_topology_replay, @@ -508,12 +509,12 @@ def __init__(self, artifact: Path, process_id: str, native_runtime: Any) -> None ) try: self._execution = json.loads( - (artifact / records[0].path).read_text(encoding="utf-8") + self._payloads.payloads.read_text(records[0].path) ) representative_physics = json.loads( - confined_path(artifact, physics_path).read_text(encoding="utf-8") + self._payloads.payloads.read_text(physics_path) ) - except (OSError, json.JSONDecodeError) as exc: + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc: raise ArtifactError( f"could not load exact-runtime metadata: {exc}" ) from exc @@ -841,7 +842,7 @@ def _derive_model_parameters( "exact parameter derivation requires one compiled model" ) model = json.loads( - confined_path(self._artifact, model_records[0].path).read_text() + ArtifactPayloads(manifest).read_text(model_records[0].path) )["ir"] self._exact_compiled_model = model definitions = { diff --git a/pyamplicol/tests/unit/test_artifact_payloads.py b/pyamplicol/tests/unit/test_artifact_payloads.py new file mode 100644 index 0000000000..0ef010ae56 --- /dev/null +++ b/pyamplicol/tests/unit/test_artifact_payloads.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: 0BSD +from __future__ import annotations + +import json +import shutil +from pathlib import Path + +import pytest + +from pyamplicol import ArtifactError +from pyamplicol.artifacts import load_manifest +from pyamplicol.artifacts.manifest import ArtifactPayloads, validate_payloads +from pyamplicol.generation.evaluator_container import PacbinReader +from pyamplicol.runtime._evaluator_payloads import ExactEvaluatorPayloadResolver +from pyamplicol.runtime.eager_exact._contracts import _PayloadIndex +from tests.unit.test_artifacts import _build_packable + +_EXECUTION = "processes/dd_to_z/execution.json" +_BINDING = "processes/dd_to_z/recurrence-binding.bin" + + +@pytest.mark.parametrize("pack", [False, True]) +def test_payloads_resolve_loose_and_packed_layouts(tmp_path: Path, pack: bool) -> None: + root = tmp_path / "artifact" + _build_packable(root, pack=pack) + manifest = load_manifest(root, verify_payloads=True) + payloads = ArtifactPayloads(manifest) + + assert payloads.is_packed(_EXECUTION) is pack + assert (payloads.container_path is not None) is pack + assert json.loads(payloads.read_text(_EXECUTION)) == {"plan": [1, 2, 3]} + assert payloads.read_bytes(_BINDING) == b"binding" + # nested containers and top-level payloads stay loose + assert not payloads.is_packed("processes/dd_to_z/eager-runtime.pacbin") + assert ( + payloads.read_bytes("physics/process.json") + == (root / "physics/process.json").read_bytes() + ) + assert _PayloadIndex.from_manifest(manifest).read_json(_EXECUTION, "execution") == { + "plan": [1, 2, 3] + } + assert ExactEvaluatorPayloadResolver(manifest).payloads.read_bytes(_BINDING) == ( + b"binding" + ) + + +def _corrupt_member(root: Path, logical_path: str) -> None: + container = root / "payloads.pacbin" + with PacbinReader.open(container, verify_payloads=False) as reader: + offset = reader.member(logical_path).offset + data = bytearray(container.read_bytes()) + data[offset] ^= 0xFF + container.chmod(0o644) + container.write_bytes(bytes(data)) + + +def test_packed_payload_reads_verify_member_digests(tmp_path: Path) -> None: + root = tmp_path / "artifact" + _build_packable(root, pack=True) + _corrupt_member(root, _BINDING) + manifest = load_manifest(root) + payloads = ArtifactPayloads(manifest) + + assert json.loads(payloads.read_text(_EXECUTION)) == {"plan": [1, 2, 3]} + with pytest.raises(ArtifactError, match="digest mismatch"): + payloads.read_bytes(_BINDING) + with pytest.raises(ArtifactError, match="payload container verification"): + validate_payloads(manifest) + + +def test_packed_payloads_may_not_also_be_files(tmp_path: Path) -> None: + root = tmp_path / "artifact" + _build_packable(root, pack=True) + shadow = root / _EXECUTION + shadow.parent.mkdir(parents=True, exist_ok=True) + shadow.write_text('{"plan": []}', encoding="utf-8") + + with pytest.raises(ArtifactError, match="also stored as a file"): + load_manifest(root, verify_payloads=True) + + +def _pack_artifact(source: Path, root: Path) -> None: + """Copy ``source`` and move its packable payloads into payloads.pacbin.""" + + from pyamplicol.artifacts.manifest import compute_artifact_id + from pyamplicol.artifacts.security import sha256_file + from pyamplicol.generation.artifact_writer import _packs_payload + from pyamplicol.generation.evaluator_container import ( + PacbinMemberKind, + PacbinMemberSource, + write_pacbin, + ) + + shutil.copytree(source, root) + raw = json.loads((root / "artifact.json").read_text(encoding="utf-8")) + packed = [ + record["path"] for record in raw["payloads"] if _packs_payload(record["path"]) + ] + assert packed + container = root / "payloads.pacbin" + write_pacbin( + container, + [ + PacbinMemberSource(path, PacbinMemberKind.DECLARED_PAYLOAD, root / path) + for path in packed + ], + ) + for path in packed: + (root / path).unlink() + raw["payloads"] = sorted( + [ + *raw["payloads"], + { + "executable": False, + "media_type": "application/octet-stream", + "path": "payloads.pacbin", + "role": "payload-container", + "sha256": sha256_file(container), + "size_bytes": container.stat().st_size, + }, + ], + key=lambda record: record["path"], + ) + raw["artifact_id"] = compute_artifact_id(raw) + (root / "artifact.json").write_text(json.dumps(raw), encoding="utf-8") + + +def test_inspection_reads_packed_payloads(tmp_path: Path) -> None: + from pyamplicol.artifacts import inspect_artifact + + source = ( + Path(__file__).resolve().parents[2] + / "src/pyamplicol/assets/selftest/portable-64le/artifact" + ) + root = tmp_path / "packed" + _pack_artifact(source, root) + assert ArtifactPayloads(load_manifest(root)).is_packed( + "processes/d_dbar_to_z/physics.json" + ) + assert not (root / "processes/d_dbar_to_z/physics.json").exists() + + loose = inspect_artifact(source) + packed = inspect_artifact(root) + assert packed.integrity == "verified" + assert packed.processes == loose.processes diff --git a/pyamplicol/tests/unit/test_artifacts.py b/pyamplicol/tests/unit/test_artifacts.py index ef8a7d537d..08537bdde1 100644 --- a/pyamplicol/tests/unit/test_artifacts.py +++ b/pyamplicol/tests/unit/test_artifacts.py @@ -1095,3 +1095,118 @@ def test_process_group_update_waits_for_the_generation_lock( writer.join() assert read_process_groups(root) is not None load_manifest(root, verify_payloads=True) + + +def _build_packable(path: Path, *, pack: bool) -> None: + packs = generation_artifact_writer._packs_payload if pack else None + with ArtifactBuilder(path, packs=packs) as builder: + builder.add_json( + "physics/process.json", + {"schema_version": 1}, + role="runtime-physics", + process_id="dd_to_z", + ) + builder.add_json( + "processes/dd_to_z/execution.json", + {"plan": [1, 2, 3]}, + role="evaluator-manifest", + process_id="dd_to_z", + ) + for relative, content in ( + ("processes/dd_to_z/recurrence-binding.bin", b"binding"), + ("processes/dd_to_z/eager-runtime.pacbin", b"nested container"), + ("recurrence/schedules/ab/recurrence-runtime.pacbin", b"schedule"), + ): + builder.add_bytes( + relative, + content, + role="evaluator-manifest", + media_type="application/octet-stream", + process_id="dd_to_z", + ) + assert builder.root is not None + if pack: + # packable payloads are spooled, never written as files + assert not (builder.root / "processes/dd_to_z/execution.json").exists() + assert not (builder.root / "recurrence").exists() + assert builder.read_staged("processes/dd_to_z/recurrence-binding.bin") == ( + b"binding" + ) + if pack: + builder.seal_payload_container() + assert builder.read_staged("processes/dd_to_z/recurrence-binding.bin") == ( + b"binding" + ) + with pytest.raises(RuntimeError, match="sealed"): + builder.add_bytes( + "processes/dd_to_z/late.bin", + b"late", + role="evaluator-manifest", + media_type="application/octet-stream", + ) + builder.finalize( + kind="pyamplicol-process", + producer=_producer(), + model=_model(), + configuration=_configuration(), + processes=[_process()], + default_process_id="dd_to_z", + runtime=_runtime(), + ) + + +def test_packed_payloads_keep_records_and_artifact_identity(tmp_path: Path) -> None: + _build_packable(tmp_path / "loose", pack=False) + _build_packable(tmp_path / "packed", pack=True) + loose = load_manifest(tmp_path / "loose") + packed = load_manifest(tmp_path / "packed") + + assert packed.artifact_id == loose.artifact_id + assert {record.path for record in packed.payloads} == { + record.path for record in loose.payloads + } | {"payloads.pacbin"} + assert generation_artifact_writer._validate_payload_container(packed) == { + "processes/dd_to_z/execution.json", + "processes/dd_to_z/recurrence-binding.bin", + "recurrence/schedules/ab/recurrence-runtime.pacbin", + } + root = tmp_path / "packed" + assert not [ + path for path in root.rglob("*") if path.is_dir() and not any(path.iterdir()) + ] + assert not (root / "recurrence").exists() + assert not (root / "processes/dd_to_z/execution.json").exists() + # a nested container stays a file for the loaders that open it + assert (root / "processes/dd_to_z/eager-runtime.pacbin").is_file() + with PacbinReader.open(root / "payloads.pacbin") as reader: + assert ( + reader.read_member( + "processes/dd_to_z/recurrence-binding.bin", length=len(b"binding") + ) + == b"binding" + ) + + +def test_payload_container_must_match_its_records(tmp_path: Path) -> None: + from pyamplicol.generation.evaluator_container import ( + PacbinMemberKind, + PacbinMemberSource, + write_pacbin_atomic, + ) + + root = tmp_path / "packed" + _build_packable(root, pack=True) + container = root / "payloads.pacbin" + container.chmod(0o644) + write_pacbin_atomic( + container, + [ + PacbinMemberSource( + "processes/dd_to_z/recurrence-binding.bin", + PacbinMemberKind.DECLARED_PAYLOAD, + io.BytesIO(b"tampered"), + ) + ], + ) + with pytest.raises(ValueError, match="invalid payload container member"): + generation_artifact_writer._validate_payload_container(load_manifest(root)) diff --git a/pyamplicol/tests/unit/test_config_models.py b/pyamplicol/tests/unit/test_config_models.py index dd19a57651..f0e9f29c8d 100644 --- a/pyamplicol/tests/unit/test_config_models.py +++ b/pyamplicol/tests/unit/test_config_models.py @@ -40,7 +40,7 @@ def test_schema_v1_registry_contains_every_contract_leaf() -> None: - assert len(FIELD_REGISTRY) == 83 + assert len(FIELD_REGISTRY) == 85 assert "generation.umami_grouping" not in FIELD_REGISTRY assert "evaluator.jit.direct_translation" not in FIELD_REGISTRY assert FIELD_REGISTRY["action"].required diff --git a/pyamplicol/tests/unit/test_flavor_blind_expansion.py b/pyamplicol/tests/unit/test_flavor_blind_expansion.py index 3fa20c4429..bbdf43c122 100644 --- a/pyamplicol/tests/unit/test_flavor_blind_expansion.py +++ b/pyamplicol/tests/unit/test_flavor_blind_expansion.py @@ -251,6 +251,34 @@ def test_flavor_blind_warns_when_couplings_distinguish_the_flavours( assert dijet["s c > s c"][1].warning is None +def test_strict_flavor_blind_generates_the_members_it_cannot_verify( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(licensing, "detect_symbolica_license", _restricted_license) + + assumed = _expanded("p p > j j", flavor_scheme=4, coupling_order_policy="explicit") + strict = _expanded( + "p p > j j", + flavor_scheme=4, + coupling_order_policy="explicit", + flavor_blind_strict=True, + ) + + members = _members(strict) + assert all(member.warning is None for _, member in members.values()) + generated = {entry.request.expression for entry in strict} + assert "u u~ > g g" in generated + assert "u u~ > g g" not in {entry.request.expression for entry in assumed} + assert members["s c > s c"][0] == "d u > d u" + # Every subprocess is still generated or routed exactly once. + assert len(generated) + len(members) == len(assumed) + len(_members(assumed)) + # Pure QCD needs no split. + qcd = _expanded("p p > j j", flavor_scheme=4, flavor_blind_strict=True) + assert [e.request.name for e in qcd] == [ + e.request.name for e in _expanded("p p > j j", flavor_scheme=4) + ] + + def test_flavor_blind_keeps_processes_with_explicit_aliases( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/pyamplicol/tests/unit/test_recurrence_direct_artifact_metadata.py b/pyamplicol/tests/unit/test_recurrence_direct_artifact_metadata.py index f6c0f377f3..ac845c70b9 100644 --- a/pyamplicol/tests/unit/test_recurrence_direct_artifact_metadata.py +++ b/pyamplicol/tests/unit/test_recurrence_direct_artifact_metadata.py @@ -1084,7 +1084,7 @@ def test_direct_recurrence_inspection_exposes_arena_counters(tmp_path: Path) -> _write_prepared_pack(tmp_path) execution = _execution_manifest() result = inspection._recurrence_execution_inspection( - SimpleNamespace(root=tmp_path), + SimpleNamespace(root=tmp_path, payloads=()), execution, ) @@ -1133,7 +1133,7 @@ def test_direct_recurrence_inspection_exposes_contracted_color_summary( } result = inspection._recurrence_execution_inspection( - SimpleNamespace(root=tmp_path), + SimpleNamespace(root=tmp_path, payloads=()), execution, ) @@ -1160,7 +1160,7 @@ def test_inspection_requires_v2_metadata_without_packet_fallback( with pytest.raises(ArtifactError, match="regenerate the recurrence artifact"): inspection._recurrence_execution_inspection( - SimpleNamespace(root=tmp_path), + SimpleNamespace(root=tmp_path, payloads=()), execution, ) @@ -1174,7 +1174,7 @@ def test_v2_inspection_tolerates_pending_direct_counter_wiring( execution = _execution_manifest(process) result = inspection._recurrence_execution_inspection( - SimpleNamespace(root=tmp_path), + SimpleNamespace(root=tmp_path, payloads=()), execution, ) diff --git a/pyamplicol/tests/unit/test_repository_policy.py b/pyamplicol/tests/unit/test_repository_policy.py index 38dc6ed7d6..8273b51dd9 100644 --- a/pyamplicol/tests/unit/test_repository_policy.py +++ b/pyamplicol/tests/unit/test_repository_policy.py @@ -112,6 +112,8 @@ def test_release_source_roots_contain_no_symlinks() -> None: for source_root in FIRST_PARTY_ROOTS for path in source_root.rglob("*") if _eligible(path) and path.is_symlink() + # shared with madspace, copied by content (_SHARED_SYMLINKS) + and path != ROOT / "src/pyamplicol/assets/api_templates/umami/umami.h" ] assert not links, "release source contains symlinks:\n" + "\n".join(links) diff --git a/pyamplicol/tests/unit/test_umami_metadata.py b/pyamplicol/tests/unit/test_umami_metadata.py index 5c03863708..507257026c 100644 --- a/pyamplicol/tests/unit/test_umami_metadata.py +++ b/pyamplicol/tests/unit/test_umami_metadata.py @@ -741,6 +741,7 @@ def test_writer_append_recovers_old_compact_otf_sdk(tmp_path, monkeypatch): written = {} builder = SimpleNamespace( staged_path=lambda path: tmp_path / path, + read_staged=lambda path: (tmp_path / path).read_bytes(), payload_records=lambda: [], add_bytes=lambda path, content, **kwargs: written.setdefault(path, content), ) @@ -965,6 +966,7 @@ def test_writer_exports_only_ungenerated_group_members(tmp_path): written = {} builder = SimpleNamespace( staged_path=lambda path: tmp_path / path, + read_staged=lambda path: (tmp_path / path).read_bytes(), payload_records=lambda: [], add_bytes=lambda path, content, **kwargs: written.setdefault(path, content), ) diff --git a/pyamplicol/tests/unit/test_umami_phasespace_partners.py b/pyamplicol/tests/unit/test_umami_phasespace_partners.py new file mode 100644 index 0000000000..db997f6d85 --- /dev/null +++ b/pyamplicol/tests/unit/test_umami_phasespace_partners.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: 0BSD +"""The phase-space channels that integrate one leading-colour order together: +where its colour singlets attach to the colour lines, which non-cyclic order +its lines take, and which s-channel resonances it maps out.""" + +from __future__ import annotations + +import logging +import time +from dataclasses import replace +from types import SimpleNamespace + +from pyamplicol.artifacts.umami import ( + _MAX_PARTNERS, + _configurations, + _line_blocks, + _phasespace_partners, +) +from pyamplicol.generation.s_channel_resonances import s_channel_resonances +from pyamplicol.models.base import Particle +from tests.unit.test_umami_metadata import export, process + +# all-outgoing colour representations +Q, QBAR, G, S = 3, -3, 8, 1 + + +def test_line_blocks_follow_the_word(): + roles = [QBAR, Q, G, Q, QBAR, G] + assert _line_blocks([1, 2, 0, 3, 5, 4], roles) == [[1, 2, 0], [3, 5, 4]] + # a gluon trace has no open lines + assert _line_blocks([0, 1, 2], [G, G, G]) is None + + +def test_singlets_follow_the_antiquark_in_every_order(): + # q q~ > e+ e- g, the word q g q~ crossed: slots 1 (q~ in) ... 0 (q in) + roles = [QBAR, Q, S, S, G] + partners = _phasespace_partners([1, 4, 0], [2, 3], roles, [()]) + assert [order for order, _ in partners] == [[1, 4, 0, 2, 3], [1, 4, 0, 3, 2]] + assert all(configuration == () for _, configuration in partners) + + +def test_singlets_are_distributed_over_the_lines(): + # two lines and one singlet: after the first or the second antiquark + roles = [Q, QBAR, Q, QBAR, S] + partners = _phasespace_partners([0, 1, 2, 3], [4], roles, [()]) + assert [order for order, _ in partners] == [[0, 1, 2, 3, 4], [0, 1, 4, 2, 3]] + + +def test_three_lines_take_both_cyclic_orders(): + roles = [Q, QBAR] * 3 + partners = _phasespace_partners(list(range(6)), [], roles, [()]) + assert [order for order, _ in partners] == [ + [0, 1, 2, 3, 4, 5], + [0, 1, 4, 5, 2, 3], + ] + + +def test_a_resonance_keeps_its_decay_products_together(): + roles = [QBAR, Q, S, S, G] + z = (23, frozenset({2, 3})) + configurations, _ = _configurations([z]) + partners = _phasespace_partners([1, 4, 0], [2, 3], roles, configurations) + # the non-resonant placements and the Z, with its legs as one unit + assert partners == [ + ([1, 4, 0, 2, 3], ()), + ([1, 4, 0, 3, 2], ()), + ([1, 4, 0, 2, 3], ((23, frozenset({2, 3})),)), + ] + + +def test_a_colour_singlet_decaying_to_a_complete_line(): + roles = [Q, QBAR, Q, QBAR, G] + configurations, _ = _configurations([(23, frozenset({2, 3}))]) + # the Z decays to the line 2 3 of this order, but not to legs of two lines + assert len(_phasespace_partners([0, 1, 2, 3], [], roles, configurations)) == 2 + assert len(_phasespace_partners([0, 3, 2, 1], [], roles, configurations)) == 1 + + +def test_overlapping_resonances_are_alternatives(): + # e+ e- e+ ve: a Z from either positron, with the W from the other + z1, z2 = (23, frozenset({2, 4})), (23, frozenset({3, 4})) + w1, w2 = (24, frozenset({2, 5})), (24, frozenset({3, 5})) + configurations, complete = _configurations([z1, z2, w1, w2]) + assert complete + assert configurations == [(), (z1, w2), (z2, w1), (z1,), (z2,), (w1,), (w2,)] + + +def test_nested_resonances_stay_contiguous(): + roles = [QBAR, Q, S, S, S, S] + h = (25, frozenset({2, 3, 4, 5})) + z1, z2 = (23, frozenset({2, 4})), (23, frozenset({3, 5})) + configurations, _ = _configurations([h, z1, z2]) + partners = _phasespace_partners([1, 0], [2, 3, 4, 5], roles, configurations) + resonant = [order for order, configuration in partners if len(configuration) == 3] + assert resonant == [[1, 0, 2, 4, 3, 5]] + + +def test_every_compatible_set_of_resonances_is_a_configuration(): + # H > Z Z > 4l: each resonance alone and every nested or disjoint set + h = (25, frozenset({2, 3, 4, 5})) + z1, z2 = (23, frozenset({2, 4})), (23, frozenset({3, 5})) + configurations, complete = _configurations([h, z1, z2]) + assert complete + assert configurations == [ + (), + (h, z1, z2), + (h, z1), + (h, z2), + (z1, z2), + (h,), + (z1,), + (z2,), + ] + + +def test_many_resonances_are_capped_not_rejected(): + # 24 pairwise disjoint resonances have 2^24 configurations + resonances = [(23, frozenset({2 * i, 2 * i + 1})) for i in range(24)] + started = time.perf_counter() + configurations, complete = _configurations(resonances, limit=100) + assert time.perf_counter() - started < 1 + assert not complete + assert len(configurations) == 100 + assert configurations[0] == () + assert [len(c) for c in configurations[1:]] == sorted( + (len(c) for c in configurations[1:]), reverse=True + ) + assert _configurations(resonances, limit=100) == (configurations, complete) + + +def test_partners_are_truncated_deterministically(): + # eight singlets: 8! orders per configuration, enumerated only up to the cap + roles = [QBAR, Q] + [S] * 8 + singlets = list(range(2, 10)) + resonances = [(23, frozenset({2 * i, 2 * i + 1})) for i in range(1, 5)] + configurations, _ = _configurations(resonances) + started = time.perf_counter() + partners = _phasespace_partners([1, 0], singlets, roles, configurations) + assert time.perf_counter() - started < 0.5 + assert len(partners) <= _MAX_PARTNERS + assert partners[0] == ([1, 0, *singlets], ()) + # every configuration keeps its first placements, the largest ones first + assert [c for c in dict.fromkeys(c for _, c in partners)] == configurations + assert partners == _phasespace_partners([1, 0], singlets, roles, configurations) + + +def test_unmapped_singlet_resonances_are_dropped_with_one_warning(caplog): + # q q~ > s s g: an s decaying to an s and the gluon, and a coloured one + p = process(pdgs=(991, -991, 995, 995, 990), words=((2, 5, 1),)) + p.physics["extensions"]["s_channel_resonances"] = [ + {"pdg": 23, "labels": [3, 4], "color": 1}, + {"pdg": 25, "labels": [3, 5], "color": 1}, + {"pdg": 25, "labels": [4, 5], "color": 1}, + {"pdg": 6, "labels": [4, 5], "color": 3}, + ] + with caplog.at_level(logging.WARNING, logger="pyamplicol.artifacts.umami"): + (data,) = export(p) + (warning,) = caplog.records + assert "25 > 995 990, 25 > 995 990" in warning.getMessage() + (channel,) = data["channels"] + assert { + tuple(r["pdg"] for r in ps["resonances"]) + for ps in channel["phasespace_channels"] + } == {(), (23,)} + # a coloured resonance is not followed; the launch checks it + assert [(r["pdg"], r["color"]) for r in channel["colored_resonances"]] == [(6, 3)] + + +def test_singlet_resonance_inside_a_gluon_trace_is_dropped_with_a_warning(caplog): + p = process() + p.physics["extensions"]["s_channel_resonances"] = [ + {"pdg": 25, "labels": [3, 4], "color": 1} + ] + with caplog.at_level(logging.WARNING, logger="pyamplicol.artifacts.umami"): + (data,) = export(p) + (warning,) = caplog.records + assert "25 > 990 990" in warning.getMessage() + (channel,) = data["channels"] + assert all(not ps["resonances"] for ps in channel["phasespace_channels"]) + + +def test_line_resonance_of_some_colour_orders_is_not_warned(caplog): + p = process(pdgs=(991, -991, 992, -992), words=((2, 1, 3, 4), (2, 4, 3, 1))) + p.physics["extensions"]["s_channel_resonances"] = [ + {"pdg": 23, "labels": [3, 4], "color": 1} + ] + sectors = ( + { + "id": 0, + "kind": "open-lines", + "word_labels": [2, 1, 3, 4], + "open_color_lines": [ + { + "fundamental_label": 2, + "adjoint_labels": [], + "antifundamental_label": 1, + }, + { + "fundamental_label": 3, + "adjoint_labels": [], + "antifundamental_label": 4, + }, + ], + }, + { + "id": 1, + "kind": "open-lines", + "word_labels": [2, 4, 3, 1], + "open_color_lines": [ + { + "fundamental_label": 2, + "adjoint_labels": [], + "antifundamental_label": 4, + }, + { + "fundamental_label": 3, + "adjoint_labels": [], + "antifundamental_label": 1, + }, + ], + }, + ) + with caplog.at_level(logging.WARNING, logger="pyamplicol.artifacts.umami"): + (data,) = export(replace(p, color_sectors=sectors)) + assert not caplog.records + resonant = [ + ps + for c in data["channels"] + for ps in c["phasespace_channels"] + if ps["resonances"] + ] + assert [r["pdg"] for ps in resonant for r in ps["resonances"]] == [23] + + +class _Model: + """f f~ > l- l+ through any vector of ``particles``.""" + + def __init__(self, particles, parameters): + self.particles = {p.pdg: p for p in particles} + self.parameters = parameters + + def vertices_accepting(self, left, right, color_accuracy): + if {left, right} != {11, -11}: + return [] + return [ + SimpleNamespace(particles=(left, right, pdg)) + for pdg in self.particles + if abs(pdg) != 11 + ] + + def skip_duplicate_vertex_orientation(self, vertex): + return False + + def vertex_coupling_orders(self, vertex): + return (("QED", 1),) + + def particle(self, pdg): + return self.particles[pdg] + + def anti_particle(self, pdg): + return self.particles[pdg].anti_pdg + + def charge(self, pdg): + return self.particles[pdg].charge + + def color_rep(self, pdg): + return self.particles[pdg].color_rep + + def runtime_mass_parameter_name(self, pdg): + return self.parameters.get((pdg, "mass")) + + def runtime_width_parameter_name(self, pdg): + return self.parameters.get((pdg, "width")) + + +def test_candidates_are_kept_unless_their_mass_or_width_is_zero_in_the_model(): + def vector(pdg, mass, width): + return Particle(pdg, pdg, 2, 4, 1, mass, width) + + particles = [ + Particle(11, -11, 1, 4, 1, 50.0, 0.0, -1.0), + Particle(-11, 11, 1, 4, 1, 50.0, 0.0, 1.0), + vector(22, 0.0, 0.0), # massless + vector(23, 91.0, 2.5), + vector(97, 60.0, 1.0), # below the threshold of the default masses + vector(98, 200.0, 0.0), # no width in the model + vector(99, 200.0, 0.0), # zero width on the parameter card + ] + parameters = {(11, "mass"): "ML", (-11, "mass"): "ML", (99, "width"): "WX"} + parameters |= {(pdg, kind): kind for pdg in (23, 97, 98, 99) for kind in ("mass",)} + parameters |= {(pdg, "width"): "W" for pdg in (23, 97)} + + def leg(label, pdg, outgoing, side): + return SimpleNamespace( + label=label, pdg=pdg, outgoing_pdg=outgoing, is_final=side == "final" + ) + + ir = SimpleNamespace( + color_accuracy="lc", + legs=( + leg(1, 11, -11, "initial"), + leg(2, -11, 11, "initial"), + leg(3, 11, 11, "final"), + leg(4, -11, -11, "final"), + ), + ) + found = s_channel_resonances(ir, _Model(particles, parameters)) + assert found == [{"pdg": pdg, "labels": [3, 4], "color": 1} for pdg in (23, 97, 99)] diff --git a/tests/acceptance_tests/test_auto_width.py b/tests/acceptance_tests/test_auto_width.py index 9dd218d631..e8ce737276 100644 --- a/tests/acceptance_tests/test_auto_width.py +++ b/tests/acceptance_tests/test_auto_width.py @@ -32,7 +32,7 @@ import madgraph.interface.master_interface as mgcmd import madgraph.models.check_param_card as check_param_card -from madgraph.iolibs.template_files.mg7 import bootstrap as mg7_bootstrap +from madgraph.mg7 import bootstrap as mg7_bootstrap pjoin = os.path.join @@ -214,7 +214,7 @@ def test_options_of_the_widths(self): with open(card, 'w') as stream: stream.write(text) - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch os.chdir(run_dir) mg7_launch.compute_auto_widths(card) os.chdir(self.cwd) diff --git a/tests/acceptance_tests/test_cmd_madevent.py b/tests/acceptance_tests/test_cmd_madevent.py index 4221b8e783..44a0299d90 100755 --- a/tests/acceptance_tests/test_cmd_madevent.py +++ b/tests/acceptance_tests/test_cmd_madevent.py @@ -1523,7 +1523,7 @@ def test_relaunch_switch_defaults_mg7(self): mg.exec_cmd('output mg7 %s' % out_dir) launcher = importlib.import_module( - 'madgraph.iolibs.template_files.mg7.launch') + 'madgraph.mg7.launch') tool_cards = ('pythia8_card.dat', 'madspin_card.dat', 'delphes_card.dat', 'reweight_card.dat', 'rivet_card.dat', 'madanalysis5_parton_card.dat', diff --git a/tests/acceptance_tests/test_madspace_examples.py b/tests/acceptance_tests/test_madspace_examples.py index 7983ba7b14..ef95ab345b 100644 --- a/tests/acceptance_tests/test_madspace_examples.py +++ b/tests/acceptance_tests/test_madspace_examples.py @@ -55,7 +55,8 @@ # in the subprocess env run_doc_example() builds. os.environ.setdefault('KMP_DUPLICATE_LIB_OK', 'TRUE') -_EXAMPLES_DIR = pjoin(MG5DIR, 'docs', 'source', 'madspace', 'examples') +_DOCS_DIR = pjoin(MG5DIR, 'docs', 'source') +_EXAMPLES_DIR = pjoin(_DOCS_DIR, 'madspace', 'examples') _MADSPACE_INSTALL = pjoin(MG5DIR, 'madspace', 'install') _CODE_BLOCK_RE = re.compile(r'^(?P[ \t]*)\.\.[ \t]+code-block::[ \t]*python[ \t]*$') @@ -106,6 +107,30 @@ def extract_python_blocks(rst_path): return blocks +def extract_named_block(rst_path, name): + """Return the dedented body of the code block whose directive carries + ``:name: ``, or None if there is none.""" + with open(rst_path) as f: + lines = f.read().splitlines() + for i, line in enumerate(lines): + if line.strip() != ':name: %s' % name: + continue + j = i + while not lines[j].lstrip().startswith('.. code-block::'): + j -= 1 + base_indent = len(lines[j]) - len(lines[j].lstrip()) + k = i + 1 + while k < len(lines) and lines[k].strip().startswith(':'): + k += 1 + body = [] + for line in lines[k:]: + if line.strip() and len(line) - len(line.lstrip()) <= base_indent: + break + body.append(line) + return textwrap.dedent('\n'.join(body)).strip() + return None + + def has_madspace(): try: import madspace @@ -147,7 +172,7 @@ def has_madnis(): return False -def run_doc_example(test, page_name, cwd=None): +def run_doc_example(test, page_name, cwd=None, examples_dir=_EXAMPLES_DIR): """Extract the python blocks of ``examples/.rst``, concatenate them into one script, and run it as a subprocess. Fails the test if the script raises or if the page has no python blocks at all (a moved or @@ -157,7 +182,7 @@ def run_doc_example(test, page_name, cwd=None): ``ggttg_process_dir`` below) run with that directory as its working directory, matching what the page's shell commands build. Without it, a fresh empty temporary directory is used and cleaned up afterwards.""" - rst_path = pjoin(_EXAMPLES_DIR, page_name + '.rst') + rst_path = pjoin(examples_dir, page_name + '.rst') if not os.path.isfile(rst_path): test.fail('documentation page not found: %s' % rst_path) @@ -367,5 +392,44 @@ def test_madspace_example_flow_training(self): run_doc_example(self, 'flow-training', cwd=scratch) +class TestGridpackExamples(unittest.TestCase): + + def test_gridpack_example_event_stream(self): + """docs/source/event_stream.rst -- creating a fix_max_weight gridpack + for p p > e+ e- with the MadGraph commands of the page, then filling + an m_ee histogram from its bin/event_stream.py. Needs madspace, numpy, + a C++ compiler and lhapdf/NNPDF40_lo_as_01180.""" + if not has_madspace(): + self.skipTest('madspace unavailable') + if not has_pdf_set('NNPDF40_lo_as_01180'): + self.skipTest('NNPDF40_lo_as_01180 LHAPDF data not found ' + '(set $LHAPDF_DATA_PATH)') + rst_path = pjoin(_DOCS_DIR, 'event_stream.rst') + commands = extract_named_block(rst_path, 'event-stream-mg5-commands') + self.assertTrue(commands, 'MadGraph commands not found in %s' % rst_path) + + with tempfile.TemporaryDirectory(prefix='doc_event_stream_') as scratch: + script = pjoin(scratch, 'mg5_commands.txt') + with open(script, 'w') as f: + f.write('set automatic_html_opening False --no_save\n') + f.write(commands + '\n') + env = dict(os.environ) + if os.path.isdir(_MADSPACE_INSTALL): + env['PYTHONPATH'] = os.pathsep.join( + [_MADSPACE_INSTALL] + + ([env['PYTHONPATH']] if env.get('PYTHONPATH') else []) + ) + result = subprocess.run( + [sys.executable, pjoin(MG5DIR, 'bin', 'madgraph'), '-f', script], + cwd=scratch, env=env, capture_output=True, text=True, + ) + self.assertEqual( + result.returncode, 0, + 'MadGraph commands failed\n--- stdout ---\n%s\n--- stderr ---\n%s' + % (result.stdout[-5000:], result.stderr[-5000:]), + ) + run_doc_example(self, 'event_stream', cwd=scratch, examples_dir=_DOCS_DIR) + + if __name__ == '__main__': unittest.main() diff --git a/tests/acceptance_tests/test_mg7_amplicol.py b/tests/acceptance_tests/test_mg7_amplicol.py new file mode 100644 index 0000000000..598fe0f0aa --- /dev/null +++ b/tests/acceptance_tests/test_mg7_amplicol.py @@ -0,0 +1,246 @@ +################################################################################ +# +# 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 (in pyamplicol/install) and +LHAPDF; the tests skip otherwise. Knobs: MG7_AMPLICOL_EVENTS (default 10000) +and MG7_AMPLICOL_NSIGMA (default 4). + + ./tests/test_manager.py test_mg7_amplicol -pA -t0 +""" + +import glob +import gzip +import 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, gridpack=False, output_format=None, **amplicol): + card = open(path).read() + if gridpack: + card = re.sub(r'(?m)^save_gridpack = false', 'save_gridpack = true', card) + card = re.sub(r'(?m)^compress = true', 'compress = false', card) + # with events, for the cross section of the run + card = re.sub(r'(?m)^run_mode = ".*"', 'run_mode = "regular"', card) + if output_format: + card = re.sub(r'(?m)^output_format = ".*"', 'output_format = "%s"' % + output_format, card) + card = card.replace('fixed_ren_scale = false', 'fixed_ren_scale = true') + card = card.replace('fixed_fact_scale = false', 'fixed_fact_scale = true') + card = re.sub(r'(?m)^(pdf[12]) = ".*"$', r'\1 = "%s"' % _PDF, card) + 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 + + +def _events(run_dir): + """The LHE events of the last run.""" + events = sorted(glob.glob(pjoin(run_dir, 'Events', '*', 'events.lhe.gz'))) + assert events, run_dir + return events[-1] + + +def _resonances(lhe_path): + """The (pdg, mass) of the status-2 particles of every event.""" + with gzip.open(lhe_path, 'rt') as stream: + text = stream.read() + return [[(int(particle[0]), float(particle[10])) + for particle in (line.split() for line in event.strip().splitlines()[1:]) + if len(particle) > 10 and particle[1] == '2'] + for event in re.findall(r'\n(.*?)', text, re.S)] + + +class MG7AmplicolTest(unittest.TestCase): + + def setUp(self): + from madgraph.mg7 import bootstrap + # not imported: Symbolica would hold its lock against the pyamplicol worker + missing = [name for name, installed in ( + ('madspace', bootstrap.madspace_is_installed()), + ('pyamplicol', bootstrap.pyamplicol_is_installed())) if not installed] + if missing: + self.skipTest('%s not installed' % ' and '.join(missing)) + if not misc.resolve_lhapdf().find_set(_PDF): + self.skipTest('%s LHAPDF data not found' % _PDF) + self.path = tempfile.mkdtemp(prefix='mg7_amplicol_') + + def tearDown(self): + shutil.rmtree(self.path, ignore_errors=True) + + def run_mg7(self, name, processes, gridpack=False, **amplicol): + """Output and run the processes, returns the cross section, its error + and the run directory.""" + run_dir = pjoin(self.path, name) + 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'), gridpack, **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'] + # a failed generation can still exit with 0, after the survey + self.assertTrue(result['done'], open(log).read()[-3000:]) + return float(result['mean']), float(result['error']), run_dir + + def assert_compatible(self, first, second, factor=1.): + (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): + """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 assert_resonances(self, lhe_path, pdg, mass, width): + """Some events write the resonance, within bw_cutoff widths of its + mass, and no other one.""" + self.assertTrue(_color_balanced(lhe_path)) + resonances = [r for event in _resonances(lhe_path) for r in event] + self.assertGreater(len(resonances), 0.2 * _EVENTS) + self.assertEqual({r[0] for r in resonances}, {pdg}) + self.assertLess(max(abs(m - mass) for _, m in resonances), 15 * width) + + def test_pp_eej(self): + """Colour singlets attached to the quark line in either order and a Z + resonance: three phase-space channels per colour order, weighted by + their mapping densities, and the Z in the event record.""" + madmatrix = self.run_mg7('mm', ['p p > e+ e- j'])[:2] + amplicol = self.run_mg7('amp', ['p p > e+ e- j [amplicol]'], + output_format='lhe') + self.assert_compatible(amplicol[:2], madmatrix) + self.assert_resonances(_events(amplicol[2]), 23, 91.188, 2.441404) + + def test_pp_wpj(self): + """A charged resonance: W+, never W-, from the charge of its decay.""" + madmatrix = self.run_mg7('mm', ['p p > e+ ve j'])[:2] + amplicol = self.run_mg7('amp', ['p p > e+ ve j [amplicol]'], + output_format='lhe') + self.assert_compatible(amplicol[:2], madmatrix) + self.assert_resonances(_events(amplicol[2]), 24, 80.419, 2.0476) + + def test_mixed(self): + """[amplicol] and madmatrix processes in one output.""" + jj = self.run_mg7('jj', ['p p > j j'])[:2] + 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)) + + def test_gridpack(self): + """A gridpack of [amplicol] and madmatrix processes, with the second + pass unweighting to full colour, reproduces the cross section of its + run. The [amplicol] process has several phase-space channels per colour + order and writes its Z resonance.""" + mean, error, run_dir = self.run_mg7( + 'grid', ['p p > e+ e- j [amplicol]', 'p p > mu+ mu- j'], gridpack=True, + mode='unweight') + gridpack = sorted(glob.glob(pjoin(run_dir, 'Events', '*', 'gridpack')))[-1] + self.assertTrue(glob.glob(pjoin(gridpack, 'SubProcesses', 'A*', 'lc'))) + log = pjoin(self.path, 'gridpack.log') + with open(log, 'w') as stream: + status = subprocess.call( + [sys.executable, pjoin(gridpack, 'bin', 'generate_events'), + '--seed', '11', '--output_dir', pjoin(self.path, 'grid_events')], + cwd=self.path, stdout=stream, stderr=subprocess.STDOUT) + self.assertEqual(status, 0, open(log).read()[-3000:]) + with open(pjoin(self.path, 'grid_events', 'info.json')) as stream: + result = json.load(stream)['process'] + self.assertTrue(result['done'], open(log).read()[-3000:]) + self.assert_compatible((float(result['mean']), float(result['error'])), + (mean, error)) + self.assert_resonances(pjoin(self.path, 'grid_events', 'events.lhe.gz'), + 23, 91.188, 2.441404) diff --git a/tests/acceptance_tests/test_mg7_custom_functions.py b/tests/acceptance_tests/test_mg7_custom_functions.py index 72fc3d0ed0..9fbcb8034d 100644 --- a/tests/acceptance_tests/test_mg7_custom_functions.py +++ b/tests/acceptance_tests/test_mg7_custom_functions.py @@ -208,7 +208,7 @@ def _import_madspace(): """The madspace of this MadGraph, as bin/generate_events imports it, or None. Done before anything imports another madspace from the path.""" try: - from madgraph.iolibs.template_files.mg7.bootstrap import ( + from madgraph.mg7.bootstrap import ( ensure_madspace, drop_install_path) ensure_madspace(interactive=False) import madspace @@ -508,7 +508,7 @@ def test_init_custom_func_languages_mg7(self): self.skipTest('madspace with custom functions unavailable') compilers = {'c': 'cc', 'cpp': 'c++', 'fortran': 'gfortran', 'cuda': 'nvcc', 'hip': 'hipcc'} - from madgraph.iolibs.template_files.mg7 import init_custom_func + from madgraph.mg7 import init_custom_func tested = [] for language, compiler in compilers.items(): diff --git a/tests/unit_tests/interface/test_amplicol_syntax.py b/tests/unit_tests/interface/test_amplicol_syntax.py new file mode 100644 index 0000000000..9ee4f4945c --- /dev/null +++ b/tests/unit_tests/interface/test_amplicol_syntax.py @@ -0,0 +1,329 @@ +################################################################################ +# +# 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]') + self.assertRaises(InvalidCmd, split, 'p p > j j [QCD] [amplicol]') + # the pre-check must not route it to aMC@NLO with 'amplicol' as order + process_type = MGCmd.MasterCmd.extract_process_type + self.assertRaises(InvalidCmd, process_type, 'p p > t t~ [QCD] [amplicol]') + + def test_process_type(self): + process_type = MGCmd.MasterCmd.extract_process_type + 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) + + +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~'])) + # every initial state, and every other particle of each final-state leg + self.assertEqual(len(processes), 81 + 2 * 8) + + def test_coupling_orders(self): + """The orders MG selects by default, as the largest non-QCD orders.""" + 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) + + def test_default_histograms_see_every_final_state_particle(self): + """The representative processes cover every particle of a final-state + leg: the default histograms of p p > e+ e- j j are of jets, which + every subprocess has, not of gluons, which most of them lack.""" + import madgraph.iolibs.export_amplicol as export_amplicol + import madgraph.various.banner as banner_mod + self.cmd.exec_cmd('generate p p > e+ e- j j [amplicol]', printcmd=False, + precmd=True) + processes = export_amplicol.representative_processes( + self.cmd._curr_amplicol_defs, self.cmd._curr_model) + run_card = banner_mod.RunCardMG7() + run_card.create_default_for_process( + {'ninitial': 2, 'nexternal': 6}, '', [processes]) + histograms = list(run_card.dynamic_sections['histograms']) + self.assertIn('jet_2-pt', histograms) + self.assertFalse([name for name in histograms if name.startswith('g_')]) + + def test_mirrored_beams(self): + """pyamplicol keeps one process per unordered initial state; the other + beam order is integrated as the mirror where the definition allows it.""" + 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)) + + def test_selector_rows_are_folded_into_weighted_options(self): + """The rows of a UMAMI selector share its evaluation: they become the + options of one flavour per mirror flag, written in proportion to their + factors, which add up to the flavour factor.""" + import madgraph.iolibs.export_amplicol as export_amplicol + metadata = { + 'pdg_ids': [[1, -1, 2, -2, 21], [1, -1, 3, -3, 21], + [1, -1, 4, -4, 21], [-1, 1, 3, -3, 21]], + 'provider': {'incoming_count': 2, 'masses': [0.] * 5, + 'supports_alpha_s': True}, + 'channels': [{ + 'phasespace_order': [0, 2, 3, 4, 1], + 'phasespace_channels': [ + {'phasespace_order': [0, 2, 3, 4, 1], 'resonances': []}], + 'processes': [ + {'multichannels': [0], 'color_flows': 0, 'matrix_elements': [ + {'pdg_ids': 0, 'factor': 1}, {'pdg_ids': 2, 'factor': 2}, + {'pdg_ids': 1, 'factor': 1}]}, + {'multichannels': [0], 'color_flows': 1, 'matrix_elements': [ + {'pdg_ids': 3, 'factor': 1}]}, + ], + }], + 'color_flows': [], 'helicities': [], 'runtime_processes': [], + } + partons = {21, 1, 2, 3, 4, -1, -2, -3, -4} + entries = export_amplicol.subprocess_entries( + metadata, self.cmd._curr_model, 'SubProcesses/A1', 'me.so', 'ratio.so', + 3, [partons, partons]) + self.assertEqual(len(entries), 1) + flavors = entries[0]['flavors'] + # d d~ > s s~ g has its mirror exported, the other rows are mirrored + common = {'channels': [0], 'resonance_pdgs': [[]]} + self.assertEqual(flavors, [ + {'index': 0, 'options': [[1, -1, 2, -2, 21], [1, -1, 4, -4, 21]], + 'weights': [1, 2], 'factor': 3, 'mirror': True, 'color_index': 0, + **common}, + {'index': 0, 'options': [[1, -1, 3, -3, 21]], 'weights': [1], + 'factor': 1, 'mirror': False, 'color_index': 0, **common}, + {'index': 1, 'options': [[-1, 1, 3, -3, 21]], 'weights': [1], + 'factor': 1, 'mirror': False, 'color_index': 1, **common}, + ]) + + def test_resonance_charges_split_the_selector_rows(self): + """A selector integrated over phase-space channels with resonances + writes them with the charge of the decay products of each row: the + charge-conjugated rows of a W selector become flavours of their own.""" + import madgraph.iolibs.export_amplicol as export_amplicol + metadata = { + 'pdg_ids': [[2, -1, -11, 12, 21], [4, -3, -11, 12, 21], + [1, -2, 11, -12, 21]], + 'provider': {'incoming_count': 2, 'masses': [0.] * 5, + 'supports_alpha_s': True}, + 'channels': [{ + 'phasespace_order': [0, 4, 1, 2, 3], + 'phasespace_channels': [ + {'phasespace_order': [0, 4, 1, 2, 3], 'resonances': []}, + {'phasespace_order': [0, 4, 1, 2, 3], + 'resonances': [{'pdg': 24, 'legs': [2, 3]}]}], + 'colored_resonances': [{'pdg': 6, 'legs': [2, 3, 4], 'color': 3}], + 'processes': [ + {'multichannels': [0, 1], 'color_flows': 0, + 'matrix_elements': [{'pdg_ids': index, 'factor': 1} + for index in range(3)]}, + ], + }], + 'color_flows': [], 'helicities': [], 'runtime_processes': [], + } + partons = {21, 1, 2, 3, 4, -1, -2, -3, -4} + entries = export_amplicol.subprocess_entries( + metadata, self.cmd._curr_model, 'SubProcesses/A1', 'me.so', 'ratio.so', + 3, [partons, partons]) + flavors = entries[0]['flavors'] + self.assertEqual([f['options'] for f in flavors], [ + [[2, -1, -11, 12, 21], [4, -3, -11, 12, 21]], [[1, -2, 11, -12, 21]]]) + self.assertEqual([f['resonance_pdgs'] for f in flavors], + [[[], [24]], [[], [-24]]]) + self.assertEqual(entries[0]['phasespace_channels'][1], + {'order': [0, 4, 1, 2, 3], + 'resonances': [{'pdg': 24, 'legs': [2, 3]}]}) + # the launch rejects a coloured resonance that goes on shell + self.assertEqual(entries[0]['colored_resonances'], + [{'pdg': 6, 'legs': [2, 3, 4], 'color': 3}]) + + +class AmplicolLaunchResonancesTest(unittest.TestCase): + """The s-channel resonances of an [amplicol] subprocess are fixed at output; + the launch keeps those that go on shell with the param card values.""" + + # q q~ > e- e+ mu- mu+ through Z Z, Z and H > Z Z + order = [0, 2, 3, 4, 5, 1] + z_e, z_mu = {'pdg': 23, 'legs': [2, 3]}, {'pdg': 23, 'legs': [4, 5]} + h = {'pdg': 25, 'legs': [2, 3, 4, 5]} + + def setUp(self): + try: + from madgraph.mg7 import subprocess_amplicol + except ImportError: + self.skipTest('madspace is not installed') + self.module = subprocess_amplicol + self.meta = { + 'incoming': [2, -2], 'outgoing': [11, -11, 13, -13], + 'phasespace_channels': [ + {'order': self.order, 'resonances': []}, + {'order': self.order, 'resonances': [self.z_e, self.z_mu]}, + {'order': self.order, 'resonances': [self.h, self.z_e]}, + {'order': self.order, 'resonances': [self.h]}, + ], + 'colored_resonances': [{'pdg': 6, 'legs': [2, 3], 'color': 3}], + 'flavors': [ + {'channels': [0, 1, 2, 3], + 'resonance_pdgs': [[], [23, 23], [25, 23], [25]]}, + {'channels': [3, 1], 'resonance_pdgs': [[25], [23, 23]]}, + ], + } + + def active(self, masses, widths): + return self.module.active_resonances( + self.meta, lambda pid: masses.get(pid, 0.), lambda pid: widths.get(pid, 0.)) + + def test_all_resonances_on_shell(self): + meta = self.active({23: 91., 25: 125.}, {23: 2.5, 25: 4e-3}) + self.assertEqual(meta['phasespace_channels'], self.meta['phasespace_channels']) + self.assertEqual(meta['flavors'], self.meta['flavors']) + + def test_stable_resonance_is_dropped_and_channels_merged(self): + # without a Higgs width, H > Z Z becomes Z and H the non-resonant channel + meta = self.active({23: 91., 25: 125.}, {23: 2.5}) + self.assertEqual(meta['phasespace_channels'], [ + {'order': self.order, 'resonances': []}, + {'order': self.order, 'resonances': [self.z_e, self.z_mu]}, + {'order': self.order, 'resonances': [self.z_e]}, + ]) + self.assertEqual([(f['channels'], f['resonance_pdgs']) for f in meta['flavors']], + [([0, 1, 2], [[], [23, 23], [23]]), ([0, 1], [[], [23, 23]])]) + # the LHE diagram masks follow the merged channels + diagrams, masks = self.module.lhe_diagrams(meta) + self.assertEqual(diagrams, [(0, ()), (1, (23, 23)), (2, (23,))]) + self.assertEqual(masks, [[1., 1., 1.], [1., 1., 0.]]) + + def test_resonance_below_its_decay_threshold_is_dropped(self): + # mu+ mu- of 2 x 70 GeV: neither Z > mu+ mu- nor H > 4l are on shell + meta = self.active({13: 70., 23: 91., 25: 125.}, {23: 2.5, 25: 4e-3}) + self.assertEqual( + [c['resonances'] for c in meta['phasespace_channels']], [[], [self.z_e]]) + self.assertEqual([f['channels'] for f in meta['flavors']], [[0, 1], [0, 1]]) + self.assertEqual(self.meta['flavors'][1]['channels'], [3, 1]) + + def test_coloured_resonance_on_shell_is_an_error(self): + masses, widths = {6: 173., 23: 91.}, {23: 2.5} + self.active(masses, widths) + widths[6] = 1.5 + with self.assertRaisesRegex(RuntimeError, 'coloured s-channel resonance 6'): + self.active(masses, widths) diff --git a/tests/unit_tests/interface/test_auto_width_engine.py b/tests/unit_tests/interface/test_auto_width_engine.py index 561a63b019..3021a97bd8 100644 --- a/tests/unit_tests/interface/test_auto_width_engine.py +++ b/tests/unit_tests/interface/test_auto_width_engine.py @@ -36,7 +36,7 @@ import madgraph.core.base_objects as base_objects import madgraph.interface.madevent_interface as madevent_interface import madgraph.interface.master_interface as mgcmd -from madgraph.iolibs.template_files.mg7 import bootstrap as mg7_bootstrap +from madgraph.mg7 import bootstrap as mg7_bootstrap from madgraph.models.model_reader import Alphas_Runner pjoin = os.path.join @@ -445,7 +445,7 @@ def tearDown(self): shutil.rmtree(self.tmp, ignore_errors=True) def write_grid(self, mass, alpha_s_mz=0.118): - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch process = object.__new__(mg7_launch.MadgraphProcess) process.param_card = mock.Mock() process.param_card.get_value = lambda block, code: alpha_s_mz @@ -544,7 +544,7 @@ class DecayFlavorSplitTest(unittest.TestCase): } def setUp(self): - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch self.launch = mg7_launch def test_one_subprocess_per_flavour(self): diff --git a/tests/unit_tests/interface/test_auto_width_nlo.py b/tests/unit_tests/interface/test_auto_width_nlo.py index 173164d601..bc20c37be8 100644 --- a/tests/unit_tests/interface/test_auto_width_nlo.py +++ b/tests/unit_tests/interface/test_auto_width_nlo.py @@ -42,7 +42,7 @@ import madgraph.models.check_param_card as check_param_card import madgraph.models.import_ufo as import_ufo import madgraph.models.model_reader as model_reader -from madgraph.iolibs.template_files.mg7 import bootstrap as mg7_bootstrap +from madgraph.mg7 import bootstrap as mg7_bootstrap pjoin = os.path.join @@ -858,7 +858,8 @@ def tearDown(self): shutil.rmtree(self.tmp, ignore_errors=True) def test_one_compute_widths_per_set_of_options(self): - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch + from madgraph.mg7 import model_widths card = pjoin(self.tmp, 'param_card.dat') with open(card, 'w') as stream: stream.write('\n'.join([ @@ -871,7 +872,7 @@ def run(command, *args, **kwargs): scripts.append(stream.read()) return mock.Mock(returncode=0) - with mock.patch.object(mg7_launch, 'read_stored_model', lambda: 'sm'), \ + with mock.patch.object(model_widths, 'read_stored_model', lambda: 'sm'), \ mock.patch('subprocess.run', run): mg7_launch.compute_auto_widths(card) self.assertEqual(len(scripts), 1) @@ -884,7 +885,7 @@ def run(command, *args, **kwargs): 'compute_widths 24 --path=CARD ']) def test_no_auto_no_run(self): - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch card = pjoin(self.tmp, 'param_card.dat') with open(card, 'w') as stream: stream.write('DECAY 6 1.5\n') @@ -894,7 +895,7 @@ def test_no_auto_no_run(self): def compile_args(self, subprocesses, extra=()): """The arguments of the make that builds ``subprocesses`` directories.""" - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch process = object.__new__(mg7_launch.MadgraphProcess) process.run_card = {'run': {'device': ['cpu'], 'cpu_mode': 'simd_128', 'precision': 'color32', 'cpu_thread_pool_size': 2}} @@ -918,7 +919,7 @@ def test_the_headers_are_precompiled_from_four_subprocesses(self): def test_only_what_is_to_build_counts(self): """Six directories, but the libraries of three of them exist: below the threshold.""" - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch for i in range(3): open(pjoin(self.tmp, 'libmadmatrix_P%d_x_simd_128.so' % i), 'w').close() process = object.__new__(mg7_launch.MadgraphProcess) diff --git a/tests/unit_tests/interface/test_mg7_interference_helicity.py b/tests/unit_tests/interface/test_mg7_interference_helicity.py index 60c4cf2a71..2354a86bdd 100644 --- a/tests/unit_tests/interface/test_mg7_interference_helicity.py +++ b/tests/unit_tests/interface/test_mg7_interference_helicity.py @@ -46,7 +46,7 @@ class TestInterferenceHelicity(unittest.TestCase): def setUp(self): try: - import madgraph.iolibs.template_files.mg7.launch as launch + import madgraph.mg7.launch as launch except ImportError as error: # the mg7 runtime (madspace) is not there self.skipTest('mg7 launch not importable: %s' % error) self.launch = launch diff --git a/tests/unit_tests/interface/test_mg7_launch.py b/tests/unit_tests/interface/test_mg7_launch.py index 97d288c4a9..6d8b4b443b 100644 --- a/tests/unit_tests/interface/test_mg7_launch.py +++ b/tests/unit_tests/interface/test_mg7_launch.py @@ -41,9 +41,9 @@ import madgraph.interface.extended_cmd as ext_cmd import madgraph.interface.master_interface as mgcmd -from madgraph.iolibs.template_files.mg7 import bootstrap as mg7_bootstrap +from madgraph.mg7 import bootstrap as mg7_bootstrap -LAUNCH_MODULE = 'madgraph.iolibs.template_files.mg7.launch' +LAUNCH_MODULE = 'madgraph.mg7.launch' _MG_ROOT = str(mg7_bootstrap.MG_ROOT) @@ -116,7 +116,7 @@ def MG7Cmd(me_dir='.', options=None): # `from pkg import launch` resolves through the parent package's # attribute when the real module has already been imported (which # MG7CmdTest does), so sys.modules alone is not enough to stub it. - import madgraph.iolibs.template_files.mg7 as mg7_pkg + import madgraph.mg7 as mg7_pkg self.mg7_pkg = mg7_pkg self._saved_attr = getattr(mg7_pkg, 'launch', None) mg7_pkg.launch = stub @@ -480,7 +480,7 @@ class MG7CmdTest(unittest.TestCase): """The real run interface (needs a compiled madspace to import).""" def setUp(self): - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch self.launch = mg7_launch self.tmpdir = tempfile.TemporaryDirectory() self.me_dir = make_mg7_dir(os.path.join(self.tmpdir.name, 'PROC')) @@ -633,7 +633,7 @@ def test_no_histograms_no_plot_directory(self): def test_missing_matplotlib_does_not_fail_the_run(self): """the numbers are in info.json: a run without the backend carries on""" - from madgraph.iolibs.template_files.mg7 import plots + from madgraph.mg7 import plots saved = plots._pyplot def no_matplotlib(): @@ -990,14 +990,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.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) @@ -1028,7 +1029,7 @@ def tearDown(self): self.logger.setLevel(self.saved) def test_info_is_dropped_during_setup(self): - from madgraph.iolibs.template_files.mg7.run_interface import _quiet_setup + from madgraph.mg7.run_interface import _quiet_setup with _quiet_setup(): self.logger.info('W E L C O M E to') @@ -1039,14 +1040,14 @@ def test_info_is_dropped_during_setup(self): def test_warnings_still_get_through(self): """Quiet is not silent: anything that needs saying still says it.""" - from madgraph.iolibs.template_files.mg7.run_interface import _quiet_setup + from madgraph.mg7.run_interface import _quiet_setup with _quiet_setup(): self.logger.warning('something worth knowing') self.assertIn('something worth knowing', self.captured.getvalue()) def test_the_level_is_restored_afterwards(self): - from madgraph.iolibs.template_files.mg7.run_interface import _quiet_setup + from madgraph.mg7.run_interface import _quiet_setup with _quiet_setup(): pass @@ -1054,7 +1055,7 @@ def test_the_level_is_restored_afterwards(self): self.assertIn('back to normal', self.captured.getvalue()) def test_it_restores_the_level_after_a_failure(self): - from madgraph.iolibs.template_files.mg7.run_interface import _quiet_setup + from madgraph.mg7.run_interface import _quiet_setup try: with _quiet_setup(): @@ -1077,7 +1078,7 @@ class TestPostProcessingIsSkippedWhenThereIsNothingToDo(unittest.TestCase): def tools_for(self, switch): """The tool list run_selected_tools builds, without running it.""" - from madgraph.iolibs.template_files.mg7 import launch + from madgraph.mg7 import launch return launch.selected_tools(switch) def test_a_not_available_switch_selects_no_tool(self): @@ -1111,7 +1112,7 @@ def test_run_selected_tools_returns_before_building_anything(self): """With no tool selected it must not construct MG7RunCmd -- that is what printed the MADEVENT banner and the configuration lines.""" - from madgraph.iolibs.template_files.mg7 import launch + from madgraph.mg7 import launch class _Process(object): run_path = '/nonexistent/run_01' @@ -1147,7 +1148,7 @@ def tearDown(self): def test_store_result_bookkeeping_writes_nothing(self): import madgraph.various.banner as banner_mod - from madgraph.iolibs.template_files.mg7.run_interface import MG7RunCmd + from madgraph.mg7.run_interface import MG7RunCmd # load_results_db only needs me_dir, process and the model name # (__del__ reads the other two); building a real MG7RunCmd needs a @@ -1189,7 +1190,7 @@ class MG7MeFrameTest(unittest.TestCase): def resolve(self, me_frame, is_decay, particle_count=4, outgoing=None): """``outgoing``, if given, lists the final states of several subprocesses of one output (g g > each of them)""" - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch process = object.__new__(mg7_launch.MadgraphProcess) process.run_card = {'run': {'me_frame': me_frame}} process.is_decay = is_decay @@ -1276,15 +1277,17 @@ class TestForceLHEOutput(unittest.TestCase): the LHE format only when something that reads LHE will run.""" def setUp(self): - from madgraph.iolibs.template_files.mg7 import launch + import madgraph + from madgraph.mg7 import launch from madgraph.various.banner import RunCardMG7 self.launch = launch self.tmpdir = tempfile.TemporaryDirectory() self.cwd = os.getcwd() os.chdir(self.tmpdir.name) os.mkdir('Cards') - template = os.path.join(os.path.dirname(launch.__file__), - 'run_card.toml') + template = os.path.join(madgraph.MG5DIR, 'madgraph', 'iolibs', + 'template_files', 'mg7', 'run_card.toml') + self.assertTrue(os.path.isfile(template), template) RunCardMG7().write(self.card, template=template) def tearDown(self): @@ -1364,7 +1367,7 @@ class TestNpyToLHE(unittest.TestCase): def setUp(self): import numpy as np - from madgraph.iolibs.template_files.mg7 import npy_to_lhe + from madgraph.mg7 import npy_to_lhe self.np = np self.npy_to_lhe = npy_to_lhe self.ms = npy_to_lhe._madspace() @@ -1462,14 +1465,14 @@ def test_defaults(self): def test_mean_ignores_under_and_overflow(self): """Both histogram kinds carry the under/overflow bins first and last; they have no position and must not shift the in-range bins.""" - from madgraph.iolibs.template_files.mg7 import launch + from madgraph.mg7 import launch mean = launch.MadgraphProcess._histogram_mean # bins [0,10) and [10,20) with weights 1 and 3, plus 100 in overflow self.assertAlmostEqual(mean(0., 20., [5., 1., 3., 100.]), 12.5) self.assertIsNone(mean(0., 20., [5., 0., 0., 100.])) def test_postprocessing_off_builds_no_event_histograms(self): - from madgraph.iolibs.template_files.mg7 import launch + from madgraph.mg7 import launch from madgraph.various.banner import RunCardMG7 process = launch.MadgraphProcess.__new__(launch.MadgraphProcess) process.run_card = RunCardMG7() @@ -1489,7 +1492,7 @@ class MG7ObservablePidsTest(unittest.TestCase): flavours, consistent over every flavour a subprocess groups together.""" def setUp(self): - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch self.launch = mg7_launch @staticmethod @@ -1535,7 +1538,7 @@ class MG7DecayProductsTest(unittest.TestCase): on-shell (decay-chain) propagator decays into, as MadEvent's check_decay.""" def setUp(self): - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch self.launch = mg7_launch @staticmethod @@ -1596,7 +1599,7 @@ class MG7AutoPtjMjjCutNamesTest(unittest.TestCase): OUTGOING_PDGS = [21, 21, 25, 1, -1] def setUp(self): - from madgraph.iolibs.template_files.mg7 import launch as mg7_launch + from madgraph.mg7 import launch as mg7_launch self.launch = mg7_launch def make_process(self, cuts, xqcut=20.0, auto=True): diff --git a/tests/unit_tests/interface/test_mg7_mlm_pdf_ladder.py b/tests/unit_tests/interface/test_mg7_mlm_pdf_ladder.py index 414acbe915..647de8452e 100644 --- a/tests/unit_tests/interface/test_mg7_mlm_pdf_ladder.py +++ b/tests/unit_tests/interface/test_mg7_mlm_pdf_ladder.py @@ -31,7 +31,7 @@ class TestMLMPdfLadder(unittest.TestCase): def setUp(self): try: - import madgraph.iolibs.template_files.mg7.launch as launch + import madgraph.mg7.launch as launch except Exception as error: # madspace not built or importable here self.skipTest('mg7 launch not importable: %s' % error) self.launch = launch diff --git a/tests/unit_tests/interface/test_mg7_systematics_summary.py b/tests/unit_tests/interface/test_mg7_systematics_summary.py index de9d6b39c9..7885c96917 100644 --- a/tests/unit_tests/interface/test_mg7_systematics_summary.py +++ b/tests/unit_tests/interface/test_mg7_systematics_summary.py @@ -26,7 +26,7 @@ import inspect import unittest -import madgraph.iolibs.template_files.mg7.systematics_summary as systematics_summary +import madgraph.mg7.systematics_summary as systematics_summary class _Summary(object): @@ -142,7 +142,7 @@ class TestOneImplementation(unittest.TestCase): """Neither caller may grow its own copy back.""" def sources(self): - import madgraph.iolibs.template_files.mg7.launch as launch + import madgraph.mg7.launch as launch return { 'the Systematics box': diff --git a/tests/unit_tests/interface/test_mg7_weighting_passes.py b/tests/unit_tests/interface/test_mg7_weighting_passes.py new file mode 100644 index 0000000000..c6957f6a5b --- /dev/null +++ b/tests/unit_tests/interface/test_mg7_weighting_passes.py @@ -0,0 +1,249 @@ +################################################################################ +# +# Copyright (c) 2026 The MadGraph5_aMC@NLO Development team and Contributors +# +# This file is a part of the MadGraph5_aMC@NLO project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph5_aMC@NLO license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""The [surrogate] and [amplicol] second passes of an mg7 run: the run card +settings of each pass and the checks done at launch.""" + +import importlib.util +import json +import os +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +from madgraph.mg7 import bootstrap as mg7_bootstrap + + +@unittest.skipUnless(mg7_bootstrap.madspace_is_installed(), + 'madspace is not installed') +class FirstPassSettingsTest(unittest.TestCase): + """The first_pass_ settings apply to the subprocesses unweighted again in a + second pass only, the [generation] ones to all others.""" + + def setUp(self): + import madgraph.various.banner as banner_mod + from madgraph.mg7 import madgraph_process + from madgraph.mg7.subprocess_amplicol import AmplicolSubprocess + from madgraph.mg7.subprocess_madmatrix import MadgraphSubprocess + self.ms = madgraph_process.ms + self.process = object.__new__(madgraph_process.MadgraphProcess) + run_card = banner_mod.RunCardMG7() + run_card['generation']['max_overweight_truncation'] = 0.001 + run_card['generation']['freeze_max_weight_after'] = 500 + run_card['amplicol']['mode'] = 'unweight' + run_card['amplicol']['first_pass_max_overweight_truncation'] = 0.02 + run_card['amplicol']['first_pass_freeze_max_weight_after'] = 70 + self.process.run_card = run_card + self.process.surrogate = None + self.process.subprocess_data = [] + step = self.ms.WeightingStep([], mode='unweight') + + amplicol = object.__new__(AmplicolSubprocess) + amplicol.process = self.process + amplicol.ratio_matrix_element = object() + amplicol.weighting_steps = lambda: [step] + madmatrix = object.__new__(MadgraphSubprocess) + madmatrix.process = self.process + madmatrix.surrogates = [] + self.process.subprocesses = [amplicol, madmatrix] + + def test_mixed_output(self): + groups = self.process.weighting_groups([[0, 1], [2]]) + self.assertEqual(len(groups), 1) + self.assertEqual(list(groups[0].channel_indices), [0, 1]) + self.assertAlmostEqual(groups[0].first_pass_max_overweight_truncation, 0.02) + self.assertEqual(groups[0].first_pass_freeze_max_weight_after, 70) + + self.process.init_generator_config() + cfg = self.process.event_generator_config + self.assertAlmostEqual(cfg.max_overweight_truncation, 0.001) + self.assertEqual(cfg.freeze_max_weight_after, 500) + + def test_reweight_mode_has_no_first_pass(self): + self.process.run_card['amplicol']['mode'] = 'reweight' + groups = self.process.weighting_groups([[0, 1], [2]]) + self.assertIsNone(groups[0].first_pass_max_overweight_truncation) + self.assertIsNone(groups[0].first_pass_freeze_max_weight_after) + + +@unittest.skipUnless(mg7_bootstrap.madspace_is_installed(), + 'madspace is not installed') +class SurrogateWithAmplicolTest(unittest.TestCase): + + def test_is_refused(self): + import madgraph.various.banner as banner_mod + from madgraph.mg7 import madgraph_process + process = object.__new__(madgraph_process.MadgraphProcess) + process.run_card = banner_mod.RunCardMG7() + process.run_card['surrogate']['function'] = 'surrogate.py' + process.subprocess_data = [{'me_path': 'a'}, {'backend': 'amplicol'}] + with self.assertRaisesRegex(RuntimeError, r'\[surrogate\] is not supported'): + process.init_surrogate() + # without [amplicol] subprocesses, the next check is reached + process.subprocess_data = process.subprocess_data[:1] + with self.assertRaisesRegex(RuntimeError, 'file not found'): + process.init_surrogate() + + +_TEMPLATES = (Path(__file__).resolve().parents[3] / 'madgraph' / 'iolibs' + / 'template_files' / 'mg7') + + +def _template_module(name): + spec = importlib.util.spec_from_file_location( + f'mg7_template_{name}', _TEMPLATES / f'{name}.py') + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@unittest.skipUnless(mg7_bootstrap.madspace_is_installed(), + 'madspace is not installed') +class GridpackFirstPassTest(unittest.TestCase): + """A gridpack takes the first_pass_ settings from its card, for the groups + that unweight in a later step, also if its groups.json predates them.""" + + def load(self, groups): + setup = _template_module('gridpack_setup') + with tempfile.TemporaryDirectory() as tmp: + os.mkdir(os.path.join(tmp, 'data')) + with open(os.path.join(tmp, 'data', 'groups.json'), 'w') as f: + json.dump(groups, f) + with mock.patch.object(setup, 'gridpack_path', + lambda *parts: os.path.join(tmp, *parts)): + return setup.load_weighting_groups(0.001, 500, 0.02, 70) + + def test_old_and_new_groups(self): + ms = _template_module('gridpack_setup').ms + step = ms.WeightingStep([], mode='unweight') + stored = json.loads(ms.weighting_groups_to_json([ + ms.WeightingGroup([0], [step], 0.5, 3), ms.WeightingGroup([1], [step]), + ])) + reweight = json.loads(ms.weighting_groups_to_json([ + ms.WeightingGroup([2], [ms.WeightingStep([], mode='reweight')], 0.5, 3) + ])) + groups = self.load(stored + reweight) + for group in groups[:2]: + self.assertAlmostEqual(group.first_pass_max_overweight_truncation, 0.02) + self.assertEqual(group.first_pass_freeze_max_weight_after, 70) + self.assertAlmostEqual(group.steps[0].max_overweight_truncation, 0.001) + self.assertIsNone(groups[2].first_pass_max_overweight_truncation) + + +class FakeApi: + def __init__(self, device): + self.device = device + + def device_type(self): + return self.device + + +class FakeContext: + def __init__(self, device='cpu'): + self.device = device + self.loaded = [] + + def load_matrix_element(self, path, param_card, parameters): + self.loaded.append((path, param_card, parameters)) + return FakeApi(self.device) + + +class SurrogateLibraryTest(unittest.TestCase): + """A .so [surrogate]: path substitution, missing files, device and the + run-time parameters of the madmatrix libraries.""" + + def setUp(self): + self.loader = _template_module('custom_function_loader') + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.path = os.path.join(self.tmp.name, 'sur_{subprocess}_{device}.so') + open(self.path.format(subprocess='P1_x', device='cpp512y'), 'w').close() + + def load(self, contexts, devices=('cpu',), exports=True, subprocess='P1_x'): + with mock.patch.object(self.loader, '_exports', lambda path, name: exports): + return self.loader.load_surrogate( + None, contexts, ['cpp512y'] * len(contexts), list(devices), + self.path, subprocess, [[1, 2]], 'card.dat', {'bwcutoff': 15}, + ) + + def test_parameters_are_passed(self): + context = FakeContext() + self.load([context]) + library = self.path.format(subprocess='P1_x', device='cpp512y') + self.assertEqual(context.loaded, [(library, 'card.dat', {'bwcutoff': 15.0})]) + + def test_no_parameters_without_set_parameter(self): + context = FakeContext() + self.load([context], exports=False) + self.assertEqual(context.loaded[0][2], {}) + + def test_missing_library(self): + with self.assertRaisesRegex(RuntimeError, 'sur_P2_y_cpp512y.so not found'): + self.load([FakeContext()], subprocess='P2_y') + + def test_wrong_device(self): + with self.assertRaisesRegex(RuntimeError, 'built for cpu, but the run uses cuda'): + self.load([FakeContext('cpu')], devices=('cuda',)) + + +@unittest.skipUnless(mg7_bootstrap.madspace_is_installed(), + 'madspace is not installed') +class CompileAmplicolTest(unittest.TestCase): + """SubProcesses/check-pyamplicol runs before make, and its message is the + error of the launch.""" + + def setUp(self): + from madgraph.mg7 import madgraph_process + self.module = madgraph_process + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.subproc = os.path.join(self.tmp.name, 'SubProcesses') + os.makedirs(os.path.join(self.subproc, 'P1_gg_ttx')) + self.process = object.__new__(madgraph_process.MadgraphProcess) + self.process.run_path = self.tmp.name + self.process.subprocess_data = [{ + 'backend': 'amplicol', + 'path': os.path.join(self.subproc, 'P1_gg_ttx'), + 'me_path': 'missing.so', 'ratio_path': 'missing.so', + }] + + def write_check(self, body): + path = os.path.join(self.subproc, 'check-pyamplicol') + with open(path, 'w') as f: + f.write('#!/bin/sh\n' + body) + os.chmod(path, 0o755) + + def compile(self): + with mock.patch.object(self.module.MadgraphProcess, 'make_subprocesses') as make: + self.process.compile_amplicol() + return make + + def test_stale_pyamplicol_stops_before_make(self): + self.write_check('echo "reinstall pyamplicol 1.2" >&2\nexit 1\n') + with mock.patch.object(self.module.MadgraphProcess, 'make_subprocesses') as make: + with self.assertRaisesRegex(RuntimeError, 'reinstall pyamplicol 1.2'): + self.process.compile_amplicol() + make.assert_not_called() + + def test_matching_pyamplicol_compiles(self): + self.write_check('exit 0\n') + self.compile().assert_called_once() + + def test_older_output_without_check(self): + self.compile().assert_called_once() + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/unit_tests/interface/test_tutorials.py b/tests/unit_tests/interface/test_tutorials.py index 635601d439..8403ed3802 100644 --- a/tests/unit_tests/interface/test_tutorials.py +++ b/tests/unit_tests/interface/test_tutorials.py @@ -1924,7 +1924,7 @@ def test_the_percentages_come_from_the_shared_helper(self): """Not a third implementation: the run's own box and the scan summary read the same one.""" - from madgraph.iolibs.template_files.mg7 import systematics_summary + from madgraph.mg7 import systematics_summary up, down = systematics_summary.scale_percentages(self.INFO['systematics']) rows = self.module()._systematics_rows(self.INFO) diff --git a/tests/unit_tests/iolibs/test_mg7_discrete_flavor.py b/tests/unit_tests/iolibs/test_mg7_discrete_flavor.py index ec03a08472..4d2316c15a 100644 --- a/tests/unit_tests/iolibs/test_mg7_discrete_flavor.py +++ b/tests/unit_tests/iolibs/test_mg7_discrete_flavor.py @@ -22,14 +22,14 @@ import types import unittest -from madgraph.iolibs.template_files.mg7 import bootstrap as mg7_bootstrap +from madgraph.mg7 import bootstrap as mg7_bootstrap # Importing launch without a compiled madspace starts the madspace build, and # test_manager imports every unit test module while collecting: check first. launch = None if mg7_bootstrap.madspace_is_installed(): try: - import madgraph.iolibs.template_files.mg7.launch as launch + import madgraph.mg7.launch as launch except Exception as error: launch = None import_error = error diff --git a/tests/unit_tests/iolibs/test_mg7_event_stream.py b/tests/unit_tests/iolibs/test_mg7_event_stream.py new file mode 100644 index 0000000000..518b0b0e94 --- /dev/null +++ b/tests/unit_tests/iolibs/test_mg7_event_stream.py @@ -0,0 +1,105 @@ +"""The Python EventStream of the mg7 gridpacks (bin/event_stream.py), on a fake +madspace stream: no event is lost between the ways of reading the stream.""" + +import collections +import importlib.util +import sys +import types +import unittest +from pathlib import Path +from unittest import mock + +import numpy as np + +_TEMPLATE = (Path(__file__).resolve().parents[3] / "madgraph" / "iolibs" + / "template_files" / "mg7" / "event_stream.py") + + +def _event_stream_module(): + setup = types.ModuleType("gridpack_setup") + setup.__getattr__ = lambda name: None + spec = importlib.util.spec_from_file_location("mg7_event_stream", _TEMPLATE) + module = importlib.util.module_from_spec(spec) + with mock.patch.dict(sys.modules, {"madspace": types.ModuleType("madspace"), + "gridpack_setup": setup}): + spec.loader.exec_module(module) + return module + + +Particle = collections.namedtuple("Particle", "pdg_id px") +Event = collections.namedtuple("Event", "weight particles") + + +class FakeStream: + """Event i has weight i and i % 3 + 1 particles.""" + + def __init__(self): + self.count = 0 + + def _events(self, count): + events = [ + Event(float(i), [Particle(21, float(i + j)) for j in range(i % 3 + 1)]) + for i in range(self.count, self.count + count) + ] + self.count += count + return events + + def next_events(self, count): + return self._events(count) + + def close(self): + pass + + def next_batch(self, count): + events = self._events(count) + px = np.zeros((count, 3)) + for row, event in zip(px, events): + row[: len(event.particles)] = [p.px for p in event.particles] + return { + "weight": np.array([e.weight for e in events]), + "particle_count": np.array([len(e.particles) for e in events], np.int32), + "pdg_id": np.where(px != 0.0, 21, 0).astype(np.int32), + "px": px, + } + + +class TestEventStream(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.module = _event_stream_module() + # the fake stream has only some of the LHE fields + cls.module.EventBatch = types.SimpleNamespace + + def make_stream(self): + stream = object.__new__(self.module.EventStream) + stream._stream = FakeStream() + stream._pending = collections.deque() + return stream + + def test_no_event_is_lost_after_leaving_the_loop(self): + stream = self.make_stream() + for event in stream: + if event.weight == 4.0: + break + self.assertEqual([e.weight for e in stream.next_events(3)], [5.0, 6.0, 7.0]) + batch = stream.next_batch(4) + self.assertEqual(list(batch.weight), [8.0, 9.0, 10.0, 11.0]) + self.assertEqual(list(batch.particle_count), [3, 1, 2, 3]) + self.assertEqual(list(batch.px[0]), [8.0, 9.0, 10.0]) + self.assertEqual(list(batch.px[2]), [10.0, 11.0, 0.0]) + self.assertEqual(batch.particle_mask[2].tolist(), [True, True, False]) + self.assertEqual(batch.pdg_id.dtype, np.int32) + # the rest of the chunk, then the stream itself + batch = stream.next_batch(1000) + self.assertEqual(list(batch.weight), [float(i) for i in range(12, 1012)]) + self.assertEqual(next(iter(stream)).weight, 1012.0) + + def test_batches_needs_a_positive_size(self): + stream = self.make_stream() + for size in (0, -1): + with self.assertRaises(ValueError): + next(stream.batches(size, max_events=10)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/iolibs/test_mg7_gridpack_card.py b/tests/unit_tests/iolibs/test_mg7_gridpack_card.py new file mode 100644 index 0000000000..1da76709bd --- /dev/null +++ b/tests/unit_tests/iolibs/test_mg7_gridpack_card.py @@ -0,0 +1,75 @@ +"""Command line defaults of the standalone mg7 gridpack launcher, taken from +its Cards/grid_run_card.toml.""" + +import argparse +import importlib.util +import sys +import types +import unittest +from pathlib import Path +from unittest import mock + +_TEMPLATE = (Path(__file__).resolve().parents[3] / "madgraph" / "iolibs" + / "template_files" / "mg7" / "gridpack.py") + + + +def _run_card(**sections): + """The default run card as plain TOML data, with sections replaced (None + removes one).""" + import madgraph.various.banner as banner_mod + card = banner_mod.RunCardMG7() + data = {key: dict(card[key]) for key in ("run", "generation", "gridpack", "amplicol")} + data["generation"]["max_overweight_truncation"] = 0.001 + for key, value in sections.items(): + if value is None: + del data[key] + else: + data[key] = value + return data + + +class _Parsed(Exception): + pass + + +def _defaults(run_card): + """The parsed arguments of main() without command line options.""" + loader = types.ModuleType("custom_function_loader") + setup = types.ModuleType("gridpack_setup") + setup.__getattr__ = lambda name: None + setup.load_run_card = lambda: run_card + setup.load_madspace_data = dict + spec = importlib.util.spec_from_file_location("mg7_gridpack_card", _TEMPLATE) + module = importlib.util.module_from_spec(spec) + with mock.patch.dict(sys.modules, {"madspace": types.ModuleType("madspace"), + "custom_function_loader": loader, + "gridpack_setup": setup}): + spec.loader.exec_module(module) + + def parse_args(parser): + raise _Parsed(parser.parse_known_args([])[0]) + + with mock.patch.object(argparse.ArgumentParser, "parse_args", parse_args): + try: + module.main() + except _Parsed as parsed: + return parsed.args[0] + raise AssertionError("main() did not parse the arguments") + + +class TestGridpackCard(unittest.TestCase): + def test_card_without_amplicol_section(self): + args = _defaults(_run_card(amplicol=None)) + self.assertEqual(args.first_pass_max_overweight_truncation, 0.01) + self.assertEqual(args.first_pass_freeze_max_weight_after, 10000) + self.assertEqual(args.max_overweight_truncation, 0.001) + + def test_amplicol_section(self): + args = _defaults(_run_card(amplicol={"first_pass_max_overweight_truncation": 0.2})) + self.assertEqual(args.first_pass_max_overweight_truncation, 0.2) + self.assertEqual(args.first_pass_freeze_max_weight_after, 10000) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/iolibs/test_mg7_gridpack_compression.py b/tests/unit_tests/iolibs/test_mg7_gridpack_compression.py index 496b983afc..aa378f5937 100644 --- a/tests/unit_tests/iolibs/test_mg7_gridpack_compression.py +++ b/tests/unit_tests/iolibs/test_mg7_gridpack_compression.py @@ -22,10 +22,14 @@ def _gridpack_module(): madspace = types.ModuleType("madspace") loader = types.ModuleType("custom_function_loader") loader.load_custom_function = lambda *args, **kwargs: None + # bin/gridpack_setup.py, whose names gridpack.py imports + setup = types.ModuleType("gridpack_setup") + setup.__getattr__ = lambda name: None spec = importlib.util.spec_from_file_location("mg7_gridpack_compression", _TEMPLATE) module = importlib.util.module_from_spec(spec) with mock.patch.dict(sys.modules, {"madspace": madspace, - "custom_function_loader": loader}): + "custom_function_loader": loader, + "gridpack_setup": setup}): spec.loader.exec_module(module) return module diff --git a/tests/unit_tests/iolibs/test_mg7_madnis_one_dim.py b/tests/unit_tests/iolibs/test_mg7_madnis_one_dim.py index cf7af11beb..edb645c166 100644 --- a/tests/unit_tests/iolibs/test_mg7_madnis_one_dim.py +++ b/tests/unit_tests/iolibs/test_mg7_madnis_one_dim.py @@ -25,14 +25,14 @@ import unittest from unittest import mock -from madgraph.iolibs.template_files.mg7 import bootstrap as mg7_bootstrap +from madgraph.mg7 import bootstrap as mg7_bootstrap # Importing launch without a compiled madspace starts the madspace build, and # test_manager imports every unit test module while collecting: check first. launch = None if mg7_bootstrap.madspace_is_installed(): try: - import madgraph.iolibs.template_files.mg7.launch as launch + import madgraph.mg7.launch as launch except Exception as error: launch = None import_error = error @@ -116,7 +116,8 @@ def test_training_skips_the_one_dimensional_subprocess(self): contexts=[mock.Mock()], devices=[None], pool_sizes=[1], status_file=None, run_seed=1, build_event_generator=mock.Mock(return_value='generator')) - with mock.patch.object(launch, 'ms') as ms: + from madgraph.mg7 import madgraph_process + with mock.patch.object(madgraph_process, 'ms') as ms: ms.MultiMadnisTraining.return_value.active_channels.return_value = [[1]] launch.MadgraphProcess.train_madnis(process) self.assertEqual( diff --git a/tests/unit_tests/various/test_banner.py b/tests/unit_tests/various/test_banner.py index d03e2e7ef7..b242304205 100755 --- a/tests/unit_tests/various/test_banner.py +++ b/tests/unit_tests/various/test_banner.py @@ -2578,6 +2578,15 @@ def test_extra_jets_turn_mlm_on(self): self.assertNotIn('jet-lepton-delta_r', rc.dynamic_sections['cuts']) self.assertIn('lepton-delta_r', rc.dynamic_sections['cuts']) + def test_amplicol_output_is_no_merging(self): + rc = bannermod.RunCardMG7() + rc.create_default_for_process( + self.PC(), '', [[mg7_proc([2, -1], [-11, 12])], + [mg7_proc([2, -1], [-11, 12, 21])]], amplicol=True) + self.assertFalse(rc.is_mlm_merged()) + self.assertEqual(rc['phasespace']['xqcut'], 0.0) + self.assertIn('jet-delta_r', rc.dynamic_sections['cuts']) + def test_an_extra_photon_is_no_merging(self): rc = self.build([[mg7_proc([2, -1], [-11, 12])], [mg7_proc([2, -1], [-11, 12, 22])]]) diff --git a/tests/unit_tests/various/test_export_amplicol.py b/tests/unit_tests/various/test_export_amplicol.py new file mode 100644 index 0000000000..5e5cf8c533 --- /dev/null +++ b/tests/unit_tests/various/test_export_amplicol.py @@ -0,0 +1,189 @@ +################################################################################ +# +# Copyright (c) 2026 The MadGraph5_aMC@NLO Development team and Contributors +# +# This file is a part of the MadGraph5_aMC@NLO project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph5_aMC@NLO license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""Test the pyamplicol worker handling of export_amplicol, with fake workers +instead of pyamplicol.""" + +import io +import json +import os +import shutil +import subprocess +import sys +import tempfile +import unittest +from unittest import mock + +import madgraph +import madgraph.iolibs.export_amplicol as export_amplicol + + +class TestRunWorker(unittest.TestCase): + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmp) + + def fake_worker(self, code): + path = os.path.join(self.tmp, 'worker.py') + with open(path, 'w') as stream: + stream.write('import sys, time\n' + code) + patcher = mock.patch.object(export_amplicol, 'WORKER', path) + patcher.start() + self.addCleanup(patcher.stop) + + def test_success_prints_the_output(self): + self.fake_worker('print("generated")\n') + with mock.patch('sys.stdout', new=io.StringIO()) as out: + export_amplicol.run_worker('job.json') + self.assertIn('generated', out.getvalue()) + + def test_failure_raises(self): + self.fake_worker('sys.exit(1)\n') + with self.assertRaises(madgraph.InvalidCmd): + export_amplicol.run_worker('job.json') + + def test_busy_symbolica_is_retried_then_given_up(self): + self.fake_worker('print(%r); sys.exit(1)\n' % export_amplicol.SYMBOLICA_BUSY) + with mock.patch.object(export_amplicol, 'SYMBOLICA_RETRY', 0), \ + mock.patch.object(export_amplicol, 'SYMBOLICA_MAX_WAIT', 0.3), \ + mock.patch.object(export_amplicol, 'run_worker_once', + wraps=export_amplicol.run_worker_once) as once: + with self.assertRaisesRegex(madgraph.InvalidCmd, 'another process'): + export_amplicol.run_worker('job.json') + self.assertGreater(once.call_count, 1) + + def test_interrupt_kills_the_worker(self): + self.fake_worker('print("started", flush=True); time.sleep(60)\n') + workers = [] + popen = subprocess.Popen + + def record(*args, **kwargs): + workers.append(popen(*args, **kwargs)) + return workers[-1] + + with mock.patch('subprocess.Popen', side_effect=record), \ + mock.patch('sys.stdout.write', side_effect=KeyboardInterrupt): + with self.assertRaises(KeyboardInterrupt): + export_amplicol.run_worker('job.json') + self.assertIsNotNone(workers[0].returncode) + + +class TestCheckPyamplicol(unittest.TestCase): + + def test_stale_install_is_an_error(self): + from madgraph.mg7 import bootstrap + with mock.patch.object(bootstrap, 'ensure_pyamplicol'), \ + mock.patch.object(bootstrap, 'check_pyamplicol_source_hash', + side_effect=RuntimeError('mismatch')): + with self.assertRaisesRegex(madgraph.InvalidCmd, 'mismatch'): + export_amplicol.check_pyamplicol() + + def test_check_script_rejects_another_pyamplicol(self): + tmp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, tmp) + export_amplicol.write_rusticol_config( + tmp, {'version': '0.0.0', 'source_hash': 'stale'}) + result = subprocess.run([os.path.join(tmp, 'check-pyamplicol')], + capture_output=True, text=True) + self.assertNotEqual(result.returncode, 0) + self.assertIn('generated with pyamplicol 0.0.0', result.stderr) + + +class FakeModel(dict): + """The attributes of a restricted MG model that model_job reads.""" + + def __init__(self, path, name): + super().__init__(name=name, version_tag=path + '##0') + + def get(self, key): + import madgraph.core.base_objects as base_objects + if key in ('modelpath', 'restrict_name'): + return base_objects.Model.get(self, key) + return self[key] + + +class TestModelJob(unittest.TestCase): + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmp) + self.path = os.path.join(self.tmp, 'my-model_UFO') + os.mkdir(self.path) + + def restriction(self, name): + job = export_amplicol.model_job(FakeModel(self.path, name), + os.path.join(self.tmp, 'cache')) + return job['restriction'] + + def test_restrictions_of_a_model_name_with_dashes(self): + self.assertEqual(self.restriction('my-model_UFO'), 'default') + self.assertEqual(self.restriction('my-model_UFO-full'), 'none') + self.assertEqual(self.restriction('my-model_UFO-no_b-mass'), + os.path.join(self.path, 'restrict_no_b-mass.dat')) + + +class TestModelCache(unittest.TestCase): + """The compiled model is reused only for the same model, restriction and + pyamplicol.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmp) + self.path = os.path.join(self.tmp, 'sm') + os.mkdir(self.path) + self.cache = os.path.join(self.tmp, export_amplicol.MODEL_FILE) + self.compiled = [] + source = mock.Mock() + source.compile.side_effect = lambda prepared_output, evaluator: ( + self.compiled.append(prepared_output), + open(prepared_output, 'w').close()) + self.pyamplicol = mock.Mock() + self.pyamplicol.ModelSource.from_path.return_value = source + config = mock.Mock() + patcher = mock.patch.dict(sys.modules, {'pyamplicol': self.pyamplicol, + 'pyamplicol.config': config}) + patcher.start() + self.addCleanup(patcher.stop) + + def compile(self, version): + from madgraph.iolibs import amplicol_worker + job = export_amplicol.model_job(FakeModel(self.path, 'sm'), self.cache, + {'version': version, 'source_hash': 'h'}) + amplicol_worker.model_source(self.pyamplicol, json.loads(json.dumps(job))) + return job + + def test_cache_is_reused_for_the_same_pyamplicol(self): + first = self.compile('1') + self.assertEqual(self.compile('1')['key'], first['key']) + self.assertEqual(len(self.compiled), 1) + + def test_cache_is_rebuilt_for_another_pyamplicol(self): + self.compile('1') + self.compile('2') + self.assertEqual(len(self.compiled), 2) + + def test_cache_is_rebuilt_for_another_restriction_card(self): + card = os.path.join(self.path, 'restrict_default.dat') + with open(card, 'w') as stream: + stream.write('a') + self.compile('1') + with open(card, 'w') as stream: + stream.write('b') + self.compile('1') + self.assertEqual(len(self.compiled), 2) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/unit_tests/various/test_mg7_hwu_output.py b/tests/unit_tests/various/test_mg7_hwu_output.py index 8f8462a544..9e1c8c2162 100644 --- a/tests/unit_tests/various/test_mg7_hwu_output.py +++ b/tests/unit_tests/various/test_mg7_hwu_output.py @@ -29,7 +29,7 @@ import unittest import madgraph.various.histograms as histograms -from madgraph.iolibs.template_files.mg7 import hwu_output +from madgraph.mg7 import hwu_output def make_histogram(name='jet-pt', bin_count=2, weights=True): diff --git a/tests/unit_tests/various/test_mg7_plots.py b/tests/unit_tests/various/test_mg7_plots.py index 533dcf345d..13f157bc00 100644 --- a/tests/unit_tests/various/test_mg7_plots.py +++ b/tests/unit_tests/various/test_mg7_plots.py @@ -29,7 +29,7 @@ import unittest from madgraph import MG5DIR -from madgraph.iolibs.template_files.mg7 import plots +from madgraph.mg7 import plots def make_histogram(name='jet-pt', values=None, bin_count=4): @@ -156,7 +156,7 @@ def test_matplotlib_does_not_narrate_the_drawing(self): 'import json, logging, os, shutil, sys\n' 'logging.basicConfig(level=logging.DEBUG)\n' 'sys.path.insert(0, %r)\n' % MG5DIR + - 'from madgraph.iolibs.template_files.mg7 import plots\n' + 'from madgraph.mg7 import plots\n' 'out = %r\n' % os.path.join(self.path, 'sub') + 'plots.render([%r], out)\n' % make_histogram() ) diff --git a/tests/unit_tests/various/test_pyamplicol_install.py b/tests/unit_tests/various/test_pyamplicol_install.py index a23d0a0db5..e25980077a 100644 --- a/tests/unit_tests/various/test_pyamplicol_install.py +++ b/tests/unit_tests/various/test_pyamplicol_install.py @@ -194,5 +194,69 @@ def test_generated_files_are_ignored(self): self.assertEqual(source_hash.compute_source_hash(root), before) +class TestInstalledPyamplicol(unittest.TestCase): + """bootstrap finds the installed pyamplicol without importing it.""" + + def setUp(self): + from madgraph.mg7 import bootstrap + self.bootstrap = bootstrap + self.tmp = Path(tempfile.mkdtemp()) + self.addCleanup(shutil.rmtree, self.tmp) + patcher = mock.patch.object(bootstrap, 'PYAMPLICOL_INSTALL_DIR', self.tmp / 'none') + patcher.start() + self.addCleanup(patcher.stop) + + def system_package(self, version='7.0.0', source_hash='abc123'): + site = self.tmp / 'site' + (site / 'pyamplicol').mkdir(parents=True) + (site / 'pyamplicol' / '__init__.py').write_text('raise ImportError\n') + (site / 'pyamplicol' / '_source_hash.py').write_text( + f'SOURCE_HASH = "{source_hash}"\n') + info = site / f'pyamplicol-{version}.dist-info' + info.mkdir() + (info / 'METADATA').write_text(f'Name: pyamplicol\nVersion: {version}\n') + patcher = mock.patch.object(sys, 'path', [str(site)] + sys.path) + patcher.start() + self.addCleanup(patcher.stop) + + def test_system_install_counts_as_installed(self): + self.system_package() + self.assertTrue(self.bootstrap.pyamplicol_is_installed()) + self.assertEqual(self.bootstrap.installed_pyamplicol(), + {'version': '7.0.0', 'source_hash': 'abc123'}) + + def test_source_directory_is_not_an_install(self): + with mock.patch.object(sys, 'path', [str(MG_ROOT)]): + self.assertFalse(self.bootstrap.pyamplicol_is_installed()) + + def hash_check(self, stdout, returncode=0): + result = subprocess.CompletedProcess([], returncode, stdout, 'boom') + with mock.patch('subprocess.run', return_value=result): + return self.bootstrap.check_pyamplicol_source_hash() + + def test_matching_source_hash(self): + self.system_package() + self.assertEqual(self.hash_check('abc123\n')['version'], '7.0.0') + + def test_source_hash_mismatch_is_an_error(self): + self.system_package() + with self.assertRaisesRegex(RuntimeError, 'install pyamplicol'): + self.hash_check('def456\n') + + def test_source_hash_failure_is_an_error(self): + self.system_package() + with self.assertRaisesRegex(RuntimeError, 'boom'): + self.hash_check('', returncode=1) + + def test_build_check(self): + record = self.tmp / 'record.json' + record.write_text(json.dumps({'version': '7.0.0', 'source_hash': 'abc123'})) + self.system_package() + self.bootstrap.check_pyamplicol_build(record) + record.write_text(json.dumps({'version': '6.0.0', 'source_hash': 'old'})) + with self.assertRaisesRegex(SystemExit, 'pyamplicol 6.0.0'): + self.bootstrap.check_pyamplicol_build(record) + + if __name__ == '__main__': unittest.main()