diff --git a/.github/workflows/docs_examples.yml b/.github/workflows/docs_examples.yml index 3259fe375b..5138407ccf 100644 --- a/.github/workflows/docs_examples.yml +++ b/.github/workflows/docs_examples.yml @@ -3,13 +3,16 @@ name: MadSpace docs examples # The other MG7 acceptance workflows all carry `paths-ignore: ['docs/**']`, so a # docs-only commit currently triggers none of them. This workflow does the # opposite: it triggers ONLY on the paths that can affect the documentation -# examples, so editing docs/source/madspace/examples/*.rst is actually tested. +# examples, so editing docs/source/madspace/examples/*.rst (or the event +# streaming example in docs/source/event_stream.rst) is actually tested. on: push: branches: - '**' paths: - 'docs/source/madspace/examples/**' + - 'docs/source/event_stream.rst' + - 'madgraph/iolibs/template_files/mg7/**' - 'madspace/**' - 'tests/acceptance_tests/test_madspace_examples.py' - '.github/workflows/docs_examples.yml' @@ -18,6 +21,8 @@ on: types: [opened, synchronize, reopened] paths: - 'docs/source/madspace/examples/**' + - 'docs/source/event_stream.rst' + - 'madgraph/iolibs/template_files/mg7/**' - 'madspace/**' - 'tests/acceptance_tests/test_madspace_examples.py' - '.github/workflows/docs_examples.yml' @@ -55,4 +60,4 @@ jobs: - name: run the madspace doc examples run: | cd $GITHUB_WORKSPACE - ./tests/test_manager.py test_madspace_example_simple_mappings test_madspace_example_cuts test_madspace_example_diagram_mapping test_madspace_example_integration_order test_madspace_example_matrix_element test_madspace_example_pdf test_madspace_example_integrator test_madspace_example_integrator_madnis test_madspace_example_flow_training -pA -t0 -l INFO + ./tests/test_manager.py test_madspace_example_simple_mappings test_madspace_example_cuts test_madspace_example_diagram_mapping test_madspace_example_integration_order test_madspace_example_matrix_element test_madspace_example_pdf test_madspace_example_integrator test_madspace_example_integrator_madnis test_madspace_example_flow_training test_gridpack_example_event_stream -pA -t0 -l INFO diff --git a/docs/source/conf.py b/docs/source/conf.py index 541cbea840..38a93e0097 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -14,6 +14,17 @@ if os.path.isdir(os.path.join(_local_install, "madspace")): sys.path.insert(0, _local_install) +# bin/event_stream.py of the gridpacks, for the API reference in event_stream.rst +sys.path.insert( + 0, + os.path.abspath( + os.path.join( + os.path.dirname(__file__), + "..", "..", "madgraph", "iolibs", "template_files", "mg7", + ) + ), +) + # Configuration file for the Sphinx documentation builder. # # For the full list of built-in configuration values, see the documentation: diff --git a/docs/source/custom_functions.rst b/docs/source/custom_functions.rst new file mode 100644 index 0000000000..a9d17bb0da --- /dev/null +++ b/docs/source/custom_functions.rst @@ -0,0 +1,324 @@ +Custom functions +================ + +MG7 can call your own code for four parts of the calculation: + +``cut`` + Decides if a phase-space point passes the cuts. It is applied on top of the + ``[cuts]`` of the :doc:`run card `. + +``scale`` + The renormalization and factorization scale. It replaces ``dynamical_scale_choice``, + like ``user_dynamical_scale`` in MadEvent. + +``pdf`` + The parton densities. It replaces the LHAPDF set of ``[beam] pdf``. + +``bias`` + An event bias. More events are generated where the bias is large, and their weights + are lowered to compensate. + +You write the function either for a single phase-space point in C, C++, Fortran, CUDA or +HIP, or for a whole batch of points in Python with numpy or PyTorch. MG7 builds a compiled +function into a shared library and calls it on batches of events. The library implements +the :doc:`UMAMI interface `, the same interface the matrix elements +use. A generated wrapper takes care of that, so you only write the function itself. A +Python function needs no build. MG7 imports it and calls it through the same interface. + +Creating a custom function +-------------------------- + +Run ``bin/init_custom_func`` in the process directory. Without options, it asks which +function and language you want: + +.. code-block:: text + + $ ./bin/init_custom_func + Which custom function do you want to create? + 1. cut a cut, 1 keeps the event and 0 rejects it + 2. scale the renormalization and factorization scale + 3. pdf a parton density x f(x, Q) + 4. bias an event bias for the unweighting + Which language do you want to write it in? + 1. c C + 2. cpp C++ + 3. fortran Fortran + 4. cuda CUDA, for runs on NVIDIA GPUs + 5. hip HIP, for runs on AMD GPUs + 6. python-numpy Python with numpy, for CPU runs, no compilation + 7. python-torch Python with PyTorch, for CPU or GPU runs, no compilation + +The same choices are available as options: + +.. code-block:: bash + + ./bin/init_custom_func --function cut --language cpp + +``-f`` and ``-l`` are short for ``--function`` and ``--language``. ``-y`` takes the +defaults (``cut`` and ``cpp``) for the options you leave out instead of asking. +``--force`` overwrites an existing directory. ``--no-build`` skips the build. +``--no-set-run-card`` leaves the run card unchanged. + +For a compiled language, the script creates a directory ``custom_``: + +``custom_.`` + The function body. This is the only file you edit. + +``wrapper_cpu.c`` or ``wrapper_gpu.cu`` + The UMAMI wrapper. It unpacks the batched inputs and calls your function once per + event. + +``umami.h`` + The UMAMI header. + +``Makefile`` + Builds ``lib/custom_.so``. + +For ``python-numpy`` and ``python-torch``, the directory only holds the function, +``custom_.py``. + +By default, the script builds the default function right away and enables it in +``Cards/run_card.toml``: + +.. code-block:: toml + + [custom_functions] + cut = "lib/custom_cut.so" + +A Python function is enabled with the path of its file: + +.. code-block:: toml + + [custom_functions] + cut = "custom_cut/custom_cut.py" + +``bin/generate_events`` runs ``make`` in ``custom_`` before each run. Your +changes are therefore picked up without a manual rebuild. The build output goes to +``custom_/compile.log``. + +Function signatures +------------------- + +The cut, scale and bias functions receive the external particles of one event. The +incoming particles come first. The momenta are in the lab frame, in GeV, with components +``(E, px, py, pz)``. The PDG ids are those of the first flavor of the subprocess for cuts +and scales, like the built-in cuts, and those of the sampled flavor for the bias. + +C, C++ and CUDA/HIP: + +.. code-block:: c + + double custom_cut(int particle_count, const int* pdg_ids, const double momenta[][4]); + double custom_scale(int particle_count, const int* pdg_ids, const double momenta[][4]); + double custom_bias(int particle_count, const int* pdg_ids, const double momenta[][4]); + double custom_pdf(int beam_index, int pdg_id, double x, double q); + +In C++, the functions are declared ``extern "C"``. In CUDA and HIP, they are +``__device__`` functions and run on the GPU, one thread per event. + +Fortran: + +.. code-block:: fortran + + function custom_cut(particle_count, pdg_ids, momenta) bind(C, name="custom_cut") result(value) + use iso_c_binding + integer(c_int), value :: particle_count + integer(c_int), intent(in) :: pdg_ids(particle_count) + real(c_double), intent(in) :: momenta(0:3, particle_count) + real(c_double) :: value + +The PDF takes ``integer(c_int), value :: beam_index, pdg_id`` and +``real(c_double), value :: x, q``. + +What the functions return: + +``cut`` + 1 to keep the event, 0 to reject it. The value multiplies the event weight, so values + in between act as a weight factor. + +``scale`` + The scale in GeV, used for the renormalization scale and both factorization scales. + ``scale_factor`` is applied on top. ``fixed_ren_scale`` and ``fixed_fact_scale`` still + take precedence. + +``pdf`` + :math:`x f(x, Q)` for the parton ``pdg_id`` (21 for the gluon) in beam ``beam_index`` + (0 or 1), at momentum fraction ``x`` and factorization scale ``q``. + +``bias`` + The bias factor :math:`b > 0`. The phase space is sampled and unweighted according to + :math:`|w|\,b`. The written events carry the weight divided by :math:`b`. Distributions + filled with the event weights therefore stay unbiased, while regions with a large bias + get more events. The reported cross section does not include the bias. + +Python functions +---------------- + +A Python function is called once per batch of phase-space points, with arrays that hold +one entry per point. The file defines the function ``custom_`` and the +framework it uses: + +.. code-block:: python + + import numpy as np + + FRAMEWORK = "numpy" # or "torch" + + def custom_cut(momenta, pdg_ids): + ... + +The cut, scale and bias functions take ``(momenta, pdg_ids)``: + +``momenta`` + float64, shape ``(batch, particle_count, 4)``. The momenta ``(E, px, py, pz)`` of the + external particles in GeV, lab frame, incoming particles first. + +``pdg_ids`` + int32, shape ``(batch, particle_count)``. Their PDG ids, as for the compiled + functions. + +The PDF function takes ``(beam_index, pdg_id, x, q)``, each with shape ``(batch,)``. +``beam_index`` and ``pdg_id`` are int32, ``x`` and ``q`` are float64. + +Every function returns a float64 array of shape ``(batch,)``, with the same meaning as for +the compiled functions. MG7 checks the type, the dtype, the shape and the device of the +result. With numpy the arguments are numpy arrays, and with ``FRAMEWORK = "torch"`` they are +torch tensors. They are views of MG7's own buffers, so they must not be modified or kept +after the call. + +numpy functions run on the CPU only. torch functions run on the device of the run, so the +same file works for CPU and GPU runs. On the GPU the function runs on the CUDA or HIP stream +of the MG7 runtime, and the arguments and the result stay on the GPU. + +The same cut as in the C++ example below, with numpy: + +.. code-block:: python + + import numpy as np + + FRAMEWORK = "numpy" + + def custom_cut(momenta, pdg_ids): + pt = np.hypot(momenta[:, 2:4, 1], momenta[:, 2:4, 2]) + return np.where((pt >= 150.).all(axis=1), 1., 0.) + +:math:`H_T/2` with torch, for CPU or GPU runs: + +.. code-block:: python + + import torch + + FRAMEWORK = "torch" + + def custom_scale(momenta, pdg_ids): + final = momenta[:, 2:] + mt = torch.sqrt(torch.clamp(final[..., 0] ** 2 - final[..., 3] ** 2, min=0.)) + return mt.sum(dim=1) / 2. + +On the CPU, MG7 splits each batch over the threads of ``cpu_thread_pool_size`` and calls +the function from these threads. Python runs one call at a time, but numpy and torch +release this lock inside their array operations, so vectorized functions still run +partly in parallel. MG7 limits torch to one thread on the CPU, so that it does not start +threads of its own on top. numpy uses several threads only for linear algebra, such as +matrix products. If your function relies on it, set ``OMP_NUM_THREADS=1``. + +The Python API behind these functions is ``madspace.PythonFunction``. It can also be used +directly with madspace. + +Examples +-------- + +A cut requiring :math:`p_T > 150` GeV for both particles of a ``p p > t t~`` process, in +C++: + +.. code-block:: c++ + + #include + + extern "C" double custom_cut(int particle_count, const int* pdg_ids, const double momenta[][4]) { + for (int i = 2; i < 4; ++i) { + if (std::hypot(momenta[i][1], momenta[i][2]) < 150.) { + return 0.; + } + } + return 1.; + } + +:math:`H_T/2` as the scale, in Fortran: + +.. code-block:: fortran + + function custom_scale(particle_count, pdg_ids, momenta) bind(C, name="custom_scale") result(value) + use iso_c_binding + implicit none + integer(c_int), value :: particle_count + integer(c_int), intent(in) :: pdg_ids(particle_count) + real(c_double), intent(in) :: momenta(0:3, particle_count) + real(c_double) :: value + integer :: i + value = 0d0 + do i = 3, particle_count + value = value + sqrt(max(momenta(0, i)**2 - momenta(3, i)**2, 0d0)) + end do + value = value / 2d0 + end function custom_scale + +A bias that populates the :math:`p_T` tail of the first final-state particle, in C: + +.. code-block:: c + + #include + + double custom_bias(int particle_count, const int* pdg_ids, const double momenta[][4]) { + double pt2 = momenta[2][1] * momenta[2][1] + momenta[2][2] * momenta[2][2]; + return fmax(pt2 / (100. * 100.), 1.); + } + +An LHAPDF set called through the LHAPDF C++ API: + +.. code-block:: c++ + + #include + #include "LHAPDF/LHAPDF.h" + + extern "C" double custom_pdf(int beam_index, int pdg_id, double x, double q) { + static LHAPDF::PDF* pdf = LHAPDF::mkPDF("NNPDF40_lo_as_01180", 0); + static std::mutex mutex; + std::lock_guard lock(mutex); + return pdf->xfxQ(pdg_id, x, q); + } + +The function is called from several threads, so the call into LHAPDF is protected by a +lock. To build it, pass the LHAPDF flags to ``make``: + +.. code-block:: bash + + make -C custom_pdf CFLAGS="$(lhapdf-config --cflags)" \ + LIBS="$(lhapdf-config --libs) -Wl,-rpath,$(lhapdf-config --libdir)" + +Building +-------- + +The ``Makefile`` takes the usual ``CC``, ``CXX`` and ``FC`` compilers, with ``gfortran`` +as the default for Fortran. ``OPTFLAGS`` defaults to ``-O3``. Extra compiler flags go into +``CFLAGS`` and extra libraries into ``LIBS``. The CUDA and HIP builds detect the GPU +architecture like the matrix elements, and accept the same ``MADGRAPH_CUDA_ARCHITECTURE`` +and ``MADGRAPH_HIP_ARCHITECTURE`` overrides. + +The language decides the device the function runs on. C, C++ and Fortran functions run on +the CPU, and CUDA and HIP functions run on the GPU. Use the language that matches the +``device`` of the run. Python functions are not built. numpy functions run on the CPU, and +torch functions run on the device of the run. + +Limitations +----------- + +- With a custom PDF, :math:`\alpha_s` still comes from the PDF set of ``[beam] pdf``, and + the scale and PDF variations of ``[systematics]`` are disabled. +- With a custom bias, the weighted integration histograms of ``[run] + weighted_histograms`` include the bias. The event weights do not. +- Decays ignore a custom scale and a custom PDF. +- A gridpack stores the custom libraries and Python files it was made with. They cannot + be changed in the gridpack. A Python function in a gridpack needs the same Python + packages, such as numpy or torch, where the gridpack runs. +- Python functions with torch support only the first GPU of a machine. diff --git a/docs/source/event_stream.rst b/docs/source/event_stream.rst new file mode 100644 index 0000000000..306c7b14ca --- /dev/null +++ b/docs/source/event_stream.rst @@ -0,0 +1,151 @@ +Streaming events from Python +============================ + +A gridpack can generate events directly into a Python program, without writing any files. +This is useful when the events are consumed right away, for example by an analysis written +in Python or as training data for a neural network. The ``EventStream`` class in the +gridpack's ``bin/event_stream.py`` returns the events either as LHE event objects or as +batches of NumPy arrays or PyTorch tensors. + +Event streaming needs a gridpack created with ``run_mode = "fix_max_weight"`` (see +:ref:`gridpack-run-mode`). In this mode, the maximum weight of every channel is fixed when +the gridpack is created. An event is final as soon as it is unweighted, so it can be returned +right away without a survey or a combination step. + +Example: the e+ e- invariant mass +--------------------------------- + +This example fills a histogram of the invariant mass of the electron pair in +``p p > e+ e-``. First, create the gridpack at the MadGraph7 prompt: + +.. code-block:: text + :name: event-stream-mg5-commands + + generate p p > e+ e- + output mg7 PROC_ppee + launch PROC_ppee + set save_gridpack true + set run_mode fix_max_weight + done + +The gridpack is written to ``PROC_ppee/Events/run_01/gridpack``. Its ``bin`` directory +contains ``event_stream.py``, so add it to the Python path and import ``EventStream``. The +script can run from any directory. + +.. code-block:: python + + import sys + + import numpy as np + + sys.path.insert(0, "PROC_ppee/Events/run_01/gridpack/bin") + from event_stream import EventStream + +Creating an ``EventStream`` starts the event generation in background threads. Its +arguments, like the device, the seed and the batch sizes, default to the values in the +gridpack's ``Cards/grid_run_card.toml``. The stream is endless, so ``max_events`` limits +the number of events. Each batch is an :py:class:`EventBatch ` +with one array per LHE field. The particle fields have the shape +``(batch_size, max_particle_count)``. + +.. code-block:: python + + bins = np.linspace(60.0, 120.0, 31) + hist = np.zeros(len(bins) - 1) + weight_sum = 0.0 + + with EventStream(seed=42) as stream: + for batch in stream.batches(10000, max_events=100000): + # four-momentum of the e+ e- pair: sum over the outgoing electrons + is_electron = (np.abs(batch.pdg_id) == 11) & (batch.status_code == 1) + e, px, py, pz = ( + np.sum(np.where(is_electron, p, 0.0), axis=1) + for p in (batch.energy, batch.px, batch.py, batch.pz) + ) + m_ee = np.sqrt(np.maximum(e**2 - px**2 - py**2 - pz**2, 0.0)) + hist += np.histogram(m_ee, bins, weights=batch.weight)[0] + weight_sum += np.sum(batch.weight) + cross_section = stream.integral + print(f"cross section: {cross_section:.1f} +- {stream.integral_error:.1f} pb") + +The event weights are ``1``, except for overweight events, so they have to be normalized. +Scaling them to the cross section gives the differential cross section in pb/GeV: + +.. code-block:: python + + dsigma = hist * cross_section / weight_sum / np.diff(bins) + for low, value in zip(bins[:-1], dsigma): + print(f"{low:5.0f} GeV {value:8.3f} pb/GeV " + "#" * int(50 * value / dsigma.max())) + + peak = np.argmax(dsigma) + assert bins[peak] < 91.2 < bins[peak + 1], "the Z peak is missing" + +The output shows the Z peak (shortened):: + + cross section: 880.2 +- 0.6 pb + ... + 84 GeV 7.569 pb/GeV ## + 86 GeV 16.543 pb/GeV ##### + 88 GeV 51.287 pb/GeV ################ + 90 GeV 158.376 pb/GeV ################################################## + 92 GeV 66.743 pb/GeV ##################### + 94 GeV 18.457 pb/GeV ##### + ... + +Output formats +-------------- + +There are three ways to get events from a stream: + +- Iterating over the stream yields one ``madspace.LHEEvent`` at a time. The events are + completed in chunks of 1000 internally. +- :py:meth:`next_events ` returns a list of a given + number of ``LHEEvent`` objects. +- :py:meth:`batches ` and + :py:meth:`next_batch ` return + :py:class:`EventBatch ` objects. With ``output_type="torch"``, + their fields are PyTorch tensors on the CPU instead of NumPy arrays. + +Events have different numbers of particles if the process has several subprocesses or +includes decays. The particle fields are then padded with zeros. ``particle_count`` holds +the number of particles of every event and ``particle_mask`` marks the valid entries. + +Weights and cross section +------------------------- + +Every event weight is ``1``, or ``-1`` for an event with a negative weight. The exception are +overweight events, whose weight exceeded the fixed maximum weight of their channel. They +keep the ratio of their weight to the maximum weight, so their weight is larger than ``1``. +Treating them as ``1`` would bias the distributions in the phase-space regions where they +occur. + +To normalize a sample to the cross section, scale its weights by +``stream.integral / sum(weights)``. The cross section estimate improves while events are +generated: it combines the integrals stored in the gridpack with the weighted events +generated by the stream so far. Read it after generating the sample, as in the example +above. + +How it works +------------ + +Each channel of the gridpack keeps a buffer of unweighted events, which background threads +refill whenever it runs low. The channel of every event is drawn at random, with +probabilities given by the current cross section estimates of the channels. Generating +the phase-space points and evaluating the matrix elements happen in separate steps, as in a +normal gridpack run. + +The events are reproducible: for a given seed, the stream always returns the same sequence of +events. It does not depend on the number of threads or on how the events are split into +batches. + +API reference +------------- + +.. autoclass:: event_stream.EventStream + :members: + :special-members: __iter__ + :member-order: bysource + +.. autoclass:: event_stream.EventBatch + :members: + :member-order: bysource diff --git a/docs/source/gridpacks.rst b/docs/source/gridpacks.rst index 84b11fb660..39c1526857 100644 --- a/docs/source/gridpacks.rst +++ b/docs/source/gridpacks.rst @@ -16,8 +16,13 @@ start a run: set save_gridpack true launch -The gridpack is written to ``Events/_NN/gridpack`` after the run. It is a plain -directory. Pack it with ``tar`` to move it to another machine. +The gridpack is written to ``Events/_NN/gridpack``, or to ``gridpack`` in the +``output_dir`` of the ``[gridpack]`` section if it is set. It is a plain directory. With +``compress = true``, it is packed into ``gridpack.tar.gz`` instead, ready to be moved to +another machine. + +By default, the run stops once the gridpack is ready and does not generate events. See +:ref:`gridpack-run-mode` for the other options. A gridpack contains: @@ -33,7 +38,7 @@ A gridpack contains: The trained phase-space channels and the information needed to write LHE files and the scale and PDF weights. -``lib/`` or ``src/`` and ``SubProcesses/`` +``lib/`` or ``src/``, ``SubProcesses/`` and ``backend/`` The matrix elements, see :ref:`gridpack-matrix-elements`. ``madspace/`` @@ -42,6 +47,34 @@ A gridpack contains: ``Events/`` Empty. The output of each gridpack run goes here. +.. _gridpack-run-mode: + +Ending the run +^^^^^^^^^^^^^^ + +``run_mode`` in the ``[gridpack]`` section sets how the run that creates the gridpack +ends: + +``minimal`` (default) + No events are generated. The gridpack is stored right after the MadNIS training or, + with VEGAS, once the grid of every channel has converged. This is the fastest way + to create a gridpack. Each gridpack run then estimates the maximum weights and the + channel integrals on its own. + +``fix_max_weight`` + No events are generated. Each channel runs until ``freeze_max_weight_after`` + unweighted events are reached. Its maximum weight is stored in the gridpack and kept + fixed in every gridpack run, so all runs unweight against the same maximum. The + channel integrals are stored as well and combined with the ones of the gridpack run + to share the events between the channels from the start. Creating the gridpack takes + longer, especially for channels with a low unweighting efficiency, but the gridpack + runs do not have to estimate the maximum weights first. This helps if each run only + generates a few events. + 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. + Generating events ----------------- @@ -51,7 +84,10 @@ Run the script from anywhere: ./bin/generate_events --events 10000 --seed 42 -Each call creates a new directory ``Events/_NN``. A gridpack generates LHE files +Each call creates a new directory ``Events/_NN``. Use ``--output_dir`` to write +the output to a given directory instead, and ``--temp_output_dir`` to put the large +temporary npy files of the channels somewhere else, for example on a fast local disk of a +cluster node. Relative paths on the command line refer to the current directory. A gridpack generates LHE files by default, since it usually feeds a shower or detector simulation. The file is compressed to ``events.lhe.gz``. Use ``--output_format`` to get one of the npy formats instead. @@ -60,6 +96,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 ----------------- @@ -108,6 +153,9 @@ gridpack is created. * - - ``cut_efficiency_threshold``, ``max_cut_repetitions`` - Card only. + * - ``[gridpack]`` + - ``output_dir``, ``temp_output_dir`` + - ``--output_dir``, ``--temp_output_dir`` * - ``[systematics]`` - ``enable``, ``mur``, ``muf``, ``together``, ``dynamical_scale``, ``pdf``, ``write_inputs`` @@ -115,6 +163,8 @@ gridpack is created. An option on the command line overrides the card. The settings have the same meaning as in the main run card, see :doc:`run_card`. ``./bin/generate_events --help`` lists all options. +Relative ``output_dir`` and ``temp_output_dir`` paths in the card refer to the gridpack +directory. Both are reset to ``""`` when the gridpack is created. The card names the ``cpu_mode`` that the gridpack was built with, which can differ from the ``"auto"`` of the original run card. The matrix element libraries carry this name. If you @@ -143,7 +193,8 @@ Matrix elements that you used. A gridpack can only run on devices for which a library exists. ``include_source = true`` - The source code is copied to ``src/`` and ``SubProcesses/`` and ``lib/`` stays empty. + The source code is copied to ``src/``, ``SubProcesses/`` and ``backend/`` and ``lib/`` + stays empty. This is useful if the gridpack has to run on a different platform. The gridpack does not compile anything by itself, so build the libraries once before the first run: @@ -179,8 +230,9 @@ If both options are on, the compiled copy is used and the source is only a fallb rebuilding it. The compiled copy has to match the platform. The source code makes the gridpack portable. -The script warns if the ``madspace`` it uses differs from the one that created the gridpack. -In this case, the results can be wrong. +The script stops with an error if the ``madspace`` it uses differs from the one that +created the gridpack, since the results can be wrong in this case. Use +``--ignore_source_hash`` to run anyway, which only prints a warning. The options can be combined. The smallest gridpack sets ``include_madspace = false`` and relies on an installed ``madspace``. A fully portable gridpack sets ``include_source = true`` diff --git a/docs/source/index.rst b/docs/source/index.rst index 3c597d70cb..960ddec98c 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -27,7 +27,9 @@ If you need a stable release, use the installation first_steps run_card + custom_functions gridpacks + event_stream madevent_to_mg7 .. toctree:: diff --git a/docs/source/madspace/cpp-api.rst b/docs/source/madspace/cpp-api.rst index 64384364ef..783ce2b3d7 100644 --- a/docs/source/madspace/cpp-api.rst +++ b/docs/source/madspace/cpp-api.rst @@ -43,6 +43,7 @@ Builders that assemble compute-graph functions for channel weights, energy scale cpp/BatchSampler cpp/BufferUnweighter cpp/ChannelWeightNetwork + cpp/CustomFunction cpp/Cuts cpp/DifferentialCrossSection cpp/DiscreteHistogram @@ -69,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 --------------------------- @@ -79,10 +81,12 @@ Supporting types shared by the mappings: topologies, cuts, observables and densi :maxdepth: 1 cpp/AlphaSGrid + cpp/CustomPdf cpp/Diagram cpp/PdfGrid cpp/Propagator cpp/Topology + cpp/WeightingStep cpp/initialize_vegas_grid Compute graph @@ -125,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 @@ -155,17 +160,26 @@ 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 cpp/SystematicsConfig + cpp/TarReader + cpp/TarWriter cpp/Tensor cpp/TensorView cpp/ThreadPool cpp/ThreadResource + cpp/UmamiFunctions cpp/UnalignedRef cpp/Variation cpp/VegasGridOptimizer + cpp/WeightingGroup + cpp/WeightingPipeline cpp/available_backends cpp/compute_generation_batch_event_count cpp/cpu_device @@ -180,6 +194,8 @@ The high-level driver plus the tensor and device layer it runs on: integration, cpp/select_combine_channel_index cpp/set_lib_path cpp/set_simd_vector_size + cpp/weighting_groups_from_json + cpp/weighting_groups_to_json Utilities --------- diff --git a/docs/source/madspace/python-api.rst b/docs/source/madspace/python-api.rst index ec5d0ed19c..95caadfb9c 100644 --- a/docs/source/madspace/python-api.rst +++ b/docs/source/madspace/python-api.rst @@ -40,9 +40,11 @@ Builders that assemble compute-graph functions for channel weights, energy scale .. toctree:: :maxdepth: 1 + python/BufferUnweighter python/CachedPdf python/CachedScale python/ChannelWeightNetwork + python/CustomFunction python/CutItem python/Cuts python/DifferentialCrossSection @@ -77,12 +79,14 @@ Supporting types shared by the mappings: topologies, cuts, observables and densi :maxdepth: 1 python/AlphaSGrid + python/CustomPdf python/Decay python/Diagram python/LineRef python/PdfGrid python/Propagator python/Topology + python/WeightingStep python/initialize_vegas_grid Compute graph @@ -122,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 @@ -152,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 @@ -166,6 +172,8 @@ The high-level driver plus the tensor and device layer it runs on: integration, python/select_combine_channel_index python/set_lib_path python/set_simd_vector_size + python/weighting_groups_from_json + python/weighting_groups_to_json Utilities --------- @@ -177,3 +185,7 @@ Lower-level helpers and containers. python/NamedTypes python/NamedValues + python/PythonFunction + python/get_stream + python/set_stream + python/stream diff --git a/docs/source/run_card.rst b/docs/source/run_card.rst index 0192a54b98..0a2c95d62a 100644 --- a/docs/source/run_card.rst +++ b/docs/source/run_card.rst @@ -114,6 +114,31 @@ See :doc:`gridpacks` for details. ``include_madspace_source`` (``false``) Include the ``madspace`` source code. +``compress`` (``false``) + Pack the gridpack into ``gridpack.tar.gz`` instead of leaving it as a directory. + +``run_mode`` (``"minimal"``) + How the run that creates the gridpack ends. Only used if ``save_gridpack`` is set. + + - ``regular``: generate the events of the run, then store the gridpack. + - ``minimal``: no events. The gridpack is stored right after the MadNIS training or, + with VEGAS, once the grid of every channel has converged. + - ``fix_max_weight``: no events. Each channel runs until ``freeze_max_weight_after`` + unweighted events are reached. Its maximum weight is then fixed in the gridpack, + and its integral is used to share the events between the channels. This can be + slow for channels with a low unweighting efficiency. The gridpack also gets + ``bin/event_stream.py`` to stream events in memory from Python. + + See :ref:`gridpack-run-mode`. + +``temp_output_dir`` (``""``) + Directory for the temporary npy files of the channels. ``""`` uses the output + directory. + +``output_dir`` (``""``) + Directory for the output files. ``""`` creates a new run folder + ``Events/_NN``. If set, no run folder is created. + [beam] ------ @@ -308,6 +333,71 @@ VEGAS grids adapt the phase-space sampling during the survey. Learn the probabilities of the permutations of identical particles instead of using equal ones. +[custom_functions] +------------------ + +User-written functions, created with ``bin/init_custom_func``. Each one is a shared +library (``.so``) or a Python file (``.py``). Paths are relative to the process directory. +An empty string disables the function. See :doc:`custom_functions` for details. + +``cut`` (``""``) + Custom cut, for example ``"lib/custom_cut.so"`` or ``"custom_cut/custom_cut.py"``. It + is applied on top of the ``[cuts]``. + +``scale`` (``""``) + Custom scale. It replaces ``dynamical_scale_choice``. + +``pdf`` (``""``) + Custom parton density. It replaces the PDF set, but :math:`\alpha_s` still comes from + ``[beam] pdf``. The systematics are disabled. + +``bias`` (``""``) + Custom event bias. Events are generated with the weight times the bias and written with + the weight divided by the bias. + +.. _run-card-surrogate: + +[surrogate] +----------- + +A fast surrogate of the matrix element, for example a neural network. The integrand and +its optimization use the surrogate. The events are then reweighted or unweighted a second +time with the full matrix element, which is only evaluated for the events that pass the +first unweighting. The cross section and the events are those of the full matrix element. +The surrogate does not return channel weights, so it needs ``sde_strategy = +"denominators"``. + +``function`` (``""``) + The surrogate. An empty string disables it. + + - A ``.so`` library implementing the UMAMI interface with the inputs of the madmatrix + libraries (momenta, :math:`\alpha_s`, flavor index and the random numbers). It is + loaded once per subprocess: ``{subprocess}`` and ``{device}`` in the path are replaced + by the subprocess name (e.g. ``P1_gg_ttx``) and the device, like in + ``lib/libmadmatrix_{subprocess}_{device}.so``. + - A Python file defining ``custom_surrogate(momenta, pdg_ids, alpha_s)``, which returns + :math:`|M|^2` for a batch of events. ``momenta`` has the shape ``(batch, n, 4)``, + ``pdg_ids`` the shape ``(batch, n)`` and identifies subprocess and flavor, and + ``alpha_s`` the shape ``(batch,)``. ``FRAMEWORK = "torch"`` passes torch tensors, + which also run on GPUs, otherwise numpy arrays are passed. See + :doc:`custom_functions`. + +``mode`` (``"unweight"``) + ``"unweight"`` unweights the events a second time with respect to the ratio of the full + matrix element to the surrogate. The events then have unit weights apart from + overweights. ``"reweight"`` keeps the ratio in the event weights. + +``first_pass_max_overweight_truncation`` (``0.01``), ``first_pass_freeze_max_weight_after`` (``10000``) + ``max_overweight_truncation`` and ``freeze_max_weight_after`` of the unweighting with + the surrogate, in unweight mode. The ones of ``[generation]`` are used for the second + pass, which gives the final events. In reweight mode, ``[generation]`` sets the + unweighting with the surrogate. + +``param_card`` (``""``) + The param card a ``.so`` surrogate is initialized with. An empty string uses + ``Cards/param_card.dat``. A madmatrix library with a different param card, for example, + is a surrogate for testing. + .. _run-card-multiparticles: [multiparticles] diff --git a/madgraph/interface/tutorials/mg7.py b/madgraph/interface/tutorials/mg7.py index 3ef6c96b5f..c66196657b 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/export_cpp.py b/madgraph/iolibs/export_cpp.py index 14c0756790..e70cbfe98c 100755 --- a/madgraph/iolibs/export_cpp.py +++ b/madgraph/iolibs/export_cpp.py @@ -3325,6 +3325,18 @@ def copy_template(self, model): ) os.chmod(npy_bin, 0o755) + custom_bin = os.path.join("bin", "init_custom_func") + with open(custom_bin, "w") as f: + f.write( + "#! /usr/bin/env python3\n" + "import sys, os\n" + f"sys.path.append('{MG5DIR}')\n" + "from madgraph.iolibs.template_files.mg7.init_custom_func import main\n" + "if __name__ == '__main__':\n" + " sys.exit(main(me_dir=os.path.dirname(os.path.dirname(os.path.realpath(__file__)))))\n" + ) + os.chmod(custom_bin, 0o755) + # Recorded in Cards/me5_configuration.txt: the tools a run needs but cannot # rediscover on its own. LHAPDF above all -- bin/generate_events may be # driven from a shell that never sourced anything MadGraph-related. diff --git a/madgraph/iolibs/template_files/mg7/custom_function_loader.py b/madgraph/iolibs/template_files/mg7/custom_function_loader.py new file mode 100644 index 0000000000..b06f34150e --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_function_loader.py @@ -0,0 +1,134 @@ +"""Loads the functions of the [custom_functions] run card section. + +Shared by bin/generate_events of a process directory and of a gridpack. A path +ending in .so is a UMAMI library built by bin/init_custom_func, a path ending in +.py a Python module defining custom_ and FRAMEWORK. +""" + +import importlib.util +import logging +import os + +logger = logging.getLogger("madevent") + +# the arguments passed to custom_ in a Python module +PYTHON_INPUTS = { + "cut": ["momenta", "pdg_ids"], + "scale": ["momenta", "pdg_ids"], + "bias": ["momenta", "pdg_ids"], + "pdf": ["beam_index", "pdg_id", "momentum_fraction", "scale"], +} + +_modules = {} + + +def _import_module(path): + """Import the module at ``path`` once per process.""" + path = os.path.abspath(path) + if path not in _modules: + spec = importlib.util.spec_from_file_location( + f"mg7_custom_function_{len(_modules)}", path + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + _modules[path] = module + return _modules[path] + + +def _python_function(path, name, device_types): + """The function ``name`` of the Python module at ``path`` and its + FRAMEWORK, checked against the devices of the run.""" + module = _import_module(path) + 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 + + +def load_custom_function(ms, contexts, device_types, kind, path): + """Load the custom ``kind`` function at ``path`` into each context, in the + same order everywhere. Returns the ``MatrixElementApi`` of each context.""" + if path.endswith(".so"): + return [context.load_matrix_element(path, "") for context in contexts] + if not path.endswith(".py"): + raise RuntimeError( + f"[custom_functions] {kind} = {path!r}: expected a .so library or a .py file" + ) + + name = f"custom_{kind}" + function, framework = _python_function(path, name, device_types) + python_function = ms.PythonFunction( + function, PYTHON_INPUTS[kind], [kind], framework=framework, name=name + ) + return [context.add_matrix_element(python_function) for context in contexts] + + +def load_surrogate( + ms, contexts, backends, device_types, path, subprocess, pdg_ids, param_card +): + """Load the [surrogate] function for one subprocess into each context, and + return the ``MatrixElementApi`` of each context. + + A .so path is a UMAMI library with the inputs of the madmatrix libraries, + where ``{subprocess}`` and ``{device}`` are replaced, initialized with + ``param_card``. A .py file defines + ``custom_surrogate(momenta, pdg_ids, alpha_s)``, which returns |M|^2; + ``pdg_ids[flavor_index]`` are the PDG ids of the subprocess for each + flavor index of its matrix element.""" + if path.endswith(".so"): + return [ + context.load_matrix_element( + path.format(subprocess=subprocess, device=backend), param_card + ) + for context, backend in zip(contexts, backends) + ] + if not path.endswith(".py"): + raise RuntimeError( + f"[surrogate] function = {path!r}: expected a .so library or a .py file" + ) + function, framework = _python_function(path, "custom_surrogate", device_types) + particle_count = len(pdg_ids[0]) + if framework == "torch": + import torch + + tables = {} + + def surrogate(momenta, flavor_index, alpha_s): + device = momenta.device + if device not in tables: + tables[device] = torch.tensor(pdg_ids, dtype=torch.int32, device=device) + return function(momenta, tables[device][flavor_index.long()], alpha_s) + + else: + import numpy as np + + table = np.array(pdg_ids, dtype=np.int32) + + def surrogate(momenta, flavor_index, alpha_s): + return function(momenta, table[flavor_index], alpha_s) + + python_function = ms.PythonFunction( + surrogate, + ["momenta", "flavor_index", "alpha_s"], + ["matrix_element"], + framework=framework, + particle_count=particle_count, + name=f"custom_surrogate ({subprocess})", + ) + return [context.add_matrix_element(python_function) for context in contexts] diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/Makefile b/madgraph/iolibs/template_files/mg7/custom_functions/Makefile new file mode 100644 index 0000000000..0435e1e525 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/Makefile @@ -0,0 +1,93 @@ +# Makefile for the custom %(kind)s function, generated by bin/init_custom_func. +# Builds $(LIB) from %(body_file)s and the UMAMI wrapper. +# Add extra compiler flags to CFLAGS and libraries to LIBS, e.g. for LHAPDF: +# make CFLAGS="$$(lhapdf-config --cflags)" LIBS="$$(lhapdf-config --libs)" + +LANGUAGE = %(language)s +LIB = ../lib/custom_%(kind)s.so +BODY = %(body_file)s + +OPTFLAGS ?= -O3 +CFLAGS ?= +LIBS ?= + +# make predefines CC, CXX and FC; only replace its f77 default +ifeq ($(origin FC),default) + FC = gfortran +endif + +UNAME_S := $(shell uname -s) +ifeq ($(UNAME_S),Darwin) + SONAME_FLAGS = -Wl,-install_name,@rpath/$(notdir $(LIB)) +else + SONAME_FLAGS = -Wl,-soname,$(notdir $(LIB)) +endif + +all: $(LIB) + +ifeq ($(LANGUAGE),cuda) + + NVCC ?= nvcc + # One may pass MADGRAPH_CUDA_ARCHITECTURE (comma-separated list), as for the matrix elements + comma := , + DETECTED_CC := $(shell nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | grep -E '^[0-9]+\.[0-9]+$$' | tr -d '.' | sort -un) + MADGRAPH_CUDA_ARCHITECTURE ?= $(foreach arch,$(DETECTED_CC),$(arch)$(comma)) + CUDA_ARCH_LIST := $(subst $(comma), ,$(MADGRAPH_CUDA_ARCHITECTURE)) + ifeq ($(strip $(CUDA_ARCH_LIST)),) + CUDA_ARCH_LIST := 70 + endif + GPUARCHFLAGS := $(foreach arch,$(CUDA_ARCH_LIST),-gencode arch=compute_$(arch),code=sm_$(arch)) + +$(LIB): wrapper_gpu.cu $(BODY) umami.h + @mkdir -p $(dir $(LIB)) + $(NVCC) $(OPTFLAGS) -std=c++17 $(GPUARCHFLAGS) -Xcompiler -fPIC --shared $(CFLAGS) \ + -o $@ wrapper_gpu.cu $(SONAME_FLAGS) $(LIBS) + +else ifeq ($(LANGUAGE),hip) + + HIP_HOME ?= $(shell hipconfig --rocmpath 2>/dev/null) + HIPCC ?= $(HIP_HOME)/bin/hipcc + # One may pass MADGRAPH_HIP_ARCHITECTURE, as for the matrix elements + ifeq ($(origin MADGRAPH_HIP_ARCHITECTURE),undefined) + MADGRAPH_HIP_ARCHITECTURE := $(shell $(HIP_HOME)/bin/rocm_agent_enumerator 2>/dev/null | grep -v gfx000 | sort -u | head -1) + endif + ifeq ($(MADGRAPH_HIP_ARCHITECTURE),) + MADGRAPH_HIP_ARCHITECTURE := gfx942 + endif + +$(LIB): wrapper_gpu.cu $(BODY) umami.h + @mkdir -p $(dir $(LIB)) + $(HIPCC) $(OPTFLAGS) -std=c++17 --offload-arch=$(MADGRAPH_HIP_ARCHITECTURE) -fPIC -shared \ + $(CFLAGS) -x hip -o $@ wrapper_gpu.cu $(SONAME_FLAGS) $(LIBS) + +else + + ifeq ($(LANGUAGE),c) + BODY_COMPILE = $(CC) $(OPTFLAGS) -fPIC $(CFLAGS) + LINKER = $(CC) + else ifeq ($(LANGUAGE),cpp) + BODY_COMPILE = $(CXX) $(OPTFLAGS) -std=c++17 -fPIC $(CFLAGS) + LINKER = $(CXX) + else ifeq ($(LANGUAGE),fortran) + BODY_COMPILE = $(FC) $(OPTFLAGS) -fPIC $(CFLAGS) + LINKER = $(FC) + else + $(error unknown LANGUAGE '$(LANGUAGE)': expected c, cpp, fortran, cuda or hip) + endif + +wrapper_cpu.o: wrapper_cpu.c umami.h + $(CC) $(OPTFLAGS) -fPIC -c -o $@ wrapper_cpu.c + +body.o: $(BODY) + $(BODY_COMPILE) -c -o $@ $(BODY) + +$(LIB): wrapper_cpu.o body.o + @mkdir -p $(dir $(LIB)) + $(LINKER) -shared -o $@ wrapper_cpu.o body.o $(SONAME_FLAGS) $(LIBS) + +endif + +clean: + rm -f *.o *.mod $(LIB) + +.PHONY: all clean diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/bias.c b/madgraph/iolibs/template_files/mg7/custom_functions/bias.c new file mode 100644 index 0000000000..04e2464c61 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/bias.c @@ -0,0 +1,29 @@ +/* + * Custom event bias for MG7, in C. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the bias factor b for one phase-space point. Events are generated + * according to |w| b, and written with the weight divided by b, so physical + * distributions are unchanged while regions with large b get more events. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_bias.so. + */ + +#include + +double custom_bias(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + (void)momenta; + // Example: populate the tail of the pT distribution of the first + // final-state particle, b = (pT / 100 GeV)^2 with a floor of 1. + // double px = momenta[2][1], py = momenta[2][2]; + // return fmax((px * px + py * py) / (100. * 100.), 1.); + + return 1.; +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/bias.cpp b/madgraph/iolibs/template_files/mg7/custom_functions/bias.cpp new file mode 100644 index 0000000000..96d703705c --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/bias.cpp @@ -0,0 +1,33 @@ +/* + * Custom event bias for MG7, in C++. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the bias factor b for one phase-space point. Events are generated + * according to |w| b, and written with the weight divided by b, so physical + * distributions are unchanged while regions with large b get more events. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_bias.so. + */ + +#include + +using std::fmax; +using std::pow; +using std::sqrt; + +extern "C" double custom_bias(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + (void)momenta; + // Example: populate the tail of the pT distribution of the first + // final-state particle, b = (pT / 100 GeV)^2 with a floor of 1. + // double px = momenta[2][1], py = momenta[2][2]; + // return fmax((px * px + py * py) / (100. * 100.), 1.); + + return 1.; +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/bias.cu b/madgraph/iolibs/template_files/mg7/custom_functions/bias.cu new file mode 100644 index 0000000000..f555ba2305 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/bias.cu @@ -0,0 +1,28 @@ +/* + * Custom event bias for MG7, in CUDA or HIP. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the bias factor b for one phase-space point. Events are generated + * according to |w| b, and written with the weight divided by b, so physical + * distributions are unchanged while regions with large b get more events. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_bias.so. + * The function runs on the GPU, one thread per event. + */ + +__device__ double custom_bias(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + (void)momenta; + // Example: populate the tail of the pT distribution of the first + // final-state particle, b = (pT / 100 GeV)^2 with a floor of 1. + // double px = momenta[2][1], py = momenta[2][2]; + // return fmax((px * px + py * py) / (100. * 100.), 1.); + + return 1.; +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/bias.f90 b/madgraph/iolibs/template_files/mg7/custom_functions/bias.f90 new file mode 100644 index 0000000000..c2d63721d4 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/bias.f90 @@ -0,0 +1,28 @@ +! Custom event bias for MG7, in Fortran. Generated by bin/init_custom_func. +! +! It is called once per phase-space point and returns +! the bias factor b for one phase-space point. Events are generated +! according to |w| b, and written with the weight divided by b, so physical +! distributions are unchanged while regions with large b get more events. +! +! Parameters: +! particle_count number of external particles, incoming ones first +! pdg_ids PDG ids of the external particles +! momenta momenta(:, i) = (E, px, py, pz) of particle i in GeV, lab frame +! +! Run `make` in this directory to build ../lib/custom_bias.so. + +function custom_bias(particle_count, pdg_ids, momenta) bind(C, name="custom_bias") result(value) + use iso_c_binding + implicit none + integer(c_int), value :: particle_count + integer(c_int), intent(in) :: pdg_ids(particle_count) + real(c_double), intent(in) :: momenta(0:3, particle_count) + real(c_double) :: value + + ! Example: populate the tail of the pT distribution of the first + ! final-state particle, b = (pT / 100 GeV)^2 with a floor of 1. + ! value = max((momenta(1, 3)**2 + momenta(2, 3)**2) / 100d0**2, 1d0) + + value = 1d0 +end function custom_bias diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/bias_numpy.py b/madgraph/iolibs/template_files/mg7/custom_functions/bias_numpy.py new file mode 100644 index 0000000000..f2e97421b1 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/bias_numpy.py @@ -0,0 +1,34 @@ +""" +Custom event bias for MG7, in Python with numpy. Generated by bin/init_custom_func. + +custom_bias is called once per batch of phase-space points, with one array +entry per point. It returns a float64 array of shape (batch,) with +the bias factor b per event. Events are generated according to |w| b, +and written with the weight divided by b, so physical distributions are +unchanged while regions with large b get more events. + +It runs on the CPU only. Use the python-torch template for GPU runs. + +Parameters: +momenta float64 array of shape (batch, particle_count, 4), the momenta + (E, px, py, pz) of the external particles in GeV, lab frame, + incoming ones first +pdg_ids int32 array of shape (batch, particle_count), their PDG ids + +Enable it in Cards/run_card.toml with +[custom_functions] +bias = "custom_bias/custom_bias.py" +""" + +import numpy as np + +FRAMEWORK = "numpy" + + +def custom_bias(momenta, pdg_ids): + # Example: populate the tail of the pT distribution of the first + # final-state particle, b = (pT / 100 GeV)^2 with a floor of 1. + # pt2 = momenta[:, 2, 1] ** 2 + momenta[:, 2, 2] ** 2 + # return np.maximum(pt2 / 100. ** 2, 1.) + + return np.ones(len(momenta)) diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/bias_torch.py b/madgraph/iolibs/template_files/mg7/custom_functions/bias_torch.py new file mode 100644 index 0000000000..79777caca8 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/bias_torch.py @@ -0,0 +1,35 @@ +""" +Custom event bias for MG7, in Python with PyTorch. Generated by bin/init_custom_func. + +custom_bias is called once per batch of phase-space points, with one array +entry per point. It returns a float64 tensor of shape (batch,) with +the bias factor b per event. Events are generated according to |w| b, +and written with the weight divided by b, so physical distributions are +unchanged while regions with large b get more events. + +It runs on the device of the run, CPU or GPU. On the CPU, MG7 limits torch +to one thread, since it already calls the function from several threads. + +Parameters: +momenta float64 tensor of shape (batch, particle_count, 4), the momenta + (E, px, py, pz) of the external particles in GeV, lab frame, + incoming ones first +pdg_ids int32 tensor of shape (batch, particle_count), their PDG ids + +Enable it in Cards/run_card.toml with +[custom_functions] +bias = "custom_bias/custom_bias.py" +""" + +import torch + +FRAMEWORK = "torch" + + +def custom_bias(momenta, pdg_ids): + # Example: populate the tail of the pT distribution of the first + # final-state particle, b = (pT / 100 GeV)^2 with a floor of 1. + # pt2 = momenta[:, 2, 1] ** 2 + momenta[:, 2, 2] ** 2 + # return torch.clamp(pt2 / 100. ** 2, min=1.) + + return torch.ones(len(momenta), dtype=torch.float64, device=momenta.device) diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/cut.c b/madgraph/iolibs/template_files/mg7/custom_functions/cut.c new file mode 100644 index 0000000000..0affcab7b4 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/cut.c @@ -0,0 +1,27 @@ +/* + * Custom cut for MG7, in C. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the cut value for one phase-space point: 1 keeps the event, 0 rejects it. + * Values in between act as a weight factor. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_cut.so. + */ + +#include + +double custom_cut(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + (void)momenta; + // Example: require pT > 20 GeV for the first final-state particle. + // double px = momenta[2][1], py = momenta[2][2]; + // return px * px + py * py > 20. * 20. ? 1. : 0.; + + return 1.; +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/cut.cpp b/madgraph/iolibs/template_files/mg7/custom_functions/cut.cpp new file mode 100644 index 0000000000..67cd66014c --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/cut.cpp @@ -0,0 +1,31 @@ +/* + * Custom cut for MG7, in C++. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the cut value for one phase-space point: 1 keeps the event, 0 rejects it. + * Values in between act as a weight factor. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_cut.so. + */ + +#include + +using std::fmax; +using std::pow; +using std::sqrt; + +extern "C" double custom_cut(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + (void)momenta; + // Example: require pT > 20 GeV for the first final-state particle. + // double px = momenta[2][1], py = momenta[2][2]; + // return px * px + py * py > 20. * 20. ? 1. : 0.; + + return 1.; +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/cut.cu b/madgraph/iolibs/template_files/mg7/custom_functions/cut.cu new file mode 100644 index 0000000000..c20b6987fc --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/cut.cu @@ -0,0 +1,26 @@ +/* + * Custom cut for MG7, in CUDA or HIP. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the cut value for one phase-space point: 1 keeps the event, 0 rejects it. + * Values in between act as a weight factor. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_cut.so. + * The function runs on the GPU, one thread per event. + */ + +__device__ double custom_cut(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + (void)momenta; + // Example: require pT > 20 GeV for the first final-state particle. + // double px = momenta[2][1], py = momenta[2][2]; + // return px * px + py * py > 20. * 20. ? 1. : 0.; + + return 1.; +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/cut.f90 b/madgraph/iolibs/template_files/mg7/custom_functions/cut.f90 new file mode 100644 index 0000000000..3a9e69e137 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/cut.f90 @@ -0,0 +1,30 @@ +! Custom cut for MG7, in Fortran. Generated by bin/init_custom_func. +! +! It is called once per phase-space point and returns +! the cut value for one phase-space point: 1 keeps the event, 0 rejects it. +! Values in between act as a weight factor. +! +! Parameters: +! particle_count number of external particles, incoming ones first +! pdg_ids PDG ids of the external particles +! momenta momenta(:, i) = (E, px, py, pz) of particle i in GeV, lab frame +! +! Run `make` in this directory to build ../lib/custom_cut.so. + +function custom_cut(particle_count, pdg_ids, momenta) bind(C, name="custom_cut") result(value) + use iso_c_binding + implicit none + integer(c_int), value :: particle_count + integer(c_int), intent(in) :: pdg_ids(particle_count) + real(c_double), intent(in) :: momenta(0:3, particle_count) + real(c_double) :: value + + ! Example: require pT > 20 GeV for the first final-state particle. + ! if (momenta(1, 3)**2 + momenta(2, 3)**2 > 20d0**2) then + ! value = 1d0 + ! else + ! value = 0d0 + ! end if + + value = 1d0 +end function custom_cut diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/cut_numpy.py b/madgraph/iolibs/template_files/mg7/custom_functions/cut_numpy.py new file mode 100644 index 0000000000..9cc1af1d0e --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/cut_numpy.py @@ -0,0 +1,32 @@ +""" +Custom cut for MG7, in Python with numpy. Generated by bin/init_custom_func. + +custom_cut is called once per batch of phase-space points, with one array +entry per point. It returns a float64 array of shape (batch,) with +the cut value per event: 1 keeps the event, 0 rejects it. Values in +between act as a weight factor. + +It runs on the CPU only. Use the python-torch template for GPU runs. + +Parameters: +momenta float64 array of shape (batch, particle_count, 4), the momenta + (E, px, py, pz) of the external particles in GeV, lab frame, + incoming ones first +pdg_ids int32 array of shape (batch, particle_count), their PDG ids + +Enable it in Cards/run_card.toml with +[custom_functions] +cut = "custom_cut/custom_cut.py" +""" + +import numpy as np + +FRAMEWORK = "numpy" + + +def custom_cut(momenta, pdg_ids): + # Example: require pT > 20 GeV for the first final-state particle. + # pt = np.hypot(momenta[:, 2, 1], momenta[:, 2, 2]) + # return np.where(pt > 20., 1., 0.) + + return np.ones(len(momenta)) diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/cut_torch.py b/madgraph/iolibs/template_files/mg7/custom_functions/cut_torch.py new file mode 100644 index 0000000000..9010d689a1 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/cut_torch.py @@ -0,0 +1,33 @@ +""" +Custom cut for MG7, in Python with PyTorch. Generated by bin/init_custom_func. + +custom_cut is called once per batch of phase-space points, with one array +entry per point. It returns a float64 tensor of shape (batch,) with +the cut value per event: 1 keeps the event, 0 rejects it. Values in +between act as a weight factor. + +It runs on the device of the run, CPU or GPU. On the CPU, MG7 limits torch +to one thread, since it already calls the function from several threads. + +Parameters: +momenta float64 tensor of shape (batch, particle_count, 4), the momenta + (E, px, py, pz) of the external particles in GeV, lab frame, + incoming ones first +pdg_ids int32 tensor of shape (batch, particle_count), their PDG ids + +Enable it in Cards/run_card.toml with +[custom_functions] +cut = "custom_cut/custom_cut.py" +""" + +import torch + +FRAMEWORK = "torch" + + +def custom_cut(momenta, pdg_ids): + # Example: require pT > 20 GeV for the first final-state particle. + # pt = torch.hypot(momenta[:, 2, 1], momenta[:, 2, 2]) + # return (pt > 20.).double() + + return torch.ones(len(momenta), dtype=torch.float64, device=momenta.device) diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/pdf.c b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.c new file mode 100644 index 0000000000..6cb5f9fc2c --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.c @@ -0,0 +1,27 @@ +/* + * Custom parton density for MG7, in C. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * x f(x, q), the momentum fraction times the parton density, for the parton + * with PDG id pdg_id (21 for the gluon) in the beam beam_index (0 or 1), at + * momentum fraction x and factorization scale q in GeV. The strong coupling + * still comes from the PDF set in the run card. + * + * Parameters: + * beam_index index of the beam, 0 or 1 + * pdg_id PDG id of the parton + * x momentum fraction + * q factorization scale in GeV + * + * Run `make` in this directory to build ../lib/custom_pdf.so. + */ + +#include + +double custom_pdf(int beam_index, int pdg_id, double x, double q) { + // default: a toy density, replace it with your own + (void)beam_index; + (void)q; + double norm = pdg_id == 21 ? 3. : 0.5; + return norm * sqrt(x) * pow(1. - x, 3); +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/pdf.cpp b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.cpp new file mode 100644 index 0000000000..c605a3bb92 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.cpp @@ -0,0 +1,36 @@ +/* + * Custom parton density for MG7, in C++. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * x f(x, q), the momentum fraction times the parton density, for the parton + * with PDG id pdg_id (21 for the gluon) in the beam beam_index (0 or 1), at + * momentum fraction x and factorization scale q in GeV. The strong coupling + * still comes from the PDF set in the run card. + * + * Parameters: + * beam_index index of the beam, 0 or 1 + * pdg_id PDG id of the parton + * x momentum fraction + * q factorization scale in GeV + * + * Run `make` in this directory to build ../lib/custom_pdf.so. + */ + +#include + +using std::fmax; +using std::pow; +using std::sqrt; + +extern "C" double custom_pdf(int beam_index, int pdg_id, double x, double q) { + // Example: call LHAPDF. Add #include "LHAPDF/LHAPDF.h" above, and the + // output of `lhapdf-config --cflags --libs` to CFLAGS and LIBS in the Makefile. + // static LHAPDF::PDF* pdf = LHAPDF::mkPDF("NNPDF23_lo_as_0130_qed", 0); + // return pdf->xfxQ(pdg_id, x, q); + + // default: a toy density, replace it with your own + (void)beam_index; + (void)q; + double norm = pdg_id == 21 ? 3. : 0.5; + return norm * sqrt(x) * pow(1. - x, 3); +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/pdf.cu b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.cu new file mode 100644 index 0000000000..dec659e12a --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.cu @@ -0,0 +1,26 @@ +/* + * Custom parton density for MG7, in CUDA or HIP. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * x f(x, q), the momentum fraction times the parton density, for the parton + * with PDG id pdg_id (21 for the gluon) in the beam beam_index (0 or 1), at + * momentum fraction x and factorization scale q in GeV. The strong coupling + * still comes from the PDF set in the run card. + * + * Parameters: + * beam_index index of the beam, 0 or 1 + * pdg_id PDG id of the parton + * x momentum fraction + * q factorization scale in GeV + * + * Run `make` in this directory to build ../lib/custom_pdf.so. + * The function runs on the GPU, one thread per event. + */ + +__device__ double custom_pdf(int beam_index, int pdg_id, double x, double q) { + // default: a toy density, replace it with your own + (void)beam_index; + (void)q; + double norm = pdg_id == 21 ? 3. : 0.5; + return norm * sqrt(x) * pow(1. - x, 3); +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/pdf.f90 b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.f90 new file mode 100644 index 0000000000..f9a9e5ebfd --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/pdf.f90 @@ -0,0 +1,32 @@ +! Custom parton density for MG7, in Fortran. Generated by bin/init_custom_func. +! +! It is called once per phase-space point and returns +! x f(x, q), the momentum fraction times the parton density, for the parton +! with PDG id pdg_id (21 for the gluon) in the beam beam_index (0 or 1), at +! momentum fraction x and factorization scale q in GeV. The strong coupling +! still comes from the PDF set in the run card. +! +! Parameters: +! beam_index index of the beam, 0 or 1 +! pdg_id PDG id of the parton +! x momentum fraction +! q factorization scale in GeV +! +! Run `make` in this directory to build ../lib/custom_pdf.so. + +function custom_pdf(beam_index, pdg_id, x, q) bind(C, name="custom_pdf") result(value) + use iso_c_binding + implicit none + integer(c_int), value :: beam_index, pdg_id + real(c_double), value :: x, q + real(c_double) :: value + real(c_double) :: norm + + ! default: a toy density, replace it with your own + if (pdg_id == 21) then + norm = 3d0 + else + norm = 0.5d0 + end if + value = norm * sqrt(x) * (1d0 - x)**3 +end function custom_pdf diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/pdf_numpy.py b/madgraph/iolibs/template_files/mg7/custom_functions/pdf_numpy.py new file mode 100644 index 0000000000..309fd73369 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/pdf_numpy.py @@ -0,0 +1,37 @@ +""" +Custom parton density for MG7, in Python with numpy. Generated by bin/init_custom_func. + +custom_pdf is called once per batch of phase-space points, with one array +entry per point. It returns a float64 array of shape (batch,) with +x f(x, q), the momentum fraction times the parton density, per event. +The strong coupling still comes from the PDF set in the run card. + +It runs on the CPU only. Use the python-torch template for GPU runs. + +Parameters: +beam_index int32 array of shape (batch,), the beam, 0 or 1 +pdg_id int32 array of shape (batch,), the PDG id of the parton +x float64 array of shape (batch,), the momentum fraction +q float64 array of shape (batch,), the factorization scale in GeV + +Enable it in Cards/run_card.toml with +[custom_functions] +pdf = "custom_pdf/custom_pdf.py" +""" + +import numpy as np + +# Example: LHAPDF through its Python module. Uncomment these lines and the +# return statement in custom_pdf below. +# import lhapdf +# pdf_set = lhapdf.mkPDF("NNPDF23_lo_as_0130_qed", 0) + +FRAMEWORK = "numpy" + + +def custom_pdf(beam_index, pdg_id, x, q): + # return np.array([pdf_set.xfxQ(int(i), a, b) for i, a, b in zip(pdg_id, x, q)]) + + # default: a toy density, replace it with your own + norm = np.where(pdg_id == 21, 3., 0.5) + return norm * np.sqrt(x) * (1. - x) ** 3 diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/pdf_torch.py b/madgraph/iolibs/template_files/mg7/custom_functions/pdf_torch.py new file mode 100644 index 0000000000..b3d6c3ac1d --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/pdf_torch.py @@ -0,0 +1,31 @@ +""" +Custom parton density for MG7, in Python with PyTorch. Generated by bin/init_custom_func. + +custom_pdf is called once per batch of phase-space points, with one array +entry per point. It returns a float64 tensor of shape (batch,) with +x f(x, q), the momentum fraction times the parton density, per event. +The strong coupling still comes from the PDF set in the run card. + +It runs on the device of the run, CPU or GPU. On the CPU, MG7 limits torch +to one thread, since it already calls the function from several threads. + +Parameters: +beam_index int32 tensor of shape (batch,), the beam, 0 or 1 +pdg_id int32 tensor of shape (batch,), the PDG id of the parton +x float64 tensor of shape (batch,), the momentum fraction +q float64 tensor of shape (batch,), the factorization scale in GeV + +Enable it in Cards/run_card.toml with +[custom_functions] +pdf = "custom_pdf/custom_pdf.py" +""" + +import torch + +FRAMEWORK = "torch" + + +def custom_pdf(beam_index, pdg_id, x, q): + # default: a toy density, replace it with your own + norm = torch.where(pdg_id == 21, 3., 0.5).double() + return norm * torch.sqrt(x) * (1. - x) ** 3 diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/scale.c b/madgraph/iolibs/template_files/mg7/custom_functions/scale.c new file mode 100644 index 0000000000..5ec5fc2a19 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/scale.c @@ -0,0 +1,36 @@ +/* + * Custom energy scale for MG7, in C. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the renormalization and factorization scale in GeV for one phase-space + * point. The run card scale_factor is applied on top, and fixed scales still + * take precedence. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_scale.so. + */ + +#include + +double custom_scale(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + // Example: H_T / 2, the scalar sum of the final-state transverse masses. + // double ht = 0.; + // for (int i = 2; i < particle_count; ++i) { + // double e = momenta[i][0], pz = momenta[i][3]; + // ht += sqrt(fmax(e * e - pz * pz, 0.)); + // } + // return ht / 2.; + + // default: partonic center-of-mass energy sqrt(s_hat) + double e = momenta[0][0] + momenta[1][0]; + double px = momenta[0][1] + momenta[1][1]; + double py = momenta[0][2] + momenta[1][2]; + double pz = momenta[0][3] + momenta[1][3]; + return sqrt(fmax(e * e - px * px - py * py - pz * pz, 0.)); +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/scale.cpp b/madgraph/iolibs/template_files/mg7/custom_functions/scale.cpp new file mode 100644 index 0000000000..b91fc25c85 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/scale.cpp @@ -0,0 +1,40 @@ +/* + * Custom energy scale for MG7, in C++. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the renormalization and factorization scale in GeV for one phase-space + * point. The run card scale_factor is applied on top, and fixed scales still + * take precedence. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_scale.so. + */ + +#include + +using std::fmax; +using std::pow; +using std::sqrt; + +extern "C" double custom_scale(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + // Example: H_T / 2, the scalar sum of the final-state transverse masses. + // double ht = 0.; + // for (int i = 2; i < particle_count; ++i) { + // double e = momenta[i][0], pz = momenta[i][3]; + // ht += sqrt(fmax(e * e - pz * pz, 0.)); + // } + // return ht / 2.; + + // default: partonic center-of-mass energy sqrt(s_hat) + double e = momenta[0][0] + momenta[1][0]; + double px = momenta[0][1] + momenta[1][1]; + double py = momenta[0][2] + momenta[1][2]; + double pz = momenta[0][3] + momenta[1][3]; + return sqrt(fmax(e * e - px * px - py * py - pz * pz, 0.)); +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/scale.cu b/madgraph/iolibs/template_files/mg7/custom_functions/scale.cu new file mode 100644 index 0000000000..05679b90ce --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/scale.cu @@ -0,0 +1,35 @@ +/* + * Custom energy scale for MG7, in CUDA or HIP. Generated by bin/init_custom_func. + * + * It is called once per phase-space point and returns + * the renormalization and factorization scale in GeV for one phase-space + * point. The run card scale_factor is applied on top, and fixed scales still + * take precedence. + * + * Parameters: + * particle_count number of external particles, incoming ones first + * pdg_ids PDG ids of the external particles + * momenta momenta[i] = (E, px, py, pz) of particle i in GeV, lab frame + * + * Run `make` in this directory to build ../lib/custom_scale.so. + * The function runs on the GPU, one thread per event. + */ + +__device__ double custom_scale(int particle_count, const int* pdg_ids, const double momenta[][4]) { + (void)particle_count; + (void)pdg_ids; + // Example: H_T / 2, the scalar sum of the final-state transverse masses. + // double ht = 0.; + // for (int i = 2; i < particle_count; ++i) { + // double e = momenta[i][0], pz = momenta[i][3]; + // ht += sqrt(fmax(e * e - pz * pz, 0.)); + // } + // return ht / 2.; + + // default: partonic center-of-mass energy sqrt(s_hat) + double e = momenta[0][0] + momenta[1][0]; + double px = momenta[0][1] + momenta[1][1]; + double py = momenta[0][2] + momenta[1][2]; + double pz = momenta[0][3] + momenta[1][3]; + return sqrt(fmax(e * e - px * px - py * py - pz * pz, 0.)); +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/scale.f90 b/madgraph/iolibs/template_files/mg7/custom_functions/scale.f90 new file mode 100644 index 0000000000..6fcc401b1d --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/scale.f90 @@ -0,0 +1,35 @@ +! Custom energy scale for MG7, in Fortran. Generated by bin/init_custom_func. +! +! It is called once per phase-space point and returns +! the renormalization and factorization scale in GeV for one phase-space +! point. The run card scale_factor is applied on top, and fixed scales still +! take precedence. +! +! Parameters: +! particle_count number of external particles, incoming ones first +! pdg_ids PDG ids of the external particles +! momenta momenta(:, i) = (E, px, py, pz) of particle i in GeV, lab frame +! +! Run `make` in this directory to build ../lib/custom_scale.so. + +function custom_scale(particle_count, pdg_ids, momenta) bind(C, name="custom_scale") result(value) + use iso_c_binding + implicit none + integer(c_int), value :: particle_count + integer(c_int), intent(in) :: pdg_ids(particle_count) + real(c_double), intent(in) :: momenta(0:3, particle_count) + real(c_double) :: value + real(c_double) :: p(0:3) + + ! Example: H_T / 2, the scalar sum of the final-state transverse masses. + ! integer :: i + ! value = 0d0 + ! do i = 3, particle_count + ! value = value + sqrt(max(momenta(0, i)**2 - momenta(3, i)**2, 0d0)) + ! end do + ! value = value / 2d0 + + ! default: partonic center-of-mass energy sqrt(s_hat) + p = momenta(:, 1) + momenta(:, 2) + value = sqrt(max(p(0)**2 - p(1)**2 - p(2)**2 - p(3)**2, 0d0)) +end function custom_scale diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/scale_numpy.py b/madgraph/iolibs/template_files/mg7/custom_functions/scale_numpy.py new file mode 100644 index 0000000000..5d8b3bea49 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/scale_numpy.py @@ -0,0 +1,35 @@ +""" +Custom energy scale for MG7, in Python with numpy. Generated by bin/init_custom_func. + +custom_scale is called once per batch of phase-space points, with one array +entry per point. It returns a float64 array of shape (batch,) with +the renormalization and factorization scale in GeV per event. The run +card scale_factor is applied on top, and fixed scales still take precedence. + +It runs on the CPU only. Use the python-torch template for GPU runs. + +Parameters: +momenta float64 array of shape (batch, particle_count, 4), the momenta + (E, px, py, pz) of the external particles in GeV, lab frame, + incoming ones first +pdg_ids int32 array of shape (batch, particle_count), their PDG ids + +Enable it in Cards/run_card.toml with +[custom_functions] +scale = "custom_scale/custom_scale.py" +""" + +import numpy as np + +FRAMEWORK = "numpy" + + +def custom_scale(momenta, pdg_ids): + # Example: H_T / 2, the scalar sum of the final-state transverse masses. + # final = momenta[:, 2:] + # mt = np.sqrt(np.maximum(final[..., 0] ** 2 - final[..., 3] ** 2, 0.)) + # return mt.sum(axis=1) / 2. + + # default: partonic center-of-mass energy sqrt(s_hat) + p = momenta[:, 0] + momenta[:, 1] + return np.sqrt(np.maximum(p[:, 0] ** 2 - (p[:, 1:] ** 2).sum(axis=1), 0.)) diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/scale_torch.py b/madgraph/iolibs/template_files/mg7/custom_functions/scale_torch.py new file mode 100644 index 0000000000..d746a3b6e3 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/scale_torch.py @@ -0,0 +1,36 @@ +""" +Custom energy scale for MG7, in Python with PyTorch. Generated by bin/init_custom_func. + +custom_scale is called once per batch of phase-space points, with one array +entry per point. It returns a float64 tensor of shape (batch,) with +the renormalization and factorization scale in GeV per event. The run +card scale_factor is applied on top, and fixed scales still take precedence. + +It runs on the device of the run, CPU or GPU. On the CPU, MG7 limits torch +to one thread, since it already calls the function from several threads. + +Parameters: +momenta float64 tensor of shape (batch, particle_count, 4), the momenta + (E, px, py, pz) of the external particles in GeV, lab frame, + incoming ones first +pdg_ids int32 tensor of shape (batch, particle_count), their PDG ids + +Enable it in Cards/run_card.toml with +[custom_functions] +scale = "custom_scale/custom_scale.py" +""" + +import torch + +FRAMEWORK = "torch" + + +def custom_scale(momenta, pdg_ids): + # Example: H_T / 2, the scalar sum of the final-state transverse masses. + # final = momenta[:, 2:] + # mt = torch.sqrt(torch.clamp(final[..., 0] ** 2 - final[..., 3] ** 2, min=0.)) + # return mt.sum(dim=1) / 2. + + # default: partonic center-of-mass energy sqrt(s_hat) + p = momenta[:, 0] + momenta[:, 1] + return torch.sqrt(torch.clamp(p[:, 0] ** 2 - (p[:, 1:] ** 2).sum(dim=1), min=0.)) diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/wrapper_cpu.c b/madgraph/iolibs/template_files/mg7/custom_functions/wrapper_cpu.c new file mode 100644 index 0000000000..dbca693923 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/wrapper_cpu.c @@ -0,0 +1,195 @@ +/* + * UMAMI wrapper for the custom %(kind)s function, generated by bin/init_custom_func. + * It unpacks the batched UMAMI inputs and calls %(function_name)s once per event. + * Write your function in %(body_file)s, there is no need to edit this file. + */ + +#include "umami.h" + +#include + +#define CUSTOM_IS_PDF %(is_pdf)d +#define CUSTOM_OUTPUT_KEY %(output_key)s + +#if CUSTOM_IS_PDF +double %(function_name)s(int beam_index, int pdg_id, double x, double q); +static const int input_keys[] = { + UMAMI_IN_MOMENTUM_FRACTION, UMAMI_IN_SCALE, UMAMI_IN_PDG_ID, UMAMI_IN_BEAM_INDEX +}; +#else +double %(function_name)s(int particle_count, const int* pdg_ids, const double (*momenta)[4]); +static const int input_keys[] = {UMAMI_IN_MOMENTA, UMAMI_IN_PARTICLE_COUNT, UMAMI_IN_PDG_IDS}; +#endif +#define INPUT_COUNT (sizeof(input_keys) / sizeof(input_keys[0])) + +static bool supported_inputs[UMAMI_INPUT_KEY_COUNT]; +static bool supported_outputs[UMAMI_OUTPUT_KEY_COUNT]; +static int dummy_instance; + +static void fill_key_tables(void) { + for (size_t i = 0; i < INPUT_COUNT; ++i) { + supported_inputs[input_keys[i]] = true; + } + supported_outputs[CUSTOM_OUTPUT_KEY] = true; +} + +UmamiStatus umami_get_meta(UmamiMetaKey meta_key, void* result) { + if (meta_key == UMAMI_META_DEVICE) { + *(UmamiDevice*)result = UMAMI_DEVICE_CPU; + return UMAMI_SUCCESS; + } + return UMAMI_ERROR_UNSUPPORTED_META; +} + +UmamiStatus umami_supported_inputs(bool const** supported, int* count) { + fill_key_tables(); + *supported = supported_inputs; + *count = UMAMI_INPUT_KEY_COUNT; + return UMAMI_SUCCESS; +} + +UmamiStatus umami_required_inputs(bool const** required, int* count) { + return umami_supported_inputs(required, count); +} + +UmamiStatus umami_supported_outputs(bool const** supported, int* count) { + fill_key_tables(); + *supported = supported_outputs; + *count = UMAMI_OUTPUT_KEY_COUNT; + return UMAMI_SUCCESS; +} + +UmamiStatus umami_initialize(UmamiHandle* handle, char const* param_card_path) { + (void)param_card_path; + *handle = &dummy_instance; + return UMAMI_SUCCESS; +} + +UmamiStatus umami_set_parameter( + UmamiHandle handle, char const* name, double parameter_real, double parameter_imag +) { + (void)handle; + (void)name; + (void)parameter_real; + (void)parameter_imag; + return UMAMI_ERROR_NOT_IMPLEMENTED; +} + +UmamiStatus umami_get_parameter( + UmamiHandle handle, char const* name, double* parameter_real, double* parameter_imag +) { + (void)handle; + (void)name; + (void)parameter_real; + (void)parameter_imag; + return UMAMI_ERROR_NOT_IMPLEMENTED; +} + +UmamiStatus umami_matrix_element( + UmamiHandle handle, + size_t count, + size_t stride, + size_t offset, + size_t input_count, + UmamiInputKey const* input_keys_in, + void const* const* inputs, + size_t output_count, + UmamiOutputKey const* output_keys, + void* const* outputs +) { + (void)handle; + double* output = NULL; + for (size_t i = 0; i < output_count; ++i) { + if (output_keys[i] == CUSTOM_OUTPUT_KEY) { + output = (double*)outputs[i]; + } else if (output_keys[i] != UMAMI_OUT_GPU_STREAM) { + return UMAMI_ERROR_UNSUPPORTED_OUTPUT; + } + } + if (output == NULL) { + return UMAMI_SUCCESS; + } + +#if CUSTOM_IS_PDF + const double* x = NULL; + const double* q = NULL; + const int* pdg_id = NULL; + const int* beam_index = NULL; + for (size_t i = 0; i < input_count; ++i) { + switch (input_keys_in[i]) { + case UMAMI_IN_MOMENTUM_FRACTION: + x = (const double*)inputs[i]; + break; + case UMAMI_IN_SCALE: + q = (const double*)inputs[i]; + break; + case UMAMI_IN_PDG_ID: + pdg_id = (const int*)inputs[i]; + break; + case UMAMI_IN_BEAM_INDEX: + beam_index = (const int*)inputs[i]; + break; + default: + return UMAMI_ERROR_UNSUPPORTED_INPUT; + } + } + if (x == NULL || q == NULL || pdg_id == NULL || beam_index == NULL) { + return UMAMI_ERROR_MISSING_INPUT; + } + for (size_t i = offset; i < offset + count; ++i) { + output[i] = %(function_name)s(beam_index[i], pdg_id[i], x[i], q[i]); + } +#else + const double* momenta = NULL; + const int* particle_count = NULL; + const int* pdg_ids = NULL; + for (size_t i = 0; i < input_count; ++i) { + switch (input_keys_in[i]) { + case UMAMI_IN_MOMENTA: + momenta = (const double*)inputs[i]; + break; + case UMAMI_IN_PARTICLE_COUNT: + particle_count = (const int*)inputs[i]; + break; + case UMAMI_IN_PDG_IDS: + pdg_ids = (const int*)inputs[i]; + break; + default: + return UMAMI_ERROR_UNSUPPORTED_INPUT; + } + } + if (momenta == NULL || particle_count == NULL || pdg_ids == NULL) { + return UMAMI_ERROR_MISSING_INPUT; + } + if (count == 0) { + return UMAMI_SUCCESS; + } + // the particle count is the same for every event of a batch + int n = particle_count[offset]; + int* event_pdg_ids = malloc(n * sizeof(int)); + double (*event_momenta)[4] = malloc(n * sizeof(double[4])); + if (event_pdg_ids == NULL || event_momenta == NULL) { + free(event_pdg_ids); + free(event_momenta); + return UMAMI_ERROR; + } + for (size_t i = offset; i < offset + count; ++i) { + // column-major layout with the batch dimension contiguous + for (int j = 0; j < n; ++j) { + event_pdg_ids[j] = pdg_ids[stride * j + i]; + for (int k = 0; k < 4; ++k) { + event_momenta[j][k] = momenta[stride * (n * k + j) + i]; + } + } + output[i] = %(function_name)s(n, event_pdg_ids, (const double (*)[4])event_momenta); + } + free(event_pdg_ids); + free(event_momenta); +#endif + return UMAMI_SUCCESS; +} + +UmamiStatus umami_free(UmamiHandle handle) { + (void)handle; + return UMAMI_SUCCESS; +} diff --git a/madgraph/iolibs/template_files/mg7/custom_functions/wrapper_gpu.cu b/madgraph/iolibs/template_files/mg7/custom_functions/wrapper_gpu.cu new file mode 100644 index 0000000000..c0033ce034 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/custom_functions/wrapper_gpu.cu @@ -0,0 +1,226 @@ +// UMAMI wrapper for the custom %(kind)s function, generated by bin/init_custom_func. +// It unpacks the batched UMAMI inputs and calls %(function_name)s once per event, +// one GPU thread per event. Write your function in %(body_file)s, there is no need +// to edit this file. + +#include "umami.h" + +#ifdef __HIPCC__ +#include +typedef hipStream_t gpuStream_t; +#define gpuGetLastError hipGetLastError +#define gpuSuccess hipSuccess +#define UMAMI_GPU_DEVICE UMAMI_DEVICE_HIP +#else +#include +typedef cudaStream_t gpuStream_t; +#define gpuGetLastError cudaGetLastError +#define gpuSuccess cudaSuccess +#define UMAMI_GPU_DEVICE UMAMI_DEVICE_CUDA +#endif + +#define CUSTOM_IS_PDF %(is_pdf)d +#define CUSTOM_OUTPUT_KEY %(output_key)s + +// the function body is compiled into this translation unit +#include "%(body_file)s" + +namespace { + +constexpr int thread_count = 256; +int dummy_instance; + +#if CUSTOM_IS_PDF +constexpr UmamiInputKey input_keys[] = { + UMAMI_IN_MOMENTUM_FRACTION, UMAMI_IN_SCALE, UMAMI_IN_PDG_ID, UMAMI_IN_BEAM_INDEX +}; + +__global__ void evaluate( + std::size_t count, + const double* x, + const double* q, + const int* pdg_id, + const int* beam_index, + double* output +) { + std::size_t i = blockDim.x * blockIdx.x + threadIdx.x; + if (i < count) { + output[i] = %(function_name)s(beam_index[i], pdg_id[i], x[i], q[i]); + } +} +#else +constexpr UmamiInputKey input_keys[] = { + UMAMI_IN_MOMENTA, UMAMI_IN_PARTICLE_COUNT, UMAMI_IN_PDG_IDS +}; +// largest particle count supported by the GPU wrapper +constexpr int max_particle_count = 32; + +__global__ void evaluate( + std::size_t count, + std::size_t stride, + const double* momenta, + const int* particle_count, + const int* pdg_ids, + double* output +) { + std::size_t i = blockDim.x * blockIdx.x + threadIdx.x; + if (i >= count) { + return; + } + int n = particle_count[i]; + int event_pdg_ids[max_particle_count]; + double event_momenta[max_particle_count][4]; + // column-major layout with the batch dimension contiguous + for (int j = 0; j < n; ++j) { + event_pdg_ids[j] = pdg_ids[stride * j + i]; + for (int k = 0; k < 4; ++k) { + event_momenta[j][k] = momenta[stride * (n * k + j) + i]; + } + } + output[i] = %(function_name)s(n, event_pdg_ids, event_momenta); +} +#endif + +struct KeyTables { + bool inputs[UMAMI_INPUT_KEY_COUNT] = {}; + bool outputs[UMAMI_OUTPUT_KEY_COUNT] = {}; + KeyTables() { + for (auto key : input_keys) { + inputs[key] = true; + } + outputs[CUSTOM_OUTPUT_KEY] = true; + outputs[UMAMI_OUT_GPU_STREAM] = true; + } +}; +const KeyTables key_tables; + +} // namespace + +extern "C" { + +UmamiStatus umami_get_meta(UmamiMetaKey meta_key, void* result) { + if (meta_key == UMAMI_META_DEVICE) { + *static_cast(result) = UMAMI_GPU_DEVICE; + return UMAMI_SUCCESS; + } + return UMAMI_ERROR_UNSUPPORTED_META; +} + +UmamiStatus umami_supported_inputs(bool const** supported, int* count) { + *supported = key_tables.inputs; + *count = UMAMI_INPUT_KEY_COUNT; + return UMAMI_SUCCESS; +} + +UmamiStatus umami_required_inputs(bool const** required, int* count) { + return umami_supported_inputs(required, count); +} + +UmamiStatus umami_supported_outputs(bool const** supported, int* count) { + *supported = key_tables.outputs; + *count = UMAMI_OUTPUT_KEY_COUNT; + return UMAMI_SUCCESS; +} + +UmamiStatus umami_initialize(UmamiHandle* handle, char const*) { + *handle = &dummy_instance; + return UMAMI_SUCCESS; +} + +UmamiStatus umami_set_parameter(UmamiHandle, char const*, double, double) { + return UMAMI_ERROR_NOT_IMPLEMENTED; +} + +UmamiStatus umami_get_parameter(UmamiHandle, char const*, double*, double*) { + return UMAMI_ERROR_NOT_IMPLEMENTED; +} + +UmamiStatus umami_matrix_element( + UmamiHandle, + std::size_t count, + std::size_t stride, + std::size_t offset, + std::size_t input_count, + UmamiInputKey const* input_keys_in, + void const* const* inputs, + std::size_t output_count, + UmamiOutputKey const* output_keys, + void* const* outputs +) { + double* output = nullptr; + gpuStream_t stream = nullptr; + for (std::size_t i = 0; i < output_count; ++i) { + if (output_keys[i] == CUSTOM_OUTPUT_KEY) { + output = static_cast(outputs[i]) + offset; + } else if (output_keys[i] == UMAMI_OUT_GPU_STREAM) { + stream = static_cast(outputs[i]); + } else { + return UMAMI_ERROR_UNSUPPORTED_OUTPUT; + } + } + if (output == nullptr || count == 0) { + return UMAMI_SUCCESS; + } + unsigned int block_count = (count + thread_count - 1) / thread_count; + +#if CUSTOM_IS_PDF + const double* x = nullptr; + const double* q = nullptr; + const int* pdg_id = nullptr; + const int* beam_index = nullptr; + for (std::size_t i = 0; i < input_count; ++i) { + switch (input_keys_in[i]) { + case UMAMI_IN_MOMENTUM_FRACTION: + x = static_cast(inputs[i]) + offset; + break; + case UMAMI_IN_SCALE: + q = static_cast(inputs[i]) + offset; + break; + case UMAMI_IN_PDG_ID: + pdg_id = static_cast(inputs[i]) + offset; + break; + case UMAMI_IN_BEAM_INDEX: + beam_index = static_cast(inputs[i]) + offset; + break; + default: + return UMAMI_ERROR_UNSUPPORTED_INPUT; + } + } + if (!x || !q || !pdg_id || !beam_index) { + return UMAMI_ERROR_MISSING_INPUT; + } + evaluate<<>>( + count, x, q, pdg_id, beam_index, output + ); +#else + const double* momenta = nullptr; + const int* particle_count = nullptr; + const int* pdg_ids = nullptr; + for (std::size_t i = 0; i < input_count; ++i) { + switch (input_keys_in[i]) { + case UMAMI_IN_MOMENTA: + momenta = static_cast(inputs[i]) + offset; + break; + case UMAMI_IN_PARTICLE_COUNT: + particle_count = static_cast(inputs[i]) + offset; + break; + case UMAMI_IN_PDG_IDS: + pdg_ids = static_cast(inputs[i]) + offset; + break; + default: + return UMAMI_ERROR_UNSUPPORTED_INPUT; + } + } + if (!momenta || !particle_count || !pdg_ids) { + return UMAMI_ERROR_MISSING_INPUT; + } + evaluate<<>>( + count, stride, momenta, particle_count, pdg_ids, output + ); +#endif + return gpuGetLastError() == gpuSuccess ? UMAMI_SUCCESS : UMAMI_ERROR; +} + +UmamiStatus umami_free(UmamiHandle) { return UMAMI_SUCCESS; } + +} // extern "C" diff --git a/madgraph/iolibs/template_files/mg7/event_stream.py b/madgraph/iolibs/template_files/mg7/event_stream.py new file mode 100644 index 0000000000..3a24ab9c97 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/event_stream.py @@ -0,0 +1,299 @@ +"""In-memory event generation from a gridpack, as a Python API. + +Only available in gridpacks made with ``[gridpack] run_mode = "fix_max_weight"``: +their channels carry a fixed maximum weight, so every event is final as soon as +it is unweighted and the events can be streamed without a survey or a combine +pass. Nothing is written to disk. + +Usage:: + + import sys + sys.path.insert(0, "/bin") + from event_stream import EventStream + + with EventStream(device="cpu", seed=42) as stream: + for event in stream: # ms.LHEEvent objects, endless + ... + for batch in stream.batches(10000, output_type="torch"): + batch.px # shape (10000, max_particle_count) + +Event weights are in units of the maximum weight: +-1, except for overweight +events. To normalize a sample to the cross section, scale its weights by +``stream.integral / sum(weights)``. The event sequence is reproducible from the seed, +independent of the thread count and of how it is split into batches. +""" + +import dataclasses +import json +import sys +from typing import Any, Iterator, Optional + +import numpy as np + +from gridpack_setup import ( + SOURCE_HASH_MESSAGE, + build_lhe_meta, + gridpack_path, + load_channels, + load_context_data, + load_lhe_completer, + load_madspace_data, + load_run_card, + make_context, + resolve_seed, + source_hash_matches, +) +import madspace as ms + +__all__ = ["EventStream", "EventBatch"] + +# events completed per call into madspace while iterating over single events +_ITER_CHUNK = 1000 + + +@dataclasses.dataclass +class EventBatch: + """A batch of LHE events as a structure of arrays: event fields have shape + (batch_size,), particle fields (batch_size, max_particle_count). Particle + slots beyond ``particle_count`` are padded with zeros, see + ``particle_mask``. numpy arrays or CPU torch tensors, depending on the + ``output_type`` of :meth:`EventStream.batches`.""" + + #: Process ID of the event (IDPRUP). + process_id: Any + #: Event weight, +-1 or larger in magnitude for overweight events (XWGTUP). + weight: Any + #: Renormalization scale in GeV (SCALUP). + scale: Any + #: QED coupling (AQEDUP). + alpha_qed: Any + #: QCD coupling (AQCDUP). + alpha_qcd: Any + #: Number of particles of the event, including resonances (NUP). + particle_count: Any + #: True for the particle slots below ``particle_count``. + particle_mask: Any + #: PDG ID of the particle (IDUP). + pdg_id: Any + #: -1 for incoming, 1 for outgoing particles, 2 for resonances (ISTUP). + status_code: Any + #: Index of the first mother, counting from 1; 0 if there is none (MOTHUP). + mother1: Any + #: Index of the second mother, counting from 1; 0 if there is none (MOTHUP). + mother2: Any + #: Color line, 0 for a color singlet (ICOLUP). + color: Any + #: Anti-color line (ICOLUP). + anti_color: Any + #: x component of the momentum in GeV. + px: Any + #: y component of the momentum in GeV. + py: Any + #: z component of the momentum in GeV. + pz: Any + #: Energy in GeV. + energy: Any + #: Mass in GeV. + mass: Any + #: Proper lifetime c*tau in mm, 0 if not tracked (VTIMUP). + lifetime: Any + #: Cosine of the spin angle, or 9 if the spin is not specified (SPINUP). + spin: Any + + def __len__(self) -> int: + return len(self.weight) + + +class EventStream: + """Endless stream of unweighted events from this gridpack. + + Constructing it starts the generation in background threads, which keep + a buffer of events for every channel. The channel of every event is picked + at random according to the current cross-section estimates, which are + refined as events are generated. + + Arguments left as None take their value from Cards/grid_run_card.toml. + + Args: + device: device to generate on, like "cpu", "cuda" or "cuda:1". + seed: run seed; -1 draws a random one (see :attr:`seed`). + cpu_batch_size: generation batch size on a CPU device. + gpu_batch_size: generation batch size on a GPU device. + thread_pool_size: number of worker threads. + cut_efficiency_threshold: fraction of a batch that has to pass the cuts. + max_cut_repetitions: how often a batch is repeated to pass the cuts. + ignore_source_hash: run even if madspace differs from the version the + gridpack was made with (can lead to errors or incorrect results). + """ + + #: 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 + run_card = load_run_card() + run_args = run_card["run"] + gen_args = run_card["generation"] + madspace_data = load_madspace_data() + if not source_hash_matches(madspace_data): + if not ignore_source_hash: + raise RuntimeError( + f"{SOURCE_HASH_MESSAGE}. Pass ignore_source_hash=True to run anyway." + ) + print(f"WARNING: {SOURCE_HASH_MESSAGE}", file=sys.stderr) + + with open(gridpack_path("data", "channels.json")) as f: + channels = json.load(f) + if not any("max_weight" in channel for channel in channels.values()): + raise RuntimeError( + "EventStream needs a gridpack made with [gridpack] run_mode = " + '"fix_max_weight"' + ) + + def pick(value, default): + return default if value is None else value + + if device is None: + devices = run_args["device"] + device = devices[0] if isinstance(devices, list) else devices + if thread_pool_size is None: + is_gpu = device.split(":")[0] in ("cuda", "hip") + thread_pool_size = run_args[ + "gpu_thread_pool_size" if is_gpu else "cpu_thread_pool_size" + ] + self._context, backend = make_context( + device, run_args["cpu_mode"], thread_pool_size, thread_pool_size + ) + + config = ms.GeneratorConfig() + config.cpu_batch_size = pick(cpu_batch_size, gen_args["cpu_batch_size"]) + config.gpu_batch_size = pick(gpu_batch_size, gen_args["gpu_batch_size"]) + config.cut_efficiency_threshold = pick( + cut_efficiency_threshold, gen_args["cut_efficiency_threshold"] + ) + config.max_cut_repetitions = pick( + max_cut_repetitions, gen_args["max_cut_repetitions"] + ) + + self.seed = resolve_seed(pick(seed, run_args.get("seed", -1))) + load_context_data([self._context], [backend], madspace_data) + channel_generators = load_channels([self._context], config) + self._stream = ms.EventStream( + self._context, channel_generators, self.seed, load_lhe_completer(), config + ) + + def __iter__(self) -> Iterator["ms.LHEEvent"]: + """Endless iterator over the events as ms.LHEEvent objects. It takes + events from the stream in chunks of 1000; those of the last chunk that + are not iterated over are dropped.""" + while True: + yield from self._stream.next_events(_ITER_CHUNK) + + def next_events(self, count: int) -> list: + """The next ``count`` events as ms.LHEEvent objects.""" + return self._stream.next_events(count) + + def next_batch(self, batch_size: int, output_type: str = "numpy") -> EventBatch: + """The next ``batch_size`` events as an :class:`EventBatch`.""" + if output_type not in ("numpy", "torch"): + raise ValueError('output_type must be "numpy" or "torch"') + arrays = self._stream.next_batch(batch_size) + counts = arrays["particle_count"] + max_count = arrays["pdg_id"].shape[1] + arrays["particle_mask"] = np.arange(max_count)[None, :] < counts[:, None] + if output_type == "torch": + import torch + arrays = {key: torch.from_numpy(value) for key, value in arrays.items()} + return EventBatch(**arrays) + + def batches( + self, + batch_size: int, + output_type: str = "numpy", + max_events: Optional[int] = None, + ) -> Iterator[EventBatch]: + """Iterator over batches of ``batch_size`` events, endless unless + ``max_events`` is given (the last batch may then be smaller).""" + remaining = max_events + while remaining is None or remaining > 0: + size = batch_size if remaining is None else min(batch_size, remaining) + yield self.next_batch(size, output_type) + if remaining is not None: + remaining -= size + + @property + def integral(self) -> float: + """Current estimate of the cross section, in pb.""" + return self._stream.status().mean + + @property + def integral_error(self) -> float: + """Standard error of :attr:`integral`.""" + return self._stream.status().error + + @property + def abs_integral(self) -> float: + """Current estimate of the integral of the absolute weights.""" + return self._stream.status().mean_abs + + @property + def n_weighted(self) -> int: + """Number of weighted events the estimates are based on.""" + return self._stream.status().count + + @property + def n_weighted_after_cuts(self) -> int: + """Number of those weighted events that passed the cuts.""" + return self._stream.status().count_after_cuts + + @property + def n_events(self) -> int: + """Number of events returned so far.""" + return self._stream.event_count() + + @property + def max_particle_count(self) -> int: + """Largest number of particles per event, including resonances.""" + return self._stream.max_particle_count() + + def status(self): + """Combined ms.GeneratorStatus over all channels.""" + return self._stream.status() + + def channel_status(self) -> list: + """ms.GeneratorStatus of every channel; ``count_unweighted`` is the + number of its events returned so far.""" + return self._stream.channel_status() + + def channel_probabilities(self) -> list: + """Current probability of every channel to be picked for an event.""" + return self._stream.channel_probabilities() + + def lhe_meta(self): + """ms.LHEMeta for writing the events to an LHE file, with the current + cross section estimate.""" + return build_lhe_meta(self._stream.status(), self.seed) + + def close(self) -> None: + """Stop the background generation.""" + if self._stream is not None: + self._stream.close() + + def __enter__(self) -> "EventStream": + return self + + def __exit__(self, *exc) -> None: + self.close() + + def __del__(self) -> None: + self.close() diff --git a/madgraph/iolibs/template_files/mg7/gridpack.py b/madgraph/iolibs/template_files/mg7/gridpack.py index e7081fa56d..aceca62457 100644 --- a/madgraph/iolibs/template_files/mg7/gridpack.py +++ b/madgraph/iolibs/template_files/mg7/gridpack.py @@ -1,37 +1,29 @@ #! /usr/bin/env python3 +import argparse +import glob import gzip +import json import os import shutil -import subprocess import sys -from pathlib import Path - -# search order for madspace package: -# 1. a precompiled install (madspace/install/madspace) -# 2. bundled source that still needs to be built (madspace/install.py) -# 3. otherwise fall back to madspace is available in the environment -_GRIDPACK_DIR = Path(os.path.realpath(__file__)).parent.parent -_LOCAL_MADSPACE_DIR = _GRIDPACK_DIR / "madspace" -_LOCAL_INSTALL_DIR = _LOCAL_MADSPACE_DIR / "install" -if (_LOCAL_INSTALL_DIR / "madspace").is_dir(): - sys.path.insert(0, str(_LOCAL_INSTALL_DIR)) -elif (_LOCAL_MADSPACE_DIR / "install.py").is_file(): - print() - print("You don't have madspace installed for this gridpack") - print("Running interactive madspace installation script") - print() - - _result = subprocess.run([sys.executable, str(_LOCAL_MADSPACE_DIR / "install.py")]) - if _result.returncode != 0: - raise RuntimeError("madspace installation failed — see output above") - sys.path.insert(0, str(_LOCAL_INSTALL_DIR)) +# bin/gridpack_setup.py, next to this script: also locates madspace +from gridpack_setup import ( + SOURCE_HASH_MESSAGE, + build_lhe_meta, + load_channels, + load_context_data, + load_lhe_completer, + load_madspace_data, + load_run_card, + load_systematics, + make_context, + resolve_seed, + source_hash_matches, +) import madspace as ms -import glob -import json -import tomllib -import argparse + def resolve_verbosity(verbosity: str) -> str: """Resolve the run_card "auto" verbosity to "pretty"/"log" depending on @@ -42,47 +34,6 @@ def resolve_verbosity(verbosity: str) -> str: return verbosity -def resolve_seed(seed: int) -> int: - """Resolve the run_card "seed": -1 draws a fresh 64-bit seed via - os.urandom, any other value is used as-is.""" - if seed == -1: - return int.from_bytes(os.urandom(8), "big") - return seed - - -def build_lhe_meta(event_generator, seed: int, systematics=None): - """The LHE header/ metadata of this gridpack run: the cards, beams - and PDF the gridpack was made with (data/lhe_meta.json), with the cross - section and seed of this run. ``systematics`` adds the block, - for the writers that do not inject it themselves (header.lhe); leave it - unset for combine_to_lhe, which does.""" - headers = [] - data = {} - meta_path = os.path.join("data", "lhe_meta.json") - if os.path.exists(meta_path): - with open(meta_path) as f: - data = json.load(f) - headers = [ms.LHEHeader(name=h["name"], content=h["content"], - escape_content=h["escape_content"]) - for h in data.pop("headers")] - headers.append(ms.LHEHeader(name="MG7Seed", content=str(seed))) - if systematics is not None and systematics.weight_ids: - headers.append(ms.LHEHeader( - name="initrwgt", content=systematics.initrwgt(), escape_content=False - )) - status = event_generator.status() - # 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.""" @@ -90,30 +41,16 @@ def write_lhe_header(run_path: str, meta) -> None: del writer # closes the file (writes the closing tag) +# directory the gridpack was called from; relative command line paths refer to it +_INVOCATION_DIR = os.getcwd() + + def main() -> None: - # load run card and metadata. Use the RunCardMG7 representation when the - # madgraph package is importable; gridpacks are meant to be portable, so - # fall back to a plain tomllib parse otherwise (the card is the same TOML). - run_card_path = os.path.join("Cards", "grid_run_card.toml") - try: - from madgraph.various.banner import RunCardMG7 - run_card = RunCardMG7(run_card_path) - except ImportError: - with open(run_card_path, "rb") as f: - run_card = tomllib.load(f) + run_card = load_run_card() run_args = run_card["run"] gen_args = run_card["generation"] - param_card_path = os.path.join("Cards", "param_card.dat") - with open(os.path.join("data", "data.json")) as f: - madspace_data = json.load(f) - if madspace_data["source_hash"] != ms.SOURCE_HASH: - print() - print( - "\033[1m\033[31mWARNING\033[39m: The madspace version is not identical " - "to the one used to generate the gridpack. This can lead to errors or " - "incorrect results\033[0m" - ) - print() + grid_args = run_card.get("gridpack", {}) + madspace_data = load_madspace_data() # parse command line arguments parser = argparse.ArgumentParser() @@ -124,6 +61,22 @@ def main() -> None: "seed each run instead of fixing one here (still recorded in the " "run's info.json)" ) + parser.add_argument( + "--output_dir", type=str, default=None, + help="directory for the final output files instead of a new Events/ " + "folder (default: the run card value; relative paths refer to the " + "current directory)" + ) + parser.add_argument( + "--temp_output_dir", type=str, default=None, + help="directory for the temporary npy files (default: the run card " + "value, else the output directory)" + ) + parser.add_argument( + "--ignore_source_hash", action="store_true", + help="run even if the madspace version differs from the one used to " + "generate the gridpack (can lead to errors or incorrect results)" + ) parser.add_argument("--device", type=str, nargs="*") parser.add_argument( "--cpu_thread_pool_size", type=int, default=run_args["cpu_thread_pool_size"] @@ -149,51 +102,67 @@ def main() -> None: parser.add_argument("--cpu_batch_size", type=int, default=gen_args["cpu_batch_size"]) parser.add_argument("--gpu_batch_size", type=int, default=gen_args["gpu_batch_size"]) args = parser.parse_args() + + if 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" + "Use --ignore_source_hash to run anyway.\033[0m" + ) + print() + print(f"\033[1m\033[31mWARNING\033[39m: {message}\033[0m") + print() seed = resolve_seed(args.seed) - # initialize event directory - run_name = args.run_name - os.makedirs("Events", exist_ok=True) - run_dir_prefix = os.path.join("Events", f"{run_name}_") - existing_run_dirs = glob.glob(f"{run_dir_prefix}*") - run_index = 1 - for run_dir in existing_run_dirs: - run_index_str = run_dir[len(run_dir_prefix):] - if run_index_str.isnumeric(): - run_index = max(run_index, int(run_index_str) + 1) - while True: - try: - run_path = f"{run_dir_prefix}{run_index:02d}" - os.mkdir(run_path) - break - except FileExistsError: - run_index += 1 + # initialize output directories; command line paths are relative to the + # invocation directory, run card paths to the gridpack + def resolve_dir(cli_value, card_value): + if cli_value: + return os.path.abspath( + os.path.join(_INVOCATION_DIR, os.path.expanduser(cli_value))) + return os.path.abspath(os.path.expanduser(card_value)) if card_value else None + + output_dir = resolve_dir(args.output_dir, grid_args.get("output_dir")) + temp_dir = resolve_dir(args.temp_output_dir, grid_args.get("temp_output_dir")) + if output_dir is not None: + run_path = output_dir + os.makedirs(run_path, exist_ok=True) + else: + run_name = args.run_name + os.makedirs("Events", exist_ok=True) + run_dir_prefix = os.path.join("Events", f"{run_name}_") + existing_run_dirs = glob.glob(f"{run_dir_prefix}*") + run_index = 1 + for run_dir in existing_run_dirs: + run_index_str = run_dir[len(run_dir_prefix):] + if run_index_str.isnumeric(): + run_index = max(run_index, int(run_index_str) + 1) + while True: + try: + run_path = f"{run_dir_prefix}{run_index:02d}" + os.mkdir(run_path) + break + except FileExistsError: + run_index += 1 + if temp_dir is None: + temp_dir = run_path + else: + os.makedirs(temp_dir, exist_ok=True) # initialize context device_names = args.device if args.device else run_args["device"] - 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() @@ -207,29 +176,9 @@ def main() -> None: 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", {}) - global_dir = os.path.join("data", "globals") - for context, backend in zip(contexts, backends): - context.load_globals(global_dir) - 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 - channel_generators = [ - ms.ChannelEventGenerator.load( - os.path.join("data", "channels", file), - contexts, - event_file=os.path.join(run_path, f"events.{name}.npy"), - weight_file=os.path.join(run_path, f"weights.{name}.npy"), - config=config, - ) - for name, file in madspace_data["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, @@ -239,8 +188,8 @@ def main() -> None: ) # 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() @@ -250,24 +199,24 @@ def main() -> None: os.path.join(run_path, "events.npy"), systematics ) write_lhe_header(run_path, - build_lhe_meta(event_generator, seed, systematics)) + build_lhe_meta(event_generator.status(), seed, systematics)) # what npy_to_lhe needs to complete these events into LHE later shutil.copy(os.path.join("data", "lhe.json"), os.path.join(run_path, "lhe_completer.json")) elif output_format == "lhe_npy": - lhe_completer = ms.LHECompleter.load(os.path.join("data", "lhe.json")) + lhe_completer = load_lhe_completer() event_generator.combine_to_lhe_npy( os.path.join(run_path, "events.npy"), lhe_completer, systematics ) write_lhe_header(run_path, - build_lhe_meta(event_generator, seed, systematics)) + build_lhe_meta(event_generator.status(), seed, systematics)) elif output_format == "lhe": - lhe_completer = ms.LHECompleter.load(os.path.join("data", "lhe.json")) + lhe_completer = load_lhe_completer() lhe_path = os.path.join(run_path, "events.lhe") # the cross section is known once generate() has converged, which is # before combine_to_lhe writes the block event_generator.combine_to_lhe( - lhe_path, lhe_completer, build_lhe_meta(event_generator, seed), + lhe_path, lhe_completer, build_lhe_meta(event_generator.status(), seed), systematics ) # Ship the LHE compressed, as the launcher that produced this gridpack @@ -290,118 +239,6 @@ def main() -> None: json.dump(data, f, indent=1) -def _pdf_search_paths(stored_path): - """Directories in which the LHAPDF grids may live: the LHAPDF_DATA_PATH - entries, the paths known to the lhapdf module, then the directory the - gridpack was created with.""" - paths = [] - env = os.environ.get("LHAPDF_DATA_PATH") - if env: - paths += env.split(os.pathsep) - try: - import lhapdf - paths += list(lhapdf.paths()) - except Exception: - pass - if stored_path: - paths.append(stored_path) - return paths - - -def _locate_pdf_file(stored_file): - """Find a PDF grid/info file: the stored absolute path, else the same - / relative to one of the PDF search paths.""" - if os.path.exists(stored_file): - return stored_file - set_dir, name = os.path.split(stored_file) - set_root, set_name = os.path.split(set_dir) - for base in _pdf_search_paths(set_root): - candidate = os.path.join(base, set_name, name) - if os.path.exists(candidate): - return candidate - raise FileNotFoundError( - "PDF file %s not found (set LHAPDF_DATA_PATH to the LHAPDF data directory)" - % stored_file) - - -def load_systematics(run_card, backends=(), param_card_path=None, - me_parameters=None): - """Rebuild the ms.SystematicsCalculator saved with the gridpack - (data/systematics.json) when [systematics] enable is set; None otherwise. - The matrix elements of the mixed-order subprocesses are reloaded into a CPU - context (when a CPU library is available) for the mu_R variations.""" - path = os.path.join("data", "systematics.json") - try: - enabled = bool(run_card["systematics"]["enable"]) - except Exception: - enabled = False - if not enabled or not os.path.exists(path): - return None - with open(path) as f: - data = json.load(f) - config = ms.SystematicsConfig.from_json(json.dumps(data["config"])) - for spec in config.pdf_members: - spec.grid_file = _locate_pdf_file(spec.grid_file) - spec.info_file = _locate_pdf_file(spec.info_file) - subproc_args = [ - ms.SubprocessSystArgs.from_json(json.dumps(a)) for a in data["subproc_args"] - ] - nominal_pdf = None - if config.has_pdf: - nominal_pdf = ms.PdfGrid(_locate_pdf_file(data["nominal_grid_file"])) - # The mu_R variations always need an alpha_s grid, so a gridpack that does - # not name one cannot reweight anything -- drop the systematics rather than - # die here (gridpacks written before the launcher recorded the file for a - # run without parton luminosity). - if not data.get("nominal_info_file"): - print("WARNING systematics: the gridpack records no alpha_s .info file, " - "the scale/PDF weights are dropped") - return None - nominal_alpha_s = ms.AlphaSGrid(_locate_pdf_file(data["nominal_info_file"])) - # PDFs, alpha_s and matrix elements are evaluated on this CPU context - context = ms.Context(device=ms.cpu_device(), thread_count=1) - matrix_elements, flavor_remap = [], [] - need = [i for i, a in enumerate(subproc_args) if a.qcd_power < 0] - # the CPU backend of this run (cpu_mode, resolved when the gridpack was - # saved), else the one the launcher used - backend = data.get("me_backend") - if need and data.get("me_paths"): - cpu = [b for b in backends if not str(b).startswith(("cuda", "hip")) and b != "auto"] - if cpu: - backend = cpu[0] - if need and backend and data.get("me_paths"): - flavor_remap = data.get("flavor_remap", []) - for i, me_path in enumerate(data["me_paths"]): - if i not in need: - matrix_elements.append(None) - continue - lib = me_path.format(device=backend) - if not os.path.exists(lib): - print("WARNING systematics: %s not found, mu_R variations of the " - "mixed-order subprocesses are dropped" % lib) - matrix_elements, flavor_remap = [], [] - break - api = context.load_matrix_element(lib, param_card_path, - me_parameters or {}) - matrix_elements.append(ms.MatrixElement( - api, - [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, - ms.MatrixElement.flavor_in], - [ms.MatrixElement.matrix_element_out], - False, - # the frame the events were generated with, not whatever the - # gridpack's run card says now: it has to match the matrix - # element inside the saved channel generators - data.get("me_frame", []), - data.get("incoming_count", 2), - )) - systematics = ms.SystematicsCalculator( - config, subproc_args, nominal_pdf, nominal_alpha_s, - context, matrix_elements, flavor_remap) - for warning in systematics.warnings: - print("WARNING systematics: %s" % warning) - return systematics - if __name__ == '__main__': os.chdir(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) try: diff --git a/madgraph/iolibs/template_files/mg7/gridpack_setup.py b/madgraph/iolibs/template_files/mg7/gridpack_setup.py new file mode 100644 index 0000000000..efdfb52b28 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/gridpack_setup.py @@ -0,0 +1,301 @@ +"""Gridpack setup shared by bin/generate_events and bin/event_stream.py: locating +madspace, reading the cards and data files, and building the contexts, channels, +LHE completer and systematics of the gridpack. Paths are resolved relative to +the gridpack directory, so nothing depends on the current directory.""" + +import json +import os +import subprocess +import sys +from pathlib import Path + +GRIDPACK_DIR = Path(os.path.realpath(__file__)).parent.parent + +# search order for madspace package: +# 1. a precompiled install (madspace/install/madspace) +# 2. bundled source that still needs to be built (madspace/install.py) +# 3. otherwise fall back to madspace is available in the environment +_LOCAL_MADSPACE_DIR = GRIDPACK_DIR / "madspace" +_LOCAL_INSTALL_DIR = _LOCAL_MADSPACE_DIR / "install" +if (_LOCAL_INSTALL_DIR / "madspace").is_dir(): + sys.path.insert(0, str(_LOCAL_INSTALL_DIR)) +elif (_LOCAL_MADSPACE_DIR / "install.py").is_file(): + print() + print("You don't have madspace installed for this gridpack") + print("Running interactive madspace installation script") + print() + + _result = subprocess.run([sys.executable, str(_LOCAL_MADSPACE_DIR / "install.py")]) + if _result.returncode != 0: + raise RuntimeError("madspace installation failed — see output above") + sys.path.insert(0, str(_LOCAL_INSTALL_DIR)) + +import madspace as ms + + +def gridpack_path(*parts) -> str: + """A path inside the gridpack directory.""" + return str(GRIDPACK_DIR.joinpath(*parts)) + + +def load_run_card(): + """The gridpack run card (Cards/grid_run_card.toml). Use the RunCardMG7 + representation when the madgraph package is importable; gridpacks are meant + to be portable, so fall back to a plain tomllib parse otherwise (the card is + the same TOML).""" + run_card_path = gridpack_path("Cards", "grid_run_card.toml") + try: + from madgraph.various.banner import RunCardMG7 + return RunCardMG7(run_card_path) + except ImportError: + import tomllib + with open(run_card_path, "rb") as f: + return tomllib.load(f) + + +def load_madspace_data() -> dict: + """data/data.json: the matrix elements, their run-time parameters and the + madspace source hash the gridpack was made with.""" + with open(gridpack_path("data", "data.json")) as f: + return json.load(f) + + +def source_hash_matches(madspace_data: dict) -> bool: + return madspace_data["source_hash"] == ms.SOURCE_HASH + + +SOURCE_HASH_MESSAGE = ( + "The madspace version is not identical to the one used to generate " + "the gridpack. This can lead to errors or incorrect results" +) + + +def resolve_seed(seed: int) -> int: + """Resolve the run_card "seed": -1 draws a fresh 64-bit seed via + os.urandom, any other value is used as-is.""" + if seed == -1: + return int.from_bytes(os.urandom(8), "big") + return seed + + +def make_context(device_name: str, cpu_mode: str, cpu_thread_pool_size: int, + gpu_thread_pool_size: int): + """The context for a device name like "cpu", "cuda:1" or "hip", with the + backend its matrix elements are built for.""" + if ":" in device_name: + device_type, device_index_str = device_name.split(":") + device_index = int(device_index_str) + else: + device_type = device_name + device_index = 0 + # cpu_mode names the SIMD width of the CPU code, so it applies to the + # 'cpu' devices only: cuda/hip build the backend named after the device. + backend = cpu_mode if device_type == "cpu" else device_type + if device_type == "cuda": + device = ms.cuda_device(device_index) + pool_size = gpu_thread_pool_size + elif device_type == "hip": + device = ms.hip_device(device_index) + pool_size = gpu_thread_pool_size + else: + device = ms.cpu_device() + pool_size = cpu_thread_pool_size + return ms.Context(device=device, thread_count=pool_size), backend + + +def load_context_data(contexts, backends, madspace_data: dict) -> None: + """Load the globals and the matrix elements of the gridpack into every + context.""" + # run-time matrix-element parameters (bwcutoff) of the run that made the + # gridpack; gridpacks written before they were recorded used the default + me_parameters = madspace_data.get("me_parameters", {}) + for context 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", []): + # older gridpacks store only the path of a .so library + kind, path = entry if isinstance(entry, list) else ("", entry) + load_custom_function(ms, contexts, device_types, kind, _gridpack_relative(path)) + for context, backend in zip(contexts, backends): + for me_path in madspace_data["matrix_elements"]: + context.load_matrix_element( + _gridpack_relative(me_path.format(device=backend)), + gridpack_path("Cards", "param_card.dat"), + me_parameters, + ) + + +def load_channels(contexts, config, file_dir=None): + """The channel generators stored in data/channels.json. With `file_dir`, + their temporary event files are written there; without, events are only + kept in memory.""" + with open(gridpack_path("data", "channels.json")) as f: + channels = json.load(f) + return [ + ms.ChannelEventGenerator.load_json( + json.dumps(channel), + contexts, + event_file=os.path.join(file_dir, f"events.{name}.npy") if file_dir else "", + weight_file=os.path.join(file_dir, f"weights.{name}.npy") if file_dir else "", + config=config, + ) + for name, channel in channels.items() + ] + + +def load_lhe_completer(): + return ms.LHECompleter.load(gridpack_path("data", "lhe.json")) + + +def build_lhe_meta(status, seed: int, systematics=None): + """The LHE header/ metadata of this gridpack run: the cards, beams + and PDF the gridpack was made with (data/lhe_meta.json), with the cross + section (from the generator `status`) and seed of this run. + ``systematics`` adds the block, for the writers that do not + inject it themselves (header.lhe); leave it unset for combine_to_lhe, + which does.""" + headers = [] + data = {} + meta_path = gridpack_path("data", "lhe_meta.json") + if os.path.exists(meta_path): + with open(meta_path) as f: + data = json.load(f) + headers = [ms.LHEHeader(name=h["name"], content=h["content"], + escape_content=h["escape_content"]) + for h in data.pop("headers")] + headers.append(ms.LHEHeader(name="MG7Seed", content=str(seed))) + if systematics is not None and systematics.weight_ids: + headers.append(ms.LHEHeader( + name="initrwgt", content=systematics.initrwgt(), escape_content=False + )) + # 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 = None + if config.has_pdf: + nominal_pdf = ms.PdfGrid(_locate_pdf_file(data["nominal_grid_file"])) + # The mu_R variations always need an alpha_s grid, so a gridpack that does + # not name one cannot reweight anything -- drop the systematics rather than + # die here (gridpacks written before the launcher recorded the file for a + # run without parton luminosity). + if not data.get("nominal_info_file"): + print("WARNING systematics: the gridpack records no alpha_s .info file, " + "the scale/PDF weights are dropped") + return None + nominal_alpha_s = ms.AlphaSGrid(_locate_pdf_file(data["nominal_info_file"])) + # PDFs, alpha_s and matrix elements are evaluated on this CPU context + context = ms.Context(device=ms.cpu_device(), thread_count=1) + matrix_elements, flavor_remap = [], [] + need = [i for i, a in enumerate(subproc_args) if a.qcd_power < 0] + # the CPU backend of this run (cpu_mode, resolved when the gridpack was + # saved), else the one the launcher used + backend = data.get("me_backend") + if need and data.get("me_paths"): + cpu = [b for b in backends if not str(b).startswith(("cuda", "hip")) and b != "auto"] + if cpu: + backend = cpu[0] + if need and backend and data.get("me_paths"): + flavor_remap = data.get("flavor_remap", []) + for i, me_path in enumerate(data["me_paths"]): + if i not in need: + matrix_elements.append(None) + continue + lib = _gridpack_relative(me_path.format(device=backend)) + if not os.path.exists(lib): + print("WARNING systematics: %s not found, mu_R variations of the " + "mixed-order subprocesses are dropped" % lib) + matrix_elements, flavor_remap = [], [] + break + api = context.load_matrix_element( + lib, gridpack_path("Cards", "param_card.dat"), me_parameters or {} + ) + matrix_elements.append(ms.MatrixElement( + api, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in], + [ms.MatrixElement.matrix_element_out], + False, + # the frame the events were generated with, not whatever the + # gridpack's run card says now: it has to match the matrix + # element inside the saved channel generators + data.get("me_frame", []), + data.get("incoming_count", 2), + )) + systematics = ms.SystematicsCalculator( + config, subproc_args, nominal_pdf, nominal_alpha_s, + context, matrix_elements, flavor_remap) + for warning in systematics.warnings: + print("WARNING systematics: %s" % warning) + return systematics diff --git a/madgraph/iolibs/template_files/mg7/init_custom_func.py b/madgraph/iolibs/template_files/mg7/init_custom_func.py new file mode 100644 index 0000000000..9a40959056 --- /dev/null +++ b/madgraph/iolibs/template_files/mg7/init_custom_func.py @@ -0,0 +1,250 @@ +"""Create a skeleton for a custom cut, scale, PDF or bias function. + +The function is written for a single phase-space point in C, C++, Fortran, +CUDA or HIP. A generated UMAMI wrapper calls it on batches of events, so that +madspace can load it like a matrix element. The files go into +``custom_/`` in the process directory, and ``make`` there builds +``lib/custom_.so``. The run card option +``[custom_functions] `` takes the path of that library. + +Usage:: + + bin/init_custom_func [-f {cut,scale,pdf,bias}] [-l {c,cpp,fortran,cuda,hip}] + [-y] [--force] [--build | --no-build] + +Options that are not given are asked interactively, unless ``-y`` is passed. +""" + +import argparse +import logging +import os +import shutil +import subprocess +import sys + +logger = logging.getLogger("madevent") + +TEMPLATE_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "custom_functions") +UMAMI_HEADER = os.path.join( + os.path.dirname(os.path.dirname(os.path.realpath(__file__))), "madmatrix", "umami.h" +) + +FUNCTIONS = { + "cut": ("UMAMI_OUT_CUT", "a cut, 1 keeps the event and 0 rejects it"), + "scale": ("UMAMI_OUT_SCALE", "the renormalization and factorization scale"), + "pdf": ("UMAMI_OUT_PDF", "a parton density x f(x, Q)"), + "bias": ("UMAMI_OUT_BIAS", "an event bias for the unweighting"), +} +# language: (body template extension, description) +LANGUAGES = { + "c": ("c", "C"), + "cpp": ("cpp", "C++"), + "fortran": ("f90", "Fortran"), + "cuda": ("cu", "CUDA, for runs on NVIDIA GPUs"), + "hip": ("cu", "HIP, for runs on AMD GPUs"), + "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_key": FUNCTIONS[function][0], + "is_pdf": int(function == "pdf"), + "body_file": body_file, + "language": language, + } + os.makedirs(out_dir, exist_ok=True) + shutil.copy(os.path.join(TEMPLATE_DIR, f"{function}.{extension}"), os.path.join(out_dir, body_file)) + shutil.copy(UMAMI_HEADER, os.path.join(out_dir, "umami.h")) + wrapper = "wrapper_gpu.cu" if language in ("cuda", "hip") else "wrapper_cpu.c" + for template in (wrapper, "Makefile"): + with open(os.path.join(TEMPLATE_DIR, template)) as f: + text = f.read() % values + with open(os.path.join(out_dir, template), "w") as f: + f.write(text) + return body_file + + +def build(out_dir): + """Run make in ``out_dir``, logging to compile.log. Returns True on success.""" + log_path = os.path.join(out_dir, "compile.log") + with open(log_path, "w") as log: + result = subprocess.run(["make"], cwd=out_dir, stdout=log, stderr=subprocess.STDOUT) + if result.returncode != 0: + with open(log_path) as log: + print(log.read()) + print(f"Building failed, see {log_path}") + return False + return True + + +def set_run_card(me_dir, function, lib_path): + from madgraph.various.banner import RunCardMG7 + + card_path = os.path.join(me_dir, "Cards", "run_card.toml") + run_card = RunCardMG7(card_path) + run_card.set(f"custom_functions.{function}", lib_path) + run_card.write(card_path) + + +def main(me_dir=None, argv=None): + global _NONINTERACTIVE + args = build_parser().parse_args(argv) + if args.yes: + _NONINTERACTIVE = True + me_dir = me_dir or os.getcwd() + + function = args.function or ask_choice( + "Which custom function do you want to create?", + {name: desc for name, (_, desc) in FUNCTIONS.items()}, + "cut", + ) + language = args.language or ask_choice( + "Which language do you want to write it in?", + {name: desc for name, (_, desc) in LANGUAGES.items()}, + "cpp", + ) + + out_dir = os.path.join(me_dir, f"custom_{function}") + if os.path.exists(out_dir): + if not (args.force or (not _NONINTERACTIVE and ask_yes_no( + f"{out_dir} already exists. Overwrite it?", False))): + print(f"{out_dir} already exists, use --force to overwrite it.") + return 1 + shutil.rmtree(out_dir) + + body_file = write_files(out_dir, function, language) + lib_path = 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/iolibs/template_files/mg7/launch.py b/madgraph/iolibs/template_files/mg7/launch.py index f19ba110ef..b871e59039 100644 --- a/madgraph/iolibs/template_files/mg7/launch.py +++ b/madgraph/iolibs/template_files/mg7/launch.py @@ -38,6 +38,10 @@ from madgraph.iolibs.template_files.mg7 import plots from madgraph.iolibs.template_files.mg7 import systematics_summary from madgraph.iolibs.template_files.mg7 import npy_to_lhe +from madgraph.iolibs.template_files.mg7.custom_function_loader import ( + load_custom_function, + load_surrogate, +) from madgraph.various.banner import RunCardMG7 from madgraph.various import misc from madgraph.interface.extended_cmd import Cmd @@ -281,10 +285,14 @@ class HistItem: class MadgraphProcess: def __init__(self): + self.madnis_trained = False + self.lhe_completer = None self.load_cards() self.init_backend() self.init_event_dir() self.init_context() + self.init_custom_functions() + self.init_surrogate() self.init_cuts() self.init_histograms() self.init_generator_config() @@ -460,22 +468,34 @@ def init_backend(self) -> None: ms.set_simd_vector_size(self.run_card["run"]["simd_vector_size"]) def init_event_dir(self) -> None: - run_name = self.run_card["run"]["run_name"] - os.makedirs("Events", exist_ok=True) - run_dir_prefix = os.path.join("Events", f"{run_name}_") - existing_run_dirs = glob.glob(f"{run_dir_prefix}*") - run_index = 1 - for run_dir in existing_run_dirs: - run_index_str = run_dir[len(run_dir_prefix):] - if run_index_str.isnumeric(): - run_index = max(run_index, int(run_index_str) + 1) - while True: - try: - self.run_path = f"{run_dir_prefix}{run_index:02d}" - os.mkdir(self.run_path) - break - except FileExistsError: - run_index += 1 + output_dir = self.run_card["gridpack"]["output_dir"] + temp_dir = self.run_card["gridpack"]["temp_output_dir"] + if output_dir: + self.run_path = os.path.abspath(os.path.expanduser(output_dir)) + os.makedirs(self.run_path, exist_ok=True) + else: + run_name = self.run_card["run"]["run_name"] + os.makedirs("Events", exist_ok=True) + run_dir_prefix = os.path.join("Events", f"{run_name}_") + existing_run_dirs = glob.glob(f"{run_dir_prefix}*") + run_index = 1 + for run_dir in existing_run_dirs: + run_index_str = run_dir[len(run_dir_prefix):] + if run_index_str.isnumeric(): + run_index = max(run_index, int(run_index_str) + 1) + while True: + try: + self.run_path = f"{run_dir_prefix}{run_index:02d}" + os.mkdir(self.run_path) + break + except FileExistsError: + run_index += 1 + # temporary npy files of the channel generators + if temp_dir: + self.temp_path = os.path.abspath(os.path.expanduser(temp_dir)) + os.makedirs(self.temp_path, exist_ok=True) + else: + self.temp_path = self.run_path # Absolute on purpose. StatusFile is a C++ object that finishes its # write from its destructor (rename info.json.tmp -> info.json), and # that destructor runs whenever Python gets round to collecting the @@ -514,6 +534,104 @@ def init_context(self) -> None: self.pool_sizes.append(pool_size) self.contexts.append(ms.Context(device=device, thread_count=pool_size)) + CUSTOM_FUNCTION_KINDS = ("cut", "scale", "pdf", "bias") + + def init_custom_functions(self) -> None: + """Load the functions of [custom_functions], building the libraries + first. They are loaded before the matrix elements, so that they have the + same index in every context and in a gridpack made from this run.""" + self.custom_functions = {} + self.custom_function_paths = [] + for kind in self.CUSTOM_FUNCTION_KINDS: + path = self.run_card["custom_functions"][kind] + if not path: + continue + build_dir = f"custom_{kind}" + if path.endswith(".so") and os.path.isfile(os.path.join(build_dir, "Makefile")): + self.build_custom_function(build_dir) + if not os.path.isfile(path): + raise RuntimeError( + f"[custom_functions] {kind} = {path!r}: file not found. " + "Create a .so library or a .py file with bin/init_custom_func." + ) + apis = load_custom_function(ms, self.contexts, self.device_types, kind, path) + if len({api.index() for api in apis}) != 1: + raise RuntimeError(f"custom {kind} function loaded at different indices") + self.custom_functions[kind] = apis[0] + self.custom_function_paths.append([kind, path]) + logger.info("Using the custom %s function %s", kind, path) + + def init_surrogate(self) -> None: + """Check the [surrogate] settings. The surrogates are loaded per + subprocess, by MadgraphSubprocess.""" + settings = self.run_card["surrogate"] + self.surrogate = None + if not settings["function"]: + return + path = settings["function"] + if path.endswith(".py") and not os.path.isfile(path): + raise RuntimeError(f"[surrogate] function = {path!r}: file not found") + if self.run_card["phasespace"]["sde_strategy"] != "denominators": + raise RuntimeError( + "[surrogate] needs [phasespace] sde_strategy = \"denominators\", " + "as the surrogate does not return channel weights" + ) + if self.run_card["run"]["dummy_matrix_element"]: + raise RuntimeError("[surrogate] cannot be used with dummy_matrix_element") + if self.run_card["gridpack"]["save_gridpack"]: + raise RuntimeError("[surrogate] does not support gridpacks yet") + self.surrogate = dict(settings) + if not self.surrogate["param_card"]: + self.surrogate["param_card"] = self.param_card_path + logger.info( + "Using the surrogate %s, %s with the full matrix element", + path, + "unweighting" if settings["mode"] == "unweight" else "reweighting", + ) + + def weighting_groups(self, subprocess_channels: list[list[int]]) -> list: + """One ms.WeightingGroup per subprocess with a [surrogate]: the second + pass with the full matrix element over its channels, given by their + indices in the event generator.""" + if self.surrogate is None: + return [] + gen_args = self.run_card["generation"] + return [ + ms.WeightingGroup( + channel_indices, + [ms.WeightingStep( + subproc.step_matrix_elements, + mode=self.surrogate["mode"], + max_overweight_truncation=gen_args["max_overweight_truncation"], + freeze_max_weight_after=gen_args["freeze_max_weight_after"], + )], + ) + for subproc, channel_indices in zip(self.subprocesses, subprocess_channels) + if channel_indices + ] + + @staticmethod + def build_custom_function(build_dir: str) -> None: + """Run make in a custom_ directory, logging to compile.log.""" + log_path = os.path.join(build_dir, "compile.log") + with open(log_path, "w") as log: + returncode = subprocess.call( + ["make"], cwd=build_dir, stdout=log, stderr=subprocess.STDOUT + ) + if returncode != 0: + raise RuntimeError( + f"Compilation of {build_dir} failed with exit status {returncode}; " + f"see {log_path} for details" + ) + + def custom_function(self, kind: str, particle_count: int, pid_options: list): + """The ms.CustomFunction of [custom_functions] , or None.""" + api = self.custom_functions.get(kind) + if api is None: + return None + function_type = getattr(ms.CustomFunction, kind) + return ms.CustomFunction(api, function_type, particle_count, pid_options) + def parse_observable(self, name: str, order_observable: str) -> dict: parts = name.split("-") sum_momenta = False @@ -692,6 +810,10 @@ def init_beam(self) -> None: fact_scale2=beam_args["fact_scale2"], scale_factor=float(beam_args.get("scale_factor", 1.)), ) + if self.is_decay and self.custom_functions.pop("scale", None) is not None: + logger.warning("custom scale ignored: a decay uses the decaying mass") + if self.is_decay and self.custom_functions.pop("pdf", None) is not None: + logger.warning("custom pdf ignored: a decay has no parton densities") if self.is_decay: # One scale is available for a decay -- the decaying mass -- so use # it, fixed, whatever the card asks for. @@ -720,10 +842,17 @@ def init_beam(self) -> None: self.pdf_dir = self.lhapdf.find_set(pdf_set) if self.pdf_dir is None: raise RuntimeError(self._no_pdf_message(pdf_set)) - self.pdf_grid = ms.PdfGrid(os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}_0000.dat")) self.alphas_grid = ms.AlphaSGrid(os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}.info")) + if "pdf" in self.custom_functions: + # alpha_s still comes from the PDF set of the run card + self.pdf_grid = ms.CustomPdf(self.custom_functions["pdf"]) + else: + self.pdf_grid = ms.PdfGrid( + os.path.join(self.pdf_dir, pdf_set, f"{pdf_set}_0000.dat") + ) for context in self.contexts: - self.pdf_grid.initialize_globals(context) + if not isinstance(self.pdf_grid, ms.CustomPdf): + self.pdf_grid.initialize_globals(context) self.alphas_grid.initialize_globals(context) self.running_coupling = ms.RunningCoupling(self.alphas_grid) @@ -734,14 +863,21 @@ def systematics_enabled(self) -> bool: """The native scale/PDF weights are requested ([systematics] enable, or the legacy [generation] systematics alias).""" try: - if self.run_card["systematics"]["enable"]: - return True - except Exception: - pass - try: - return bool(self.run_card["generation"]["systematics"]) + enabled = bool(self.run_card["systematics"]["enable"]) except Exception: + enabled = False + if not enabled: + try: + enabled = bool(self.run_card["generation"]["systematics"]) + except Exception: + enabled = False + if enabled and "pdf" in getattr(self, "custom_functions", {}): + # the variations re-evaluate the PDF set, which a custom PDF replaces + if not getattr(self, "_warned_custom_pdf_systematics", False): + logger.warning("systematics are disabled with a custom pdf") + self._warned_custom_pdf_systematics = True return False + return enabled @staticmethod def pdf_set_info(info_path: str) -> dict: @@ -1157,6 +1293,14 @@ def init_generator_config(self) -> None: cfg.vegas_damping = vegas_args["damping"] cfg.max_overweight_truncation = gen_args["max_overweight_truncation"] cfg.freeze_max_weight_after = gen_args["freeze_max_weight_after"] + if self.surrogate is not None and self.surrogate["mode"] == "unweight": + # [generation] applies to the second pass, which gives the events + cfg.max_overweight_truncation = ( + self.surrogate["first_pass_max_overweight_truncation"] + ) + cfg.freeze_max_weight_after = ( + self.surrogate["first_pass_freeze_max_weight_after"] + ) cfg.start_batch_size = vegas_args["start_batch_size"] cfg.max_batch_size = vegas_args["max_batch_size"] cfg.survey_min_iters = gen_args["survey_min_iters"] @@ -1287,6 +1431,7 @@ def make_subprocesses( def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenerator: channel_generators = [] + subprocess_channels = [[] for _ in self.subprocesses] for i, (subproc, phasespace) in enumerate(zip(self.subprocesses, phasespaces)): for integrand, channel in zip( subproc.build_integrands(phasespace), phasespace.channels @@ -1295,13 +1440,14 @@ def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenera channel.event_generator = ms.ChannelEventGenerator( contexts=self.contexts, integrand=integrand, - event_file=os.path.join(self.run_path, f"events.{i}.{channel.name}.npy"), - weight_file=os.path.join(self.run_path, f"weights.{i}.{channel.name}.npy"), + event_file=os.path.join(self.temp_path, f"events.{i}.{channel.name}.npy"), + weight_file=os.path.join(self.temp_path, f"weights.{i}.{channel.name}.npy"), config=self.event_generator_config, subprocess_index=i, name=f"{i}.{channel.name}", histograms=subproc.histograms ) + subprocess_channels[i].append(len(channel_generators)) channel_generators.append(channel.event_generator) event_generator = ms.EventGenerator( @@ -1310,6 +1456,7 @@ def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenera status_file=self.status_file, config=self.event_generator_config, seed=self.run_seed, + groups=self.weighting_groups(subprocess_channels), ) unused_globals = ( set(self.contexts[0].global_names()) - event_generator.used_globals() @@ -1506,6 +1653,7 @@ def train_madnis(self) -> None: for context in self.contexts[1:]: context.copy_globals_from(self.contexts[0]) self.event_generator = self.build_event_generator(madnis_phasespaces) + self.madnis_trained = True def update_madnis_status_single( self, batch: int, batch_target: int, loss: float, lr: float, channel_count: int @@ -1589,7 +1737,30 @@ def update_madnis_status_multi( self.madnis_upper_box.print_update() self.madnis_lower_box.print_update() + @property + def gridpack_run_mode(self) -> str: + """How this run ends: a gridpack run may skip the events (see the + [gridpack] run_mode run card option); any other run is a regular one.""" + if not self.run_card["gridpack"]["save_gridpack"]: + return "regular" + return self.run_card["gridpack"]["run_mode"] + + def finish_without_events(self) -> None: + """The end of a gridpack run that writes no events: only what the + gridpack needs is prepared, then it is stored.""" + if self.gridpack_run_mode == "minimal": + # a trained MadNIS needs nothing more, a VEGAS grid has to converge + if not self.madnis_trained: + self.event_generator.optimize() + else: + self.event_generator.fix_max_weights() + self.build_systematics() + self.save_gridpack() + def generate_events(self) -> None: + if self.gridpack_run_mode != "regular": + self.finish_without_events() + return start_time = get_start_time() self.event_generator.generate() output_format = self.run_card["run"]["output_format"] @@ -1897,24 +2068,30 @@ def save_gridpack(self) -> None: os.mkdir(gridpack_path) os.mkdir(data_path) os.mkdir(events_path) - self.contexts[0].save_globals(os.path.join(data_path, "globals")) + self.contexts[0].save_globals(os.path.join(data_path, "globals.tar")) - channel_path = os.path.join(data_path, "channels") - os.mkdir(channel_path) - channel_files = {} - for channel in self.event_generator.channels(): - name = channel.status().name - file = f"channel{name}.json" - channel_files[name] = file - channel.save(os.path.join(channel_path, file)) + # one json object per channel, joined by hand so that the channel + # json is not parsed and dumped again + channel_entries = [ + f"{json.dumps(channel.status().name)}:{channel.to_json(self.gridpack_run_mode == "fix_max_weight")}" + for channel in self.event_generator.channels() + ] + with open(os.path.join(data_path, "channels.json"), "w") as f: + f.write("{" + ",".join(channel_entries) + "}") lib_path = os.path.join(gridpack_path, "lib") if self.run_card["gridpack"]["include_source"]: os.mkdir(lib_path) shutil.copytree("src", os.path.join(gridpack_path, "src")) shutil.copytree("SubProcesses", os.path.join(gridpack_path, "SubProcesses")) + shutil.copytree("backend", os.path.join(gridpack_path, "backend")) else: shutil.copytree("lib", lib_path) + for _, path in self.custom_function_paths: + target = path if os.path.isabs(path) else os.path.join(gridpack_path, path) + if not os.path.exists(target): + os.makedirs(os.path.dirname(target), exist_ok=True) + shutil.copy(path, target) if self.run_card["gridpack"]["include_madspace_source"]: shutil.copytree( @@ -1927,6 +2104,7 @@ def save_gridpack(self) -> None: shutil.copytree( _INSTALL_DIR / "madspace", os.path.join(gridpack_path, "madspace", "install", "madspace"), + ignore=shutil.ignore_patterns("include"), ) matrix_elements = [] @@ -1970,24 +2148,36 @@ def save_gridpack(self) -> None: # (bin/generate_events --output_format still selects an npy output). previous_output_format = self.run_card["run"]["output_format"] self.run_card["run"]["output_format"] = "lhe" + # the directories of this run mean nothing on the machine running the gridpack + previous_dirs = (self.run_card["gridpack"]["temp_output_dir"], + self.run_card["gridpack"]["output_dir"]) + self.run_card["gridpack"]["temp_output_dir"] = "" + self.run_card["gridpack"]["output_dir"] = "" try: self.run_card.write_gridpack_card( os.path.join(cards_path, "grid_run_card.toml")) finally: self.run_card["run"]["cpu_mode"] = previous_cpu_mode self.run_card["run"]["output_format"] = previous_output_format + (self.run_card["gridpack"]["temp_output_dir"], + self.run_card["gridpack"]["output_dir"]) = previous_dirs bin_path = os.path.join(gridpack_path, "bin") os.mkdir(bin_path) gen_events_file = os.path.join(bin_path, "generate_events") - shutil.copy( - os.path.join(os.path.dirname(__file__), "gridpack.py"), gen_events_file - ) + template_dir = os.path.dirname(__file__) + shutil.copy(os.path.join(template_dir, "gridpack.py"), gen_events_file) + 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 = { - "channels": channel_files, "matrix_elements": matrix_elements, + # loaded before the matrix elements, as in the run + "custom_functions": self.custom_function_paths, "me_parameters": me_parameters(self.run_card), "source_hash": ms.SOURCE_HASH, } @@ -2010,6 +2200,9 @@ def save_gridpack(self) -> None: systematics_data["nominal_info_file"] = os.path.join("data", name) with open(os.path.join(data_path, "systematics.json"), "w") as f: json.dump(systematics_data, f) + if self.run_card["gridpack"]["compress"]: + shutil.make_archive(gridpack_path, "gztar", self.run_path, "gridpack") + shutil.rmtree(gridpack_path) def save_gridpack_lhe_meta(self, path) -> None: """The run-independent part of the LHE header (cards, beams, PDF), for @@ -2370,11 +2563,14 @@ def __init__( self.unmerged_meta = None if unmerged_meta is None: api_path_formats = [self.meta["me_path"]] + member_metas = [self.meta] else: api_path_formats = [] + member_metas = [] for subproc in self.meta["subprocesses"]: submeta = unmerged_meta[subproc] api_path_formats.append(submeta["me_path"]) + member_metas.append(submeta) if len(api_path_formats) == 1: self.meta = submeta else: @@ -2407,8 +2603,14 @@ def cut_pids(observable_kwargs): # pdg id 0 is in no group, so the particle is never selected return [0 if i - n_in in uncut else pid for i, pid in enumerate(pids)] - self.cuts = ( - ms.Cuts([ + # custom cuts and scales see the PDG ids of the first flavor, like the + # built-in cuts, which select by flavor-independent groups + representative_pids = [self.flavor_pids()[0]] + custom_cut = self.process.custom_function( + "cut", self.particle_count, representative_pids + ) + if len(self.process.cut_data) > 0: + self.cuts = ms.Cuts([ ms.CutItem( observable=build_observable( cut_pids(cut_item.observable_kwargs), @@ -2419,10 +2621,11 @@ def cut_pids(observable_kwargs): mode=cut_item.mode, ) for cut_item in self.process.cut_data - ]) - if len(self.process.cut_data) > 0 - else None - ) + ], custom_cut) + elif custom_cut is not None: + self.cuts = ms.Cuts(self.particle_count, custom_cut) + else: + self.cuts = None # the integration histograms are functions of the momenta; the weight # distribution is a property of the final event sample, so it is only # filled by MadgraphProcess.build_event_histograms at combine time @@ -2452,7 +2655,11 @@ def cut_pids(observable_kwargs): ) self.scale = ms.EnergyScale( - particle_count=self.particle_count, **self.process.scale_kwargs + particle_count=self.particle_count, + custom_scale=self.process.custom_function( + "scale", self.particle_count, representative_pids + ), + **self.process.scale_kwargs ) if self.process.run_card["run"]["dummy_matrix_element"]: @@ -2466,6 +2673,46 @@ def cut_pids(observable_kwargs): me_parameters(self.process.run_card) ) self.matrix_elements.append(mat) + self.surrogates = [] + if self.process.surrogate is not None: + for path_format, member_meta in zip(api_path_formats, member_metas): + self.surrogates.append(self.load_surrogate(path_format, member_meta)) + + def load_surrogate(self, me_path_format: str, member_meta: dict): + """The [surrogate] of one matrix element of the subprocess.""" + name = re.sub( + r"^libmadmatrix_(.*)_\{device\}\.so$", r"\1", os.path.basename(me_path_format) + ) + # PDG ids for each flavor index of the matrix element + flavors = member_meta["flavors"] + pdg_ids = [flavors[0]["options"][0]] * (max(f["index"] for f in flavors) + 1) + for flavor in flavors: + pdg_ids[flavor["index"]] = flavor["options"][0] + settings = self.process.surrogate + apis = load_surrogate( + ms, + self.process.contexts, + self.process.backends, + self.process.device_types, + settings["function"], + name, + pdg_ids, + settings["param_card"], + ) + if len({api.index() for api in apis}) != 1: + raise RuntimeError("surrogate loaded at different indices") + return apis[0] + + def flavor_pids(self) -> list[list[int]]: + """PDG ids of the external particles for every flavor of the + 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: @@ -3099,8 +3346,37 @@ def build_integrands( flavor_mirror.append(flav["mirror"]) cross_sections = [] - for matrix_element in self.matrix_elements: - if matrix_element: + self.step_matrix_elements = [] + for me_index, matrix_element in enumerate(self.matrix_elements): + if self.surrogates: + # the surrogate in the integrand, the full matrix element in + # the second pass, with the random numbers passed in + mat = ms.MatrixElement( + self.surrogates[me_index].index(), + self.particle_count, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in], + [ms.MatrixElement.matrix_element_out], + 1, + False, + self.process.me_frame, + self.process.incoming_count, + ) + self.step_matrix_elements.append(ms.MatrixElement( + matrix_element, + [ms.MatrixElement.momenta_in, ms.MatrixElement.alpha_s_in, + ms.MatrixElement.flavor_in, ms.MatrixElement.random_color_in, + ms.MatrixElement.random_helicity_in, + ms.MatrixElement.random_diagram_in], + [ms.MatrixElement.matrix_element_out, + ms.MatrixElement.color_index_out, + ms.MatrixElement.helicity_index_out, + ms.MatrixElement.diagram_index_out], + False, + self.process.me_frame, + self.process.incoming_count, + )) + elif matrix_element: mat = ms.MatrixElement( matrix_element, ms.Integrand.matrix_element_inputs, @@ -3137,10 +3413,12 @@ def build_integrands( pdf2=pdf_arg, input_momentum_fraction=True, decay=self.process.is_decay, + return_matrix_element=bool(self.surrogates), ) ) partial_weights = self.process.systematics_enabled() madnis_args = self.process.run_card["madnis"] + custom_bias = self.process.custom_function("bias", self.particle_count, flavors) integrands = [] for channel in phasespace.channels: integrands.append(ms.Integrand( @@ -3171,7 +3449,8 @@ def build_integrands( flavor_diff_xs_indices, flavor_subproc_indices, flavor_per_subproc_remap, - madnis_args["compressed_channel_weight_count"] + madnis_args["compressed_channel_weight_count"], + custom_bias, )) #print(integrands[1].function()) #for i in integrands: print(i.function()) @@ -4226,6 +4505,8 @@ def run_single(switch=None) -> "MadgraphProcess": process.survey() process.train_madnis() process.generate_events() + if process.gridpack_run_mode != "regular": + return process # run_card-driven LHE post-processing (displaced vertex + systematics) run_lhe_postprocessing(process) # run the post-processing tools (Pythia8/Delphes/MadSpin/reweight/analysis) diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 7c571d12d7..ea618c4862 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -43,6 +43,13 @@ save_gridpack = %(gridpack.save_gridpack)s include_source = %(gridpack.include_source)s include_madspace = %(gridpack.include_madspace)s include_madspace_source = %(gridpack.include_madspace_source)s +compress = %(gridpack.compress)s +# 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 +temp_output_dir = %(gridpack.temp_output_dir)s +# output files; "" creates a new run folder in Events +output_dir = %(gridpack.output_dir)s [beam] e_cm = %(beam.e_cm)s @@ -114,6 +121,27 @@ bw_cutoff = %(phasespace.bw_cutoff)s cut_decays = %(phasespace.cut_decays)s # apply the [cuts] to decay products of on-shell particles adaptive_symmetry_sampling = %(phasespace.adaptive_symmetry_sampling)s +[custom_functions] +# functions created with bin/init_custom_func: a .so library, e.g. "lib/custom_cut.so", +# or a Python file, e.g. "custom_cut/custom_cut.py" +# paths are relative to the process directory; "" disables the function +cut = %(custom_functions.cut)s +scale = %(custom_functions.scale)s # replaces dynamical_scale_choice +pdf = %(custom_functions.pdf)s # replaces the PDF set, alpha_s still comes from [beam] pdf +bias = %(custom_functions.bias)s # events are generated with |w| b and written with w / b + +[surrogate] +# a fast surrogate of the matrix element in the integrand: a .so library, which can +# contain {subprocess} and {device}, or a Python file defining +# custom_surrogate(momenta, pdg_ids, alpha_s); "" disables it +function = %(surrogate.function)s +# unweight or reweight the events with the full matrix element in a second pass +mode = %(surrogate.mode)s +# unweighting with the surrogate (unweight mode); [generation] sets the second pass +first_pass_max_overweight_truncation = %(surrogate.first_pass_max_overweight_truncation)s +first_pass_freeze_max_weight_after = %(surrogate.first_pass_freeze_max_weight_after)s +param_card = %(surrogate.param_card)s # for a .so surrogate, "" uses Cards/param_card.dat + [multiparticles] $multiparticles diff --git a/madgraph/various/banner.py b/madgraph/various/banner.py index 7619d334b8..c904f8691d 100755 --- a/madgraph/various/banner.py +++ b/madgraph/various/banner.py @@ -6755,6 +6755,16 @@ def default_setup(self): comment="include the installed madspace package") self.add_toml_param('gridpack', 'include_madspace_source', False, comment="include the madspace source code") + self.add_toml_param('gridpack', 'compress', False, + comment="compress the gridpack into a tarball") + self.add_toml_param('gridpack', 'run_mode', "minimal", + allowed=['regular', 'minimal', 'fix_max_weight'], + comment="how the run creating the gridpack ends: regular (with events), " + "minimal or fix_max_weight (no events)") + self.add_toml_param('gridpack', 'temp_output_dir', "", gridpack=True, + comment="directory for the temporary npy files; \"\" uses the output directory") + self.add_toml_param('gridpack', 'output_dir', "", gridpack=True, + comment="directory for the output files; \"\" creates a new run folder in Events") # ----------------------------- [beam] ------------------------- self.add_toml_param('beam', 'e_cm', 13000.0, @@ -6883,6 +6893,31 @@ def default_setup(self): self.add_toml_param('phasespace', 'adaptive_symmetry_sampling', True, comment="sample symmetric channels adaptively") + # ------------------------- [custom_functions] ----------------- + # user functions built with bin/init_custom_func; "" disables them + self.add_toml_param('custom_functions', 'cut', "", + comment="custom cut, a .so library or a .py file, e.g. lib/custom_cut.so") + self.add_toml_param('custom_functions', 'scale', "", + comment="custom scale, replaces dynamical_scale_choice") + self.add_toml_param('custom_functions', 'pdf', "", + comment="custom PDF, replaces the PDF set (not alpha_s)") + self.add_toml_param('custom_functions', 'bias', "", + comment="custom event bias") + + # ---------------------------- [surrogate] --------------------- + # a fast surrogate of the matrix element for the integrand, followed by + # a second pass with the full matrix element; "" disables it + self.add_toml_param('surrogate', 'function', "", + comment="surrogate matrix element, a .so library or a .py file") + self.add_toml_param('surrogate', 'mode', "unweight", allowed=["unweight", "reweight"], + comment="unweight or reweight the events with the full matrix element") + self.add_toml_param('surrogate', 'first_pass_max_overweight_truncation', 0.01, + comment="max_overweight_truncation of the unweighting with the surrogate") + self.add_toml_param('surrogate', 'first_pass_freeze_max_weight_after', 10000, + comment="freeze_max_weight_after of the unweighting with the surrogate") + self.add_toml_param('surrogate', 'param_card', "", + comment="param card for a .so surrogate; \"\" uses Cards/param_card.dat") + # ----------------------------- [madnis] ----------------------- self.add_toml_param('madnis', 'enable', False, allowed=["auto", True, False], auto=True, comment="train MadNIS networks to improve the sampling; auto decides by process") diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index 316453d193..b69730fa95 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -209,10 +209,13 @@ 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 src/driver/io.cpp + src/driver/tar.cpp src/driver/vegas_optimizer.cpp src/driver/discrete_optimizer.cpp src/driver/backend.cpp @@ -221,6 +224,7 @@ add_library( src/driver/event_histograms.cpp src/driver/madnis_training.cpp src/driver/random.cpp + src/driver/weighting_pipeline.cpp src/compgraphs/type.cpp src/compgraphs/function.cpp src/compgraphs/instruction.cpp @@ -240,6 +244,7 @@ add_library( src/phasespace/madnis.cpp src/phasespace/rambo.cpp src/phasespace/cuts.cpp + src/phasespace/custom_function.cpp src/phasespace/observable.cpp src/phasespace/histograms.cpp src/phasespace/chili.cpp @@ -256,6 +261,7 @@ add_library( src/phasespace/matrix_element.cpp src/phasespace/cross_section.cpp src/phasespace/scale.cpp + src/phasespace/weighting_step.cpp include/madspace/util.hpp include/madspace/driver.hpp include/madspace/driver/adam_optimizer.hpp @@ -267,6 +273,7 @@ add_library( include/madspace/driver/channel_generator.hpp include/madspace/driver/generator_data.hpp include/madspace/driver/io.hpp + include/madspace/driver/tar.hpp include/madspace/driver/vegas_optimizer.hpp include/madspace/driver/discrete_optimizer.hpp include/madspace/driver/backend.hpp @@ -274,6 +281,7 @@ add_library( include/madspace/driver/logger.hpp include/madspace/driver/madnis_training.hpp include/madspace/driver/random.hpp + include/madspace/driver/weighting_pipeline.hpp include/madspace/compgraphs.hpp include/madspace/compgraphs/type.hpp include/madspace/compgraphs/function.hpp @@ -312,6 +320,7 @@ add_library( include/madspace/phasespace/matrix_element.hpp include/madspace/phasespace/cross_section.hpp include/madspace/phasespace/scale.hpp + include/madspace/phasespace/weighting_step.hpp include/madspace/mixmax/mixmax.hpp include/madspace/mixmax/mixmax_skip_N17.c ) @@ -533,6 +542,8 @@ python_add_library( ${MADSPACE_GENERATED_DIR}/python/docstrings.hpp src/python/function_runtime.cpp src/python/function_runtime.hpp + src/python/python_function.cpp + src/python/python_function.hpp WITH_SOABI ) target_include_directories(_madspace_py PRIVATE ${MADSPACE_GENERATED_DIR}/python) diff --git a/madspace/generate_pyi.py b/madspace/generate_pyi.py index c1581c1d6b..03633a792a 100644 --- a/madspace/generate_pyi.py +++ b/madspace/generate_pyi.py @@ -31,7 +31,7 @@ def main(): shutil.rmtree(stage_pkg) (stage_pkg / "lib").mkdir(parents=True) - for name in ("__init__.py", "_madspace_py_loader.py"): + for name in ("__init__.py", "_madspace_py_loader.py", "python_function.py"): shutil.copy2(args.package_dir / name, stage_pkg / name) shutil.copy2(args.module, stage_pkg / args.module.name) shutil.copy2(args.lib, stage_pkg / "lib" / args.lib.name) @@ -57,6 +57,7 @@ def main(): "--enum-class-locations=Activation:MLP.Activation", "--enum-class-locations=CutMode:Cuts.CutMode", "--enum-class-locations=LRSchedule:AdamOptimizer.LRSchedule", + "--enum-class-locations=Mode:WeightingStep.Mode", # Fail the build on unresolvable names/expressions (e.g. an enum # default value needing its own --enum-class-locations entry) # instead of silently emitting a stub with holes in it. diff --git a/madspace/include/madspace/driver.hpp b/madspace/include/madspace/driver.hpp index 26afff2a88..74142e3ff4 100644 --- a/madspace/include/madspace/driver.hpp +++ b/madspace/include/madspace/driver.hpp @@ -5,6 +5,9 @@ #include "driver/context.hpp" #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" @@ -12,7 +15,8 @@ #include "driver/madnis_training.hpp" #include "driver/status_file.hpp" #include "driver/systematics.hpp" -#include "driver/event_histograms.hpp" +#include "driver/tar.hpp" #include "driver/tensor.hpp" #include "driver/thread_pool.hpp" #include "driver/vegas_optimizer.hpp" +#include "driver/weighting_pipeline.hpp" diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 1830817719..c7434ed177 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/madspace/include/madspace/driver/channel_generator.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -19,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. * @@ -39,12 +69,23 @@ class ChannelEventGenerator { const std::string& weight_file, const GeneratorConfig& config ); + /// Like @ref load, from the JSON text produced by @ref to_json instead of + /// a file. + static ChannelEventGenerator load_json( + const std::string& channel_json, + const std::vector& contexts, + const std::string& event_file, + const std::string& weight_file, + const GeneratorConfig& config + ); /** * @param contexts One context per device to run on. * @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. @@ -77,12 +118,30 @@ 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); + /// Set the integral of the absolute weights known from an earlier run, with + /// its standard error. Both must be positive. + void set_abs_integral_prior(double mean, double error); + /// Whether @ref set_abs_integral_prior was called. + bool has_abs_integral_prior() const { return _abs_integral_prior.has_value(); } + /// The integral of the absolute weights: the mean of this run, combined + /// with the prior (if any) weighted by the inverse variances. Used to share + /// the events between channels. 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 @@ -103,6 +162,104 @@ class ChannelEventGenerator { int particle_layout_extra_flags() const { return _particle_layout_extra_flags; } /// Layout of @ref event_file. const DataLayout& event_file_layout() const { return _event_file_layout; } + /// Layout of @ref weight_file. + const DataLayout& weight_layout() const { + return _step_settings.empty() ? weight_file_layout : step_weight_file_layout; + } + /// Records per event in @ref weight_file, none without steps. + std::size_t weight_record_count() const { + return _step_settings.empty() ? 0 : _step_settings.size() + 1; + } + /// Types of the unweighted events, the integrand outputs before `random`. + const NamedVector& event_types() const { return _event_types; } + /// Whether the matrix element of the integrand returns the color, helicity + /// and diagram index. + const std::array& provided_indices() const { return _provided_indices; } + + /// Integral of the final weights: after the last re- or unweighting step, + /// or the weights of the integrand without steps. + struct IntegralEstimate { + double mean; + double variance; + double mean_abs; + double variance_abs; + std::size_t count; + }; + /// See @ref IntegralEstimate. With steps, but no events through all of + /// them yet, the integral of the first step. + IntegralEstimate result_integral() const; + /// The integral of the first step, see @ref IntegralEstimate. + IntegralEstimate first_step_integral() const; + + /// Set the re- and unweighting steps after the unweighting of this channel, + /// see @ref WeightingStep. Their events are generated by the + /// @ref EventGenerator, before any events of this channel. + void set_weighting_steps(const std::vector& steps); + /// The re- and unweighting steps. + const std::vector& weighting_steps() const { return _step_settings; } + /// Maximum weight of step `step`, counted from 1. + double step_max_weight(std::size_t step) const { + return _step_states.at(step - 1).max_weight; + } + /// 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 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`. @@ -141,11 +298,35 @@ class ChannelEventGenerator { /// Append `unweighted_events` to @ref event_file, recording /// `job_max_weight`. void write_events(const TensorVec& unweighted_events, double job_max_weight); + /// Pack `unweighted_events` into `event_buffer` (layout @ref + /// event_file_layout) and their weights into `weight_buffer` (layout + /// `weight_file_layout`), resizing both. + void fill_event_buffers( + const TensorVec& unweighted_events, + EventBuffer& event_buffer, + EventBuffer& weight_buffer + ) const; /// Serialize this generator's state to `file_name`; see @ref load. void save(const std::string& file_name) const; + /// The JSON text that @ref save writes; see @ref load_json. With + /// `include_estimates`, the current maximum weight and the integral of the + /// absolute weights are stored as well. @ref load_json then keeps the + /// maximum weight fixed and uses the integral as a prior, see + /// @ref abs_integral_estimate. + std::string to_json(bool include_estimates = false) const; private: + struct StepState { + double max_weight = 0.; + bool max_weight_fixed = false; + std::vector large_weights; + double count_accepted = 0.; + }; + void apply_truncation_budget(); + void update_status_integrals(); + void unweight_step_file(MixMaxRandom& rand_gen); + StepChainResult step_chain(EventBuffer& weight_buffer, std::size_t index) const; ChannelEventGenerator( const std::vector& contexts, std::size_t particle_count, @@ -173,6 +354,7 @@ class ChannelEventGenerator { RuntimePtr vegas_histogram = nullptr; RuntimePtr discrete_histogram = nullptr; RuntimePtr observable_histograms = nullptr; + RuntimePtr step_random = nullptr; }; struct FieldIndices { @@ -181,6 +363,8 @@ class ChannelEventGenerator { int ren_scale, alpha_qcd; int x1, fact_scale1, x2, fact_scale2, partial_weight_product; int subprocess_index; + int bias; + int matrix_element; int random, rest; }; @@ -191,8 +375,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; @@ -205,6 +390,9 @@ class ChannelEventGenerator { RunningIntegral _cross_section; RunningIntegral _abs_cross_section; double _max_weight = 0.; + bool _max_weight_fixed = false; + // mean and standard error + std::optional> _abs_integral_prior; std::size_t _count_requested_opt = 0; // Monotonic per-job counters keying each job's deterministic random stream; // never reset. Separate for survey/generate so neither depends on the other. @@ -223,6 +411,14 @@ class ChannelEventGenerator { std::vector _histograms; std::unordered_set _used_globals; FieldIndices _field_indices; + NamedVector _event_types; + std::array _provided_indices{true, true, true}; + std::vector _step_settings; + std::vector _step_states; + std::vector _loaded_step_max_weights; + double _first_step_count = 0.; + StepIntegral _final_integral; + std::size_t _event_epoch = 0; friend void to_json(nlohmann::json& j, const ChannelEventGenerator& channel); }; diff --git a/madspace/include/madspace/driver/context.hpp b/madspace/include/madspace/driver/context.hpp index 521cfb4493..0d1ed0eeb5 100644 --- a/madspace/include/madspace/driver/context.hpp +++ b/madspace/include/madspace/driver/context.hpp @@ -1,7 +1,9 @@ #pragma once #include +#include #include +#include #include #include @@ -12,6 +14,31 @@ namespace madspace { +/** + * Table of UMAMI entry points, one per function of the UMAMI C interface. + * + * Used to build a @ref MatrixElementApi without a shared library, with @ref + * Context::add_matrix_element. Each entry accepts a plain C function pointer + * or a capturing callable. `get_meta`, `initialize`, `matrix_element` and + * `free` are required. The others are optional: missing key queries report + * `UMAMI_ERROR_NOT_IMPLEMENTED`, and a missing `set_parameter` only fails if + * parameters have to be set. + */ +struct UmamiFunctions { + std::function> get_meta; + std::function> + supported_inputs; + std::function> + required_inputs; + std::function> + supported_outputs; + std::function> initialize; + std::function> set_parameter; + std::function> + matrix_element; + std::function> free; +}; + /** * Loaded UMAMI shared library for one matrix element. * @@ -20,8 +47,9 @@ namespace madspace { * MatrixElement for the compute-graph function built on top of it. A * `MatrixElementApi` dynamically loads the library, resolves its exported * UMAMI entry points, and exposes the metadata queries and the raw @ref call. - * Obtained from @ref Context::load_matrix_element, never constructed - * directly. + * Obtained from @ref Context::load_matrix_element, or from @ref + * Context::add_matrix_element for entry points given as @ref UmamiFunctions, + * never constructed directly. * * **References** * - [1] S. Hageböck et al., "Data-parallel leading-order event generation in @@ -36,7 +64,7 @@ class MatrixElementApi { /// Device the underlying library was built for. DeviceType device_type() const { UmamiDevice dev; - check_umami_status(_get_meta(UMAMI_META_DEVICE, &dev)); + check_umami_status(_functions.get_meta(UMAMI_META_DEVICE, &dev)); switch (dev) { case UMAMI_DEVICE_CPU: return DeviceType::cpu; @@ -51,30 +79,31 @@ class MatrixElementApi { /// Number of external particles. std::size_t particle_count() const { int count; - check_umami_status(_get_meta(UMAMI_META_PARTICLE_COUNT, &count)); + check_umami_status(_functions.get_meta(UMAMI_META_PARTICLE_COUNT, &count)); return count; } /// Number of Feynman diagrams. std::size_t diagram_count() const { int count; - check_umami_status(_get_meta(UMAMI_META_DIAGRAM_COUNT, &count)); + check_umami_status(_functions.get_meta(UMAMI_META_DIAGRAM_COUNT, &count)); return count; } /// Number of helicity configurations. std::size_t helicity_count() const { int count; - check_umami_status(_get_meta(UMAMI_META_HELICITY_COUNT, &count)); + check_umami_status(_functions.get_meta(UMAMI_META_HELICITY_COUNT, &count)); return count; } /// Index of this matrix element in its @ref Context, in load order. std::size_t index() const { return _index; } - /// Path of the loaded UMAMI shared library. + /// Path of the loaded UMAMI shared library, or the name given to @ref + /// Context::add_matrix_element. const std::string& file_name() const { return _file_name; } /// UMAMI input keys this library can accept. std::vector supported_inputs() const { bool const* data; int count; - check_umami_status(_supported_inputs(&data, &count)); + check_umami_status(_functions.supported_inputs(&data, &count)); std::vector result(UMAMI_INPUT_KEY_COUNT, false); for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i) { result[i] = data[i]; @@ -85,7 +114,7 @@ class MatrixElementApi { std::vector required_inputs() const { bool const* data; int count; - check_umami_status(_required_inputs(&data, &count)); + check_umami_status(_functions.required_inputs(&data, &count)); std::vector supported = supported_inputs(); for (int i = 0; i < count && i < UMAMI_INPUT_KEY_COUNT; ++i) { if (data[i] && !supported[i]) { @@ -106,7 +135,7 @@ class MatrixElementApi { std::vector supported_outputs() const { bool const* data; int count; - check_umami_status(_supported_outputs(&data, &count)); + check_umami_status(_functions.supported_outputs(&data, &count)); std::vector result(UMAMI_OUTPUT_KEY_COUNT, false); for (int i = 0; i < count && i < UMAMI_OUTPUT_KEY_COUNT; ++i) { result[i] = data[i]; @@ -128,7 +157,7 @@ class MatrixElementApi { UmamiOutputKey const* output_keys, void* const* outputs ) const { - check_umami_status(_matrix_element( + check_umami_status(_functions.matrix_element( handle, count, stride, @@ -154,18 +183,27 @@ class MatrixElementApi { std::size_t index = 0, const std::unordered_map& parameters = {} ); + MatrixElementApi( + UmamiFunctions functions, + const std::string& name, + const std::string& param_card, + ThreadPool& thread_pool, + DevicePtr device, + std::size_t index = 0, + const std::unordered_map& parameters = {} + ); + void init_functions(const std::unordered_map& parameters); + void init_instances( + ThreadPool& thread_pool, + DevicePtr device, + const std::string& param_card, + const std::unordered_map& parameters + ); void check_umami_status(UmamiStatus status) const; [[noreturn]] void throw_error(const std::string& message) const; std::unique_ptr> _shared_lib; - decltype(&umami_get_meta) _get_meta; - decltype(&umami_supported_inputs) _supported_inputs; - decltype(&umami_required_inputs) _required_inputs; - decltype(&umami_supported_outputs) _supported_outputs; - decltype(&umami_initialize) _initialize; - decltype(&umami_set_parameter) _set_parameter; - decltype(&umami_matrix_element) _matrix_element; - decltype(&umami_free) _free; + UmamiFunctions _functions; using InstanceType = std::unique_ptr>; ThreadResource _instances; std::string _file_name; @@ -217,6 +255,16 @@ class Context { const std::string& param_card, const std::unordered_map& parameters = {} ); + /// Add a matrix element from the UMAMI entry points in `functions`, + /// without a shared library. `name` identifies it in error messages, and + /// `param_card` and `parameters` are passed on as in @ref + /// load_matrix_element. + const MatrixElementApi& add_matrix_element( + UmamiFunctions functions, + const std::string& name, + const std::string& param_card = "", + const std::unordered_map& parameters = {} + ); /// Create and zero-initialize a new global named `name`. Tensor define_global( const std::string& name, @@ -244,10 +292,11 @@ class Context { /// The matrix element previously loaded at position `index`; see @ref /// load_matrix_element. const MatrixElementApi& matrix_element(std::size_t index) const; - /// Write every global to `dir`, one `.npy` file per name. - void save_globals(const std::string& dir) const; - /// Define a global from every `.npy` file in `dir`, named after the file. - void load_globals(const std::string& dir); + /// Write every global to the tar archive `file`, as one `.npy` file per name. + void save_globals(const std::string& file) const; + /// Define a global from every `.npy` file in the tar archive `file`, named + /// after the file. + void load_globals(const std::string& file); /// The device every global and matrix element of this context lives on. DevicePtr device() { return _device; } /// The thread pool backing this context's device-local resources. diff --git a/madspace/include/madspace/driver/event_generator.hpp b/madspace/include/madspace/driver/event_generator.hpp index 7ad0aae2bb..eaaf1d511a 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; @@ -73,6 +79,15 @@ class EventGenerator { /// Generate unweighted events until every channel reaches its target /// count. void generate(); + /// Run only the VEGAS and channel-weight optimization of every channel + /// that still needs it, until all are optimized. No events are kept. + void optimize(); + /// Like @ref generate, but every channel with a nonzero cross section runs + /// until it has `freeze_max_weight_after` unweighted events, regardless of + /// its share of the cross section, so that its maximum weight is final + /// (see @ref ChannelEventGenerator::max_weight). The events are not meant + /// to be combined. + void fix_max_weights(); /** * Combine every channel's generated events into a single unweighted * `.npy` file. @@ -175,6 +190,7 @@ class EventGenerator { // dispatch, drained with priority by start_jobs(). std::vector> _context_unweight_queue; ResultQueue _result_queue; + std::unique_ptr _pipeline; // Base seed for reproducible event generation. std::uint64_t _seed; @@ -184,6 +200,13 @@ class EventGenerator { // Salted by this counter so repeated calls don't replay the same stream. std::size_t _unweight_call_index = 0; + // What a run of the generate loop does; see run_generate(). + enum class GenerateMode { events, optimize_only, fix_max_weights }; + void run_generate(GenerateMode mode); + // Channels with an optimizer, and how many of them are optimized. + std::pair optimized_channel_count() const; + GenerateMode _generate_mode = GenerateMode::events; + // Scheduling context for the running survey()/generate() call, read by // start_jobs() to derive job seeds. bool _survey_job = false; @@ -213,25 +236,27 @@ 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 + std::size_t start_pipeline_jobs(); // Layout of the combined events as read from the channel files DataLayout combined_layout() const; // Layout of the written events: the reweighting inputs are dropped unless diff --git a/madspace/include/madspace/driver/event_stream.hpp b/madspace/include/madspace/driver/event_stream.hpp new file mode 100644 index 0000000000..7839f886df --- /dev/null +++ b/madspace/include/madspace/driver/event_stream.hpp @@ -0,0 +1,184 @@ +#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 every channel. `count_unweighted` is the number of + /// events returned so far. + GeneratorStatus status() const; + /// Per-channel status; `count_unweighted` is the number of events of the + /// channel returned so far. + std::vector channel_status() const; + /// Current probability of each channel to be picked for the next event. + std::vector channel_probabilities() const; + /// Number of events returned so far. + std::size_t event_count() const; + /// Largest particle count of the returned LHE events. + std::size_t max_particle_count() const { + return _lhe_completer.max_particle_count(); + } + +private: + // A job whose events are buffered, waiting to be consumed. + struct ClosedJob { + GeneratorBatchJob job; + EventBuffer events; + EventBuffer weights; + bool opened = false; + std::size_t next_event = 0; + // 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..e3fbd51c99 --- /dev/null +++ b/madspace/include/madspace/driver/event_utils.hpp @@ -0,0 +1,71 @@ +#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); + +/** + * 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. + */ +void copy_channel_event( + EventBuffer& in_buffer, + std::size_t in_index, + EventBuffer& out_buffer, + std::size_t out_index, + double weight, + int subprocess_index, + int extra_flags +); + +/// Build `lhe_event` from event `event_index` of a combined buffer and +/// complete it with `lhe_completer`. +void fill_lhe_event( + LHECompleter& lhe_completer, + LHEEvent& lhe_event, + EventBuffer& buffer, + std::size_t event_index, + MixMaxRandom& rand_gen +); + +/// Write the particles of `lhe_event` to event `event_index` of a buffer with +/// LHE particle layout, padding the remaining slots with empty particles. +void write_lhe_particles( + const LHEEvent& lhe_event, EventBuffer& buffer, std::size_t event_index +); + +/// Sum the integrals and event counts of `channels` into `status`. Throws if +/// an integral is not finite. +void sum_channel_status( + const std::vector>& channels, + GeneratorStatus& status +); + +} // namespace madspace diff --git a/madspace/include/madspace/driver/generator_data.hpp b/madspace/include/madspace/driver/generator_data.hpp index ee2c4e87e9..447a5aec5f 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. * @@ -192,6 +243,10 @@ struct GeneratorStatus { std::size_t count_target; /// Number of survey iterations run. std::size_t iterations; + /// Number of optimization iterations in a row without enough improvement; + /// the channel counts as optimized when this reaches + /// @ref GeneratorConfig::optimization_patience. + std::size_t iters_without_improvement; /// Whether the survey/optimization phase has finished. bool optimized; /// Whether generation of this subprocess is complete. @@ -243,6 +298,8 @@ struct GeneratorBatchJob { // for them and isn't read. std::size_t split_job_count; Tensor weights; + // per-event bias, empty without a custom bias + Tensor biases; TensorVec events; TensorVec unweighted_events; TensorVec hists; @@ -251,6 +308,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; @@ -265,6 +324,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 5e631ad78d..a9b9b9592a 100644 --- a/madspace/include/madspace/driver/io.hpp +++ b/madspace/include/madspace/driver/io.hpp @@ -6,12 +6,19 @@ #include #include "madspace/driver/lhe_output.hpp" +#include "madspace/driver/tar.hpp" #include "madspace/driver/tensor.hpp" namespace madspace { +/// Read a tensor from the npy file `file`. Tensor load_tensor(const std::string& file); +/// Write `tensor` to the npy file `file`. void save_tensor(const std::string& file, Tensor tensor); +/// Read a tensor from the npy file `name` inside the tar archive `tar`. +Tensor load_tensor(TarReader& tar, const std::string& name); +/// Append `tensor` as an npy file called `name` to the tar archive `tar`. +void save_tensor(TarWriter& tar, const std::string& name, Tensor tensor); struct FieldLayout { static constexpr const char* i32 = " cluster_scale() { return &_data[_offsets[0] + 0]; } + // weight files with re- and unweighting steps store one record per step + UnalignedRef step_ratio() { return &_data[0]; } + UnalignedRef step_max_weight() { return &_data[8]; } + void from_lhe_particle(const LHEParticle& particle) { lhe_pdg_id() = particle.pdg_id; lhe_status_code() = particle.status_code; @@ -149,13 +160,12 @@ class EventRecord { static constexpr int f_beam2 = 32; static constexpr int f_partial_weights = 64; static constexpr int f_syst_weights = 128; + static constexpr int f_bias = 256; // `syst_weight_count` columns of systematic variation weights are appended // when f_syst_weights is set; `syst_weight_ids` names them (rwgt_). - static std::vector layout( - int fields, - const std::vector& syst_weight_ids = {} - ) { + static std::vector + layout(int fields, const std::vector& syst_weight_ids = {}) { std::vector ret; if (fields & f_weight) { ret.push_back({"weight", FieldLayout::f64, 0}); @@ -214,6 +224,9 @@ class EventRecord { ret.push_back({std::format("rwgt_{}", id), FieldLayout::f64, 7}); } } + if (fields & f_bias) { + ret.push_back({"bias", FieldLayout::f64, 8}); + } return ret; } @@ -257,6 +270,9 @@ class EventRecord { return &_data[_offsets[6] + 8 * index]; } + // event bias, divided out of the weight when the channels are combined + UnalignedRef bias() { return &_data[_offsets[7] + 0]; } + void from_lhe_event(const LHEEvent& event) { lhe_process_id() = event.process_id; lhe_weight() = event.weight; @@ -327,6 +343,13 @@ inline const DataLayout weight_file_layout = DataLayout( ParticleRecord::layout(ParticleRecord::f_none) ); +/// Weight file of a channel with re- and unweighting steps: the weight of the +/// first step and, as "particles", the ratio and the maximum weight of each step. +inline const DataLayout step_weight_file_layout = DataLayout( + EventRecord::layout(EventRecord::f_weight), + {{"ratio", FieldLayout::f64, 0}, {"max_weight", FieldLayout::f64, 0}} +); + class EventBuffer { public: EventBuffer( diff --git a/madspace/include/madspace/driver/random.hpp b/madspace/include/madspace/driver/random.hpp index 821d054a25..e617ecd9f4 100644 --- a/madspace/include/madspace/driver/random.hpp +++ b/madspace/include/madspace/driver/random.hpp @@ -32,6 +32,9 @@ struct DerivedSeed { madnis_unweight, madnis_sample_buffer, global_init, + // per-event random numbers of the re- and unweighting steps; at most + // 16 types fit into the seed + weighting_step_random, }; std::array seed_parts; diff --git a/madspace/include/madspace/driver/tar.hpp b/madspace/include/madspace/driver/tar.hpp new file mode 100644 index 0000000000..1d4ffa7453 --- /dev/null +++ b/madspace/include/madspace/driver/tar.hpp @@ -0,0 +1,66 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace madspace { + +/// Writes a POSIX ustar archive, readable by any standard tar tool. Entries are +/// regular files, added one at a time. Names longer than 100 characters are +/// split over the ustar `prefix` and `name` fields; names that do not fit even +/// then are rejected. The archive is finished by @ref close or the destructor. +class TarWriter { +public: + /// Create (or truncate) the archive at `file_name`. + explicit TarWriter(const std::string& file_name); + TarWriter(const TarWriter&) = delete; + TarWriter& operator=(const TarWriter&) = delete; + ~TarWriter(); + /// Append a file called `name` with the given contents. + void add(const std::string& name, const char* data, std::size_t size); + /// Write the end-of-archive marker and close the file. + void close(); + +private: + std::string _file_name; + std::ofstream _stream; + bool _closed; +}; + +/// Reads the regular files of a tar archive (ustar or GNU, as written by +/// standard tools). The archive is indexed when opened; entries are then read +/// by name. Entries other than regular files are skipped, and long names +/// stored in GNU or pax extension headers are not interpreted. +class TarReader { +public: + /// Open and index the archive at `file_name`. + explicit TarReader(const std::string& file_name); + /// The names of all files in the archive, in sorted order. + std::vector names() const; + /// Whether the archive contains a file called `name`. + bool contains(const std::string& name) const; + /// The size in bytes of the file called `name`. + std::size_t size(const std::string& name) const; + /// Read the file called `name` into `data`, which must hold `size(name)` + /// bytes. + void read(const std::string& name, char* data); + /// The contents of the file called `name`. + std::vector read(const std::string& name); + +private: + struct Entry { + std::uint64_t offset; + std::uint64_t size; + }; + const Entry& entry(const std::string& name) const; + + std::string _file_name; + std::ifstream _stream; + std::map _entries; +}; + +} // namespace madspace diff --git a/madspace/include/madspace/driver/tensor.hpp b/madspace/include/madspace/driver/tensor.hpp index c6fe9a20eb..0feb424426 100644 --- a/madspace/include/madspace/driver/tensor.hpp +++ b/madspace/include/madspace/driver/tensor.hpp @@ -747,7 +747,8 @@ class Tensor { return copy(*impl->device, hint); } - /// Whether the elements are stored without gaps in row-major order. + /// Whether the elements are stored without gaps in column-major (Fortran) + /// order, i.e. the first index varies fastest. bool is_contiguous() const { return impl->contiguous_dims == impl->shape.size(); } /// Number of leading dimensions that are contiguous. diff --git a/madspace/include/madspace/driver/weighting_pipeline.hpp b/madspace/include/madspace/driver/weighting_pipeline.hpp new file mode 100644 index 0000000000..a61a6c6a47 --- /dev/null +++ b/madspace/include/madspace/driver/weighting_pipeline.hpp @@ -0,0 +1,160 @@ +#pragma once + +#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; +}; + +/// 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, or of all collected events with + /// `flush`. Returns the number of started jobs. + std::size_t start_jobs(ResultQueue& result_queue, bool flush); + /// Process the result of the finished job `job_id`. + void commit_job(std::size_t job_id); + /// Whether events are still collected or in jobs. + bool busy() const { return _chunk_count > 0; } + /// Drop all events, after a failed job. + void reset(); + +private: + // the unweighted events of one job of a channel + struct Chunk { + std::size_t channel_index; + std::size_t 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; + RuntimePtr concat; + std::vector steps; + }; + struct Group { + std::vector settings; + std::vector runtimes; + // [context][stage - 1], stage step_count + 1 copies the events to the CPU + std::vector>> buffers; + std::vector> buffer_sizes; + std::vector multiple_matrix_elements; + int momenta, alpha_s, flavor, me_index, color, helicity, diagram, random; + }; + struct Job { + std::size_t group; + std::size_t stage; + std::size_t context_index; + std::vector chunks; + }; + + void advance(ChunkPtr chunk, std::size_t stage); + void set_ready(ChunkPtr chunk, std::size_t stage); + void commit_stage(ChunkPtr chunk, std::size_t stage); + void decide(Chunk& chunk, std::size_t step); + void run_job(Job& job); + void gather_events(Chunk& chunk, ContextRuntimes& runtimes); + std::size_t batch_size(std::size_t context_index) const; + + std::vector _contexts; + std::vector> _channels; + GeneratorConfig _config; + std::vector _groups; + std::vector _channel_groups; + std::vector _channel_sequences; + // per channel and stage: the next sequence to commit, and the finished + // chunks waiting for their turn + std::vector> _next_sequences; + std::vector>> _ready; + std::unordered_map _jobs; + std::size_t _job_counter = 0; + std::size_t _chunk_count = 0; + Sink _sink; +}; + +} // namespace madspace diff --git a/madspace/include/madspace/phasespace.hpp b/madspace/include/madspace/phasespace.hpp index ca6cb3cf8d..7ffa30a15a 100644 --- a/madspace/include/madspace/phasespace.hpp +++ b/madspace/include/madspace/phasespace.hpp @@ -6,6 +6,7 @@ #include "phasespace/channel_weights.hpp" #include "phasespace/color_ordered_mapping.hpp" #include "phasespace/cross_section.hpp" +#include "phasespace/custom_function.hpp" #include "phasespace/cuts.hpp" #include "phasespace/discrete_flow.hpp" #include "phasespace/discrete_sampler.hpp" @@ -29,3 +30,4 @@ #include "phasespace/two_particle.hpp" #include "phasespace/unweighter.hpp" #include "phasespace/vegas.hpp" +#include "phasespace/weighting_step.hpp" diff --git a/madspace/include/madspace/phasespace/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/custom_function.hpp b/madspace/include/madspace/phasespace/custom_function.hpp new file mode 100644 index 0000000000..8b04abee39 --- /dev/null +++ b/madspace/include/madspace/phasespace/custom_function.hpp @@ -0,0 +1,103 @@ +#pragma once + +#include "madspace/driver/context.hpp" +#include "madspace/phasespace/base.hpp" + +namespace madspace { + +/** + * A user-supplied function of the event momenta, loaded as a UMAMI library. + * + * Custom cuts, scales and biases use the UMAMI interface as a generic batched + * function. The library is loaded like a matrix element, with + * Context::load_matrix_element. It receives the momenta + * (`UMAMI_IN_MOMENTA`), the particle count (`UMAMI_IN_PARTICLE_COUNT`) and the + * PDG ids of the external particles (`UMAMI_IN_PDG_IDS`). It returns one value + * per event, under `UMAMI_OUT_CUT`, `UMAMI_OUT_SCALE` or `UMAMI_OUT_BIAS`. + * + * With a single PDG-id option the ids are constant. With several options the + * ids are selected per event by the `flavor_index` argument. + * + * `batch` is the leading batch dimension. + * + * **Arguments** + * - `momenta` – `float`, shape `(batch, particle_count, 4)` – the event momenta. + * - `flavor_index` – `int`, shape `(batch,)` – index into the PDG-id options. + * Present only with more than one option. + * + * **Returns** + * - `cut`, `scale` or `bias` – `float`, shape `(batch,)` – the function value. + */ +class CustomFunction : public FunctionGenerator { +public: + /// What the custom function computes. + enum CustomFunctionType { + cut, ///< Acceptance factor, 0 rejects the event. + scale, ///< Energy scale. + bias ///< Event bias. + }; + + /** + * @param function_index Index of the loaded library in the context. + * @param type What the function computes. + * @param particle_count Number of external particles. + * @param pdg_id_options PDG ids of the external particles, one list per + * flavor option. Empty passes zeros. + */ + CustomFunction( + std::size_t function_index, + CustomFunctionType type, + std::size_t particle_count, + const nested_vector2& pdg_id_options = {} + ); + /** + * Take the index from an already-loaded library. + * @param function_api The loaded library. + * @param type What the function computes. + * @param particle_count Number of external particles. + * @param pdg_id_options See the other constructor. + */ + CustomFunction( + const MatrixElementApi& function_api, + CustomFunctionType type, + std::size_t particle_count, + const nested_vector2& pdg_id_options = {} + ) : + CustomFunction(function_api.index(), type, particle_count, pdg_id_options) {} + + /// Index of the loaded library in the context. + std::size_t function_index() const { return _function_index; } + /// What the function computes. + CustomFunctionType type() const { return _type; } + /// Whether the PDG ids are selected per event by a `flavor_index` argument. + bool dynamic_pdg_ids() const { return _pdg_id_options.size() > 1; } + +private: + NamedVector build_function_impl( + FunctionBuilder& fb, const NamedVector& args + ) const override; + + std::size_t _function_index; + CustomFunctionType _type; + std::size_t _particle_count; + nested_vector2 _pdg_id_options; +}; + +/** + * A user-supplied parton density, loaded as a UMAMI library. + * + * Passed in place of a @ref PdfGrid. The library receives the momentum + * fraction, the factorization scale, the parton PDG id and the beam index. It + * returns @f$x f(x, q)@f$ under `UMAMI_OUT_PDF`. + */ +struct CustomPdf { + /// @param function_index Index of the loaded library in the context. + CustomPdf(std::size_t function_index) : function_index(function_index) {} + /// @param function_api The loaded library. + CustomPdf(const MatrixElementApi& function_api) : + function_index(function_api.index()) {} + /// Index of the loaded library in the context. + std::size_t function_index; +}; + +} // namespace madspace diff --git a/madspace/include/madspace/phasespace/cuts.hpp b/madspace/include/madspace/phasespace/cuts.hpp index f3ff571ac4..1933855ed9 100644 --- a/madspace/include/madspace/phasespace/cuts.hpp +++ b/madspace/include/madspace/phasespace/cuts.hpp @@ -2,6 +2,7 @@ #include "madspace/compgraphs.hpp" #include "madspace/phasespace/base.hpp" +#include "madspace/phasespace/custom_function.hpp" #include "madspace/phasespace/observable.hpp" #include @@ -17,7 +18,9 @@ namespace madspace { * a list of @ref Cuts::CutItem "cut items" [1]. Each item evaluates one * @ref Observable for its particle selection and requires the result to lie in * `[min, max]`. When the selection yields several values, @ref CutMode decides - * whether all or any of them must pass. See Sec. 3.2.4 of [1]. + * whether all or any of them must pass. See Sec. 3.2.4 of [1]. An optional + * @ref CustomFunction of type CustomFunction::cut multiplies its value into the + * mask. * * The accessors re-expose the active lower bounds in the layout the * phase-space mappings consume. `pt_min()` and `eta_max()` are indexed by @@ -57,11 +60,19 @@ class Cuts : public FunctionGenerator { CutMode mode = CutMode::all; }; - /// @param cut_data The cut items to apply. - Cuts(const std::vector& cut_data); - /// Build a pass-through mask with no cuts. + /// @param cut_data The cut items to apply. + /// @param custom_cut Optional user-supplied cut function. + Cuts( + const std::vector& cut_data, + const std::optional& custom_cut = std::nullopt + ); + /// Build a mask with no cut items. /// @param particle_count Number of external particles. - Cuts(std::size_t particle_count); + /// @param custom_cut Optional user-supplied cut function. + Cuts( + std::size_t particle_count, + const std::optional& custom_cut = std::nullopt + ); /// Largest required partonic center-of-mass energy, or 0 if unconstrained. double sqrt_s_min() const; /// Per-outgoing-particle maximum pseudorapidity (infinity where inactive). @@ -96,6 +107,7 @@ class Cuts : public FunctionGenerator { ) const; std::vector _cut_data; + std::optional _custom_cut; }; } // namespace madspace diff --git a/madspace/include/madspace/phasespace/integrand.hpp b/madspace/include/madspace/phasespace/integrand.hpp index 0ca15b8cac..3da4e6e9c8 100644 --- a/madspace/include/madspace/phasespace/integrand.hpp +++ b/madspace/include/madspace/phasespace/integrand.hpp @@ -3,6 +3,7 @@ #include "madspace/phasespace/channel_weight_network.hpp" #include "madspace/phasespace/channel_weights.hpp" #include "madspace/phasespace/cross_section.hpp" +#include "madspace/phasespace/custom_function.hpp" #include "madspace/phasespace/discrete_flow.hpp" #include "madspace/phasespace/discrete_sampler.hpp" #include "madspace/phasespace/flow.hpp" @@ -40,11 +41,23 @@ 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 no @p prop_chan_weights. * * **References** * - [1] T. Heimel, O. Mattelaer, R. Winterhalder, "MadSpace", @@ -57,6 +70,8 @@ class Integrand : public FunctionGenerator { /// The optional adaptive sampler for a set of discrete choices. using AdaptiveDiscrete = std::variant; + /// The optional parton density: an LHAPDF grid or a user-supplied library. + using PdfSource = std::variant; /// The @ref MatrixElement inputs an `Integrand` supplies. inline static const std::vector @@ -68,7 +83,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, @@ -91,7 +106,8 @@ class Integrand : public FunctionGenerator { * the flavour choice. * @param pid_options Allowed initial-parton flavour * combinations. - * @param pdf_grid Optional shared @ref PdfGrid. + * @param pdf_grid Optional shared @ref PdfGrid or + * @ref CustomPdf. * @param running_coupling Optional shared * @ref RunningCoupling. * @param energy_scale Optional shared @ref EnergyScale. @@ -126,6 +142,10 @@ class Integrand : public FunctionGenerator { * @param flavor_per_subproc_remap Per-sub-process flavour remap. * @param compressed_channel_weight_count Channel-weight entries kept per * event. + * @param custom_bias Optional user-supplied event + * bias of type CustomFunction::bias. + * Its PDG-id options must match + * @p pid_options. */ Integrand( const PhaseSpaceMapping& mapping, @@ -134,7 +154,7 @@ class Integrand : public FunctionGenerator { const AdaptiveDiscrete& discrete_sym = std::monostate{}, const AdaptiveDiscrete& discrete_flavor = std::monostate{}, const nested_vector2& pid_options = {}, - const std::optional& pdf_grid = std::nullopt, + const PdfSource& pdf_grid = std::monostate{}, const std::optional& running_coupling = std::nullopt, const std::optional& energy_scale = std::nullopt, const std::optional& prop_chan_weights = std::nullopt, @@ -155,7 +175,8 @@ class Integrand : public FunctionGenerator { const std::vector& flavor_diff_xs_indices = {}, const std::vector& flavor_subproc_indices = {}, const std::vector& flavor_per_subproc_remap = {}, - std::size_t compressed_channel_weight_count = 50 + std::size_t compressed_channel_weight_count = 50, + const std::optional& custom_bias = std::nullopt ); /// Total number of particles, incoming and outgoing. std::size_t particle_count() const { return _mapping.particle_count(); } @@ -259,6 +280,7 @@ class Integrand : public FunctionGenerator { std::vector _flavor_diff_xs_indices; std::vector _flavor_subproc_indices; std::vector _flavor_per_subproc_remap; + std::optional _custom_bias; friend class IntegrandProbability; friend class IntegrandChannelPart; diff --git a/madspace/include/madspace/phasespace/matrix_element.hpp b/madspace/include/madspace/phasespace/matrix_element.hpp index 70d2c7afc8..8d8a5a72f5 100644 --- a/madspace/include/madspace/phasespace/matrix_element.hpp +++ b/madspace/include/madspace/phasespace/matrix_element.hpp @@ -40,15 +40,21 @@ class MatrixElement : public FunctionGenerator { public: /// A quantity the matrix element can be asked to consume. enum MatrixElementInput { - momenta_in, ///< Final-state momenta. - alpha_s_in, ///< Strong coupling. - flavor_in, ///< Flavour assignment. - random_color_in, ///< Random number for the color draw. - random_helicity_in, ///< Random number for the helicity draw. - random_diagram_in, ///< Random number for the diagram draw. - helicity_in, ///< Fixed helicity. - channel_in, ///< Multi-channel index. - diagram_in ///< Fixed diagram. + momenta_in, ///< Final-state momenta. + alpha_s_in, ///< Strong coupling. + flavor_in, ///< Flavour assignment. + random_color_in, ///< Random number for the color draw. + random_helicity_in, ///< Random number for the helicity draw. + random_diagram_in, ///< Random number for the diagram draw. + helicity_in, ///< Fixed helicity. + channel_in, ///< Multi-channel index. + diagram_in, ///< Fixed diagram. + particle_count_in, ///< Number of external particles (custom functions). + pdg_ids_in, ///< PDG ids of the external particles. + momentum_fraction_in, ///< Momentum fraction of an incoming parton. + scale_in, ///< Energy scale, e.g. the factorization scale. + pdg_id_in, ///< PDG id of a single parton. + beam_index_in ///< Index of the beam (0 or 1). }; /// A quantity the matrix element can be asked to produce. @@ -57,7 +63,14 @@ class MatrixElement : public FunctionGenerator { diagram_amp2_out, ///< Per-diagram squared amplitudes. color_index_out, ///< Sampled color. helicity_index_out, ///< Sampled helicity. - diagram_index_out ///< Sampled diagram. + diagram_index_out, ///< Sampled diagram. + cut_out, ///< Custom cut value, 0 rejects the event. + scale_out, ///< Custom energy scale. + pdf_out, ///< Custom parton density x f(x, Q). + bias_out, ///< Custom event bias. + /// Ratio of the matrix element to a reference matrix element, for a + /// @ref WeightingStep. + matrix_element_ratio_out }; /** @@ -132,6 +145,8 @@ class MatrixElement : public FunctionGenerator { const std::vector& inputs() const { return _inputs; } /// The configured output quantities. const std::vector& outputs() const { return _outputs; } + /// Whether the `random_*` inputs are drawn internally. + bool sample_random_inputs() const { return _sample_random_inputs; } /// The inputs that must be supplied by the caller (excludes the ones drawn /// internally when @p sample_random_inputs is set). std::vector external_inputs() const; diff --git a/madspace/include/madspace/phasespace/pdf.hpp b/madspace/include/madspace/phasespace/pdf.hpp index 916fe357b1..599b9e87c1 100644 --- a/madspace/include/madspace/phasespace/pdf.hpp +++ b/madspace/include/madspace/phasespace/pdf.hpp @@ -2,6 +2,7 @@ #include "madspace/driver/context.hpp" #include "madspace/phasespace/base.hpp" +#include "madspace/phasespace/custom_function.hpp" namespace madspace { @@ -57,7 +58,8 @@ struct PdfGrid { * * Interpolates a @ref PdfGrid at the sampled @f$(x, q)@f$, returning * @f$x f(x, q)@f$ for the requested flavours; the numerical output agrees with - * LHAPDF to double precision [1] (Sec. 3.2.3 of [2]). + * LHAPDF to double precision [1] (Sec. 3.2.3 of [2]). With a @ref CustomPdf, + * the user-supplied library is called once per requested flavour instead. * * `batch` is the leading batch dimension. * @@ -92,6 +94,19 @@ class PartonDensity : public FunctionGenerator { bool dynamic_pid = false, const std::string& prefix = "" ); + /** + * @param custom_pdf The user-supplied parton density. + * @param pids PDG ids of the flavours to return. + * @param beam_index Index of the beam (0 or 1), passed to the library. + * @param dynamic_pid If true, a single flavour selected per event by the + * `flavor_index` argument is returned. + */ + PartonDensity( + const CustomPdf& custom_pdf, + const std::vector& pids, + int beam_index, + bool dynamic_pid = false + ); private: NamedVector build_function_impl( @@ -99,6 +114,9 @@ class PartonDensity : public FunctionGenerator { ) const override; std::vector _pid_indices; + std::vector _pids; + std::optional _custom_pdf; + me_int_t _beam_index = 0; bool _dynamic_pid; std::string _prefix; std::vector _logx_shape; diff --git a/madspace/include/madspace/phasespace/scale.hpp b/madspace/include/madspace/phasespace/scale.hpp index d25fe2279b..8f8abab21c 100644 --- a/madspace/include/madspace/phasespace/scale.hpp +++ b/madspace/include/madspace/phasespace/scale.hpp @@ -1,6 +1,7 @@ #pragma once #include "madspace/phasespace/base.hpp" +#include "madspace/phasespace/custom_function.hpp" namespace madspace { @@ -24,6 +25,8 @@ namespace madspace { * @f$H_\mathrm{T}/2@f$ (`half_transverse_mass`) or @f$\sqrt{\hat s}@f$ * (`partonic_energy`). The dynamical scale is multiplied by `scale_factor`, * e.g. @f$H_\mathrm{T}/4@f$ is `half_transverse_mass` with `scale_factor = 0.5`. + * A @ref CustomFunction of type CustomFunction::scale replaces the dynamical + * choice, like `user_dynamical_scale` in MadGraph5. * * `batch` is the leading batch dimension. * @@ -85,6 +88,8 @@ class EnergyScale : public FunctionGenerator { * @param fact_scale2 Fixed factorization scale for the second beam. * @param scale_factor Factor on the dynamical scale; fixed scales are * left alone. + * @param custom_scale Optional user-supplied function that replaces + * the dynamical scale choice. */ EnergyScale( std::size_t particle_count, @@ -94,7 +99,8 @@ class EnergyScale : public FunctionGenerator { double ren_scale, double fact_scale1, double fact_scale2, - double scale_factor = 1. + double scale_factor = 1., + const std::optional& custom_scale = std::nullopt ); private: @@ -109,6 +115,7 @@ class EnergyScale : public FunctionGenerator { double _fact_scale1; double _fact_scale2; double _scale_factor; + std::optional _custom_scale; }; } // namespace madspace diff --git a/madspace/include/madspace/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 57d079d1c1..4eb2dba380 100644 --- a/madspace/include/madspace/umami.h +++ b/madspace/include/madspace/umami.h @@ -30,7 +30,7 @@ extern "C" { * Minor version number of the UMAMI interface. Between minor versions, new keys for * errors, devices, metadata, inputs and outputs can be added. */ -#define UMAMI_MINOR_VERSION 0 +#define UMAMI_MINOR_VERSION 2 typedef enum { /** operation was executed successfully */ @@ -93,15 +93,27 @@ typedef enum { UMAMI_IN_DIAGRAM_INDEX, /** externally selected channel index, type: `unsigned int`, shape: `()` */ UMAMI_IN_CHANNEL_INDEX, + /** number of external legs, type: `int`, shape: `()` */ + UMAMI_IN_PARTICLE_COUNT, + /** PDG ids of the external legs, type: `int`, shape: `(particle count)` */ + UMAMI_IN_PDG_IDS, + /** momentum fraction of an incoming parton, type: `double`, shape: `()` */ + UMAMI_IN_MOMENTUM_FRACTION, + /** energy scale, e.g. the factorization scale, type: `double`, shape: `()` */ + UMAMI_IN_SCALE, + /** PDG id of a single parton, type: `int`, shape: `()` */ + UMAMI_IN_PDG_ID, + /** index of the beam (0 or 1), type: `int`, shape: `()` */ + UMAMI_IN_BEAM_INDEX, } UmamiInputKey; /** Number of values in `UmamiInputKey` */ -#define UMAMI_INPUT_KEY_COUNT (UMAMI_IN_CHANNEL_INDEX + 1) +#define UMAMI_INPUT_KEY_COUNT (UMAMI_IN_BEAM_INDEX + 1) typedef enum { /** value of the matrix element, type: `double`, shape: `()` */ UMAMI_OUT_MATRIX_ELEMENT, - /** selected color index, type: `double`, shape: `(diagram count)` */ + /** squared amplitudes of the diagrams, type: `double`, shape: `(diagram count)` */ UMAMI_OUT_DIAGRAM_AMP2, /** selected color index, type: `int`, shape: `()` */ UMAMI_OUT_COLOR_INDEX, @@ -112,10 +124,21 @@ typedef enum { /** CUDA or HIP stream for asynchronous execution. Listed as an output as it is a * mutable pointer */ UMAMI_OUT_GPU_STREAM, + /** value of a custom cut, 0 rejects the event, type: `double`, shape: `()` */ + UMAMI_OUT_CUT, + /** value of a custom energy scale, type: `double`, shape: `()` */ + UMAMI_OUT_SCALE, + /** value x f(x, Q) of a custom parton density, type: `double`, shape: `()` */ + UMAMI_OUT_PDF, + /** value of a custom event bias, type: `double`, shape: `()` */ + UMAMI_OUT_BIAS, + /** 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_GPU_STREAM + 1) +#define UMAMI_OUTPUT_KEY_COUNT (UMAMI_OUT_MATRIX_ELEMENT_RATIO + 1) /** Implementation-defined pointer to a matrix element instance */ typedef void* UmamiHandle; diff --git a/madspace/madspace/_madspace_py_loader.py b/madspace/madspace/_madspace_py_loader.py index 5c8744b981..b68cff15f7 100644 --- a/madspace/madspace/_madspace_py_loader.py +++ b/madspace/madspace/_madspace_py_loader.py @@ -16,6 +16,7 @@ ) from ._madspace_py import * +from .python_function import PythonFunction @contextlib.contextmanager diff --git a/madspace/madspace/python_function.py b/madspace/madspace/python_function.py new file mode 100644 index 0000000000..3202613ea3 --- /dev/null +++ b/madspace/madspace/python_function.py @@ -0,0 +1,248 @@ +import contextlib +import ctypes + +from ._madspace_py import ( + _UMAMI_INPUT_KEYS, + _UMAMI_OUTPUT_KEYS, + Context, + DataType, + _tensor_view, +) + +# input name: (data type, shape after the batch dimension); "n" is the particle count +_INPUT_SHAPES = { + "momenta": (DataType.float, ("n", 4)), + "alpha_s": (DataType.float, ()), + "flavor_index": (DataType.int, ()), + "random_color": (DataType.float, ()), + "random_helicity": (DataType.float, ()), + "random_diagram": (DataType.float, ()), + "helicity_index": (DataType.int, ()), + "diagram_index": (DataType.int, ()), + "channel_index": (DataType.int, ()), + "particle_count": (DataType.int, ()), + "pdg_ids": (DataType.int, ("n",)), + "momentum_fraction": (DataType.float, ()), + "scale": (DataType.float, ()), + "pdg_id": (DataType.int, ()), + "beam_index": (DataType.int, ()), +} +_OUTPUT_TYPES = { + "matrix_element": DataType.float, + "color_index": DataType.int, + "helicity_index": DataType.int, + "diagram_index": DataType.int, + "cut": DataType.float, + "scale": DataType.float, + "pdf": DataType.float, + "bias": DataType.float, + "matrix_element_ratio": DataType.float, +} + + +class PythonFunction: + """ + A UMAMI implementation that calls a Python function on numpy arrays or torch + tensors, for example a custom cut. Add it to a context with + ``Context.add_matrix_element``, then use the returned ``MatrixElementApi`` like a + loaded library. + + ``function`` is called once per batch, with one array per entry of ``inputs``, + in that order. The arrays have the batch size as first dimension: + ``momenta`` has shape ``(batch, n, 4)``, ``pdg_ids`` has shape ``(batch, n)``, + and all other inputs have shape ``(batch,)``. Floats are 64-bit, integers 32-bit. + The function returns one array of shape ``(batch,)`` per entry of ``outputs``, a + single array for one output and a tuple for several. + + The arrays are views of madspace's buffers on the device of the context, so they + must not be kept after the call. numpy functions only run on the CPU, torch + functions also on CUDA and HIP GPUs, on the stream of the madspace runtime. + + Args: + function: The Python function. + inputs: Names of the UMAMI inputs passed to ``function``. + outputs: Names of the UMAMI outputs ``function`` returns. + framework: ``"numpy"`` or ``"torch"``. + particle_count: The particle count ``n`` for ``momenta`` and ``pdg_ids``, if + the caller does not pass the ``particle_count`` input. + name: Name used in error messages, by default the function's name. + """ + + def __init__( + self, + function, + inputs, + outputs, + framework="numpy", + particle_count=None, + name=None, + ): + if framework not in ("numpy", "torch"): + raise ValueError(f"framework must be 'numpy' or 'torch', not {framework!r}") + for key in inputs: + if key not in _INPUT_SHAPES: + raise ValueError( + f"unknown input {key!r}, expected one of {', '.join(_INPUT_SHAPES)}" + ) + for key in outputs: + if key not in _OUTPUT_TYPES: + raise ValueError( + f"unknown output {key!r}, expected one of {', '.join(_OUTPUT_TYPES)}" + ) + if len(outputs) == 0: + raise ValueError("at least one output is needed") + self.function = function + self.inputs = list(inputs) + self.outputs = list(outputs) + self.framework = framework + self.particle_count = particle_count + self.name = name or getattr(function, "__name__", "python_function") + self._input_keys = [_UMAMI_INPUT_KEYS[key] for key in self.inputs] + self._output_keys = [_UMAMI_OUTPUT_KEYS[key] for key in self.outputs] + self._needs_particle_count = any( + "n" in _INPUT_SHAPES[key][1] for key in self.inputs + ) + if self._needs_particle_count and particle_count is None: + # read from the particle_count input, even if the function does not take it + self._input_keys.append(_UMAMI_INPUT_KEYS["particle_count"]) + if framework == "torch": + import torch + + self._torch = torch + self._from_dlpack = torch.from_dlpack + else: + import numpy + + self._numpy = numpy + self._from_dlpack = numpy.from_dlpack + + def _view(self, address, dtype, shape, stride, offset, device): + return self._from_dlpack( + _tensor_view(address, dtype, shape, stride, offset, device) + ) + + def _numpy_output(self, address, dtype, count, offset): + # numpy.from_dlpack gives read-only arrays, so map the CPU buffer directly + ctype = ctypes.c_double if dtype == DataType.float else ctypes.c_int32 + size = ctypes.sizeof(ctype) + buffer = (ctype * count).from_address(address + offset * size) + return self._numpy.frombuffer(buffer, dtype=self._numpy.dtype(ctype)) + + def _stream(self, device, stream): + if self.framework == "torch" and device != "cpu" and stream is not None: + torch = self._torch + return torch.cuda.stream(torch.cuda.ExternalStream(stream)) + return contextlib.nullcontext() + + def _call(self, device, count, stride, offset, inputs, outputs, stream): + """Called by the UMAMI implementation with the raw buffer addresses.""" + with self._stream(device, stream): + n = self.particle_count + if self._needs_particle_count and n is None: + counts = self._view( + inputs[_UMAMI_INPUT_KEYS["particle_count"]], + DataType.int, + [count], + [1], + offset, + device, + ) + n = int(counts[0]) + args = [] + for key in self.inputs: + dtype, shape = _INPUT_SHAPES[key] + shape = [count] + [n if size == "n" else size for size in shape] + # the batch dimension is contiguous, the others are strided by it + strides = [1] + [ + stride * _prod(shape[1:i]) for i in range(1, len(shape)) + ] + args.append( + self._view( + inputs[_UMAMI_INPUT_KEYS[key]], + dtype, + shape, + strides, + offset, + device, + ) + ) + results = self.function(*args) + if len(self.outputs) == 1: + results = (results,) + elif not isinstance(results, (tuple, list)) or len(results) != len( + self.outputs + ): + raise TypeError( + f"{self.name} must return a tuple of {len(self.outputs)} arrays " + f"({', '.join(self.outputs)})" + ) + for key, result in zip(self.outputs, results): + address = outputs.get(_UMAMI_OUTPUT_KEYS[key]) + if address is None: + continue + dtype = _OUTPUT_TYPES[key] + self._check(key, result, dtype, count, device) + if self.framework == "torch": + self._view(address, dtype, [count], [1], offset, device).copy_( + result + ) + else: + self._numpy_output(address, dtype, count, offset)[:] = result + + def _check(self, key, result, dtype, count, device): + prefix = f"output {key!r} of {self.name}" + if self.framework == "torch": + torch = self._torch + if not isinstance(result, torch.Tensor): + raise TypeError( + f"{prefix} must be a torch.Tensor, not {type(result).__name__}" + ) + expected = torch.float64 if dtype == DataType.float else torch.int32 + result_device = "cpu" if result.device.type == "cpu" else "gpu" + if result_device != ("cpu" if device == "cpu" else "gpu"): + raise ValueError( + f"{prefix} must be on the {device} device, not {result.device}" + ) + else: + numpy = self._numpy + if not isinstance(result, numpy.ndarray): + raise TypeError( + f"{prefix} must be a numpy.ndarray, not {type(result).__name__}" + ) + expected = numpy.dtype( + numpy.float64 if dtype == DataType.float else numpy.int32 + ) + if result.dtype != expected: + raise TypeError(f"{prefix} must have dtype {expected}, not {result.dtype}") + if tuple(result.shape) != (count,): + raise ValueError( + f"{prefix} must have shape ({count},), not {tuple(result.shape)}" + ) + + +def _prod(values): + result = 1 + for value in values: + result *= value + return result + + +def _add_matrix_element(self, function, name=None): + """ + Add a matrix element implemented in Python, see ``PythonFunction``. Returns its + ``MatrixElementApi``. + """ + if not isinstance(function, PythonFunction): + raise TypeError("add_matrix_element takes a madspace.PythonFunction") + required = list(function._input_keys) + return self._add_python_function( + function, + function._input_keys, + required, + function._output_keys, + function.framework, + name or function.name, + ) + + +Context.add_matrix_element = _add_matrix_element diff --git a/madspace/src/compgraphs/instruction.cpp b/madspace/src/compgraphs/instruction.cpp index b05d0486ae..39e165a5c1 100644 --- a/madspace/src/compgraphs/instruction.cpp +++ b/madspace/src/compgraphs/instruction.cpp @@ -939,6 +939,8 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const { case UMAMI_IN_RANDOM_COLOR: case UMAMI_IN_RANDOM_HELICITY: case UMAMI_IN_RANDOM_DIAGRAM: + case UMAMI_IN_MOMENTUM_FRACTION: + case UMAMI_IN_SCALE: if (input_type.dtype != DataType::dt_float || input_type.shape.size() != 0) { throw std::invalid_argument( @@ -952,6 +954,9 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const { case UMAMI_IN_HELICITY_INDEX: case UMAMI_IN_DIAGRAM_INDEX: case UMAMI_IN_CHANNEL_INDEX: + case UMAMI_IN_PARTICLE_COUNT: + case UMAMI_IN_PDG_ID: + case UMAMI_IN_BEAM_INDEX: if (input_type.dtype != DataType::dt_int || input_type.shape.size() != 0) { throw std::invalid_argument( std::format( @@ -960,6 +965,15 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const { ); } break; + case UMAMI_IN_PDG_IDS: + if (input_type.dtype != DataType::dt_int || input_type.shape.size() != 1) { + throw std::invalid_argument( + std::format( + "matrix_element, argument {}: expected array of integers", i + 1 + ) + ); + } + break; default: throw std::invalid_argument( std::format( @@ -979,6 +993,11 @@ TypeVec MatrixElementInstruction::signature(const ValueVec& args) const { std::vector shape; switch (output_key) { case UMAMI_OUT_MATRIX_ELEMENT: + case UMAMI_OUT_CUT: + case UMAMI_OUT_SCALE: + case UMAMI_OUT_PDF: + case UMAMI_OUT_BIAS: + case UMAMI_OUT_MATRIX_ELEMENT_RATIO: dtype = DataType::dt_float; shape = {}; break; diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp index 9b4a717926..9d24f77720 100644 --- a/madspace/src/cpu/runtime.cpp +++ b/madspace/src/cpu/runtime.cpp @@ -87,6 +87,7 @@ void op_matrix_element( // implemented at LUMI hackathon where the coffee was indeed terrible switch (output_keys[i]) { case UMAMI_OUT_MATRIX_ELEMENT: + case UMAMI_OUT_MATRIX_ELEMENT_RATIO: device.submit([=]() mutable { double* ptr = static_cast(output_ptrs[i]); std::fill(ptr, ptr + batch_size, 1.); diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 86f471c5a2..b96e37595e 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; @@ -21,6 +22,9 @@ int event_extra_flags(const std::unordered_map& index_ if (index_map.contains("subprocess_index")) { flags |= EventRecord::f_subproc_index; } + if (index_map.contains("bias")) { + flags |= EventRecord::f_bias; + } return flags; } @@ -72,6 +76,18 @@ DerivedSeed generate_phase_seed( ); } +// An empty path means the events are only kept in memory. +std::optional open_event_file( + const std::string& path, const DataLayout& layout, std::size_t particle_count +) { + if (path.empty()) { + return std::nullopt; + } + return std::make_optional( + path, layout, particle_count, EventFile::create, true + ); +} + } // namespace ChannelEventGenerator::ChannelEventGenerator( @@ -114,13 +130,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()), @@ -146,6 +159,17 @@ ChannelEventGenerator::ChannelEventGenerator( init_used_globals(); init_runtimes(); init_field_indices(); + auto& me_outputs = integrand.diff_xs().at(0).matrix_element().outputs(); + for (auto [provided, output] : + zip(_provided_indices, + std::array{ + MatrixElement::color_index_out, + MatrixElement::helicity_index_out, + MatrixElement::diagram_index_out + })) { + provided = + std::find(me_outputs.begin(), me_outputs.end(), output) != me_outputs.end(); + } if (const auto& grid_name = integrand.vegas_grid_name(); grid_name) { _vegas_optimizer = VegasGridOptimizer(contexts, grid_name.value(), config.vegas_damping); @@ -241,10 +265,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), @@ -290,6 +313,13 @@ void ChannelEventGenerator::init_runtimes() { } void ChannelEventGenerator::init_field_indices() { + auto& unweighter_outputs = _unweighter_function.outputs(); + for (auto [key, value] : zip(unweighter_outputs.keys(), unweighter_outputs)) { + // with the common batch size, to compare and combine channels + _event_types.push_back( + key, Type(value.type.dtype, madspace::batch_size, value.type.shape) + ); + } auto index_map = _integrand_common_function.outputs().index_map(); _field_indices.weight = index_map.at("weight"); _field_indices.momenta = index_map.at("momenta"); @@ -323,18 +353,39 @@ void ChannelEventGenerator::init_field_indices() { } else { _field_indices.subprocess_index = -1; } + _field_indices.bias = index_map.contains("bias") ? index_map.at("bias") : -1; + _field_indices.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()) { @@ -346,10 +397,10 @@ void ChannelEventGenerator::unweight_file(MixMaxRandom& rand_gen) { ++accept_count; } } - _weight_file.seek(i); - _weight_file.write(buffer); + w_file.seek(i); + w_file.write(buffer); } - _unweighted_count = _weight_file.event_count(); + _unweighted_count = w_file.event_count(); _unweighted_accept_count = accept_count; _status.count_unweighted = accept_count; } @@ -361,14 +412,17 @@ void ChannelEventGenerator::integrate(const GeneratorBatchJob& job) { if (w_view[i] != 0) { ++sample_count_after_cuts; } - _cross_section.push(w_view[i]); + // With a bias b the weights are w b: the absolute integral, which drives + // the adaptation, unweighting and channel allocation, keeps the bias, + // while the reported cross section divides it out. + double weight = w_view[i]; + if (_field_indices.bias != -1 && weight != 0.) { + weight /= job.biases.view()[i]; + } + _cross_section.push(weight); _abs_cross_section.push(std::abs(w_view[i])); } - _status.mean = _cross_section.mean(); - _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; @@ -414,17 +468,19 @@ void ChannelEventGenerator::optimize_vegas(const GeneratorBatchJob& job) { } } _best_rsd = std::min(rsd, _best_rsd); + _status.iters_without_improvement = _iters_without_improvement; ++_status.iterations; } 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) { @@ -435,7 +491,9 @@ double ChannelEventGenerator::channel_weight_sum(std::size_t event_count) { if (weight == 0.) { continue; } - weight_sum += std::abs(weight) / _max_weight; + weight_sum += _step_settings.empty() + ? std::abs(weight) / _max_weight + : std::abs(unweighted_weight(buffer, j)); ++unweighted_count; } if (done) { @@ -556,6 +614,9 @@ void ChannelEventGenerator::start_job( throw std::runtime_error("job cancelled"); } job.weights = job.events.at(_field_indices.weight).cpu(); + if (_field_indices.bias != -1) { + job.biases = job.events.at(_field_indices.bias).cpu(); + } // observable_histograms/vegas_histogram/discrete_histogram don't consume // random numbers, so they're never seeded. if (runtimes.observable_histograms) { @@ -594,6 +655,7 @@ void ChannelEventGenerator::start_job( void ChannelEventGenerator::prepare_unweight_job(GeneratorBatchJob& job) const { job.max_weight = _max_weight; + job.count_target = _status.count_target; } void ChannelEventGenerator::submit_unweight_job( @@ -616,6 +678,30 @@ void ChannelEventGenerator::submit_unweight_job( ); unweighter_args.push_back(Tensor(job.max_weight, context->device())); TensorVec unw_events = runtimes.unweighter->run(unweighter_args); + if (!_step_settings.empty()) { + // the events stay on the device for the steps, with the random + // numbers of the steps drawn per event + if (job.rng_seed) { + runtimes.step_random->set_seed(DerivedSeed( + job.rng_seed, + DerivedSeed::weighting_step_random, + job.rng_job_index, + job.channel_index, + job.rng_is_survey ? 1 + job.rng_survey_pass : 0 + )); + } + auto random = runtimes.step_random->run({unw_events.at(0)}); + job.step_host.push_back(unw_events.at(_field_indices.weight).cpu()); + job.step_host.push_back( + unw_events.at(_field_indices.matrix_element).cpu() + ); + if (_field_indices.bias != -1) { + job.step_host.push_back(unw_events.at(_field_indices.bias).cpu()); + } + job.unweighted_events = std::move(unw_events); + job.unweighted_events.push_back(random.at(0)); + return; + } for (auto& item : unw_events) { job.unweighted_events.push_back(item.cpu()); } @@ -638,17 +724,33 @@ std::size_t ChannelEventGenerator::next_vegas_batch_size() { void ChannelEventGenerator::clear_events() { _status.count_unweighted = 0; - _max_weight = 0; + if (!_max_weight_fixed) { + _max_weight = 0; + } _unweighted_count = 0; _unweighted_accept_count = 0; _status.count_opt = 0; _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.; + _final_integral = {}; + ++_event_epoch; + for (auto& state : _step_states) { + if (!state.max_weight_fixed) { + state.max_weight = 0.; + } + state.large_weights.clear(); + state.count_accepted = 0.; + } for (auto& hist : _histograms) { std::fill(hist.bin_values.begin(), hist.bin_values.end(), 0.); std::fill(hist.bin_errors.begin(), hist.bin_errors.end(), 0.); @@ -656,7 +758,10 @@ void ChannelEventGenerator::clear_events() { } void ChannelEventGenerator::update_max_weight(Tensor weights) { - if (_status.count_unweighted > _config.freeze_max_weight_after) { + // with steps, the first step freezes on its own accepted events + double unweighted_count = + _step_settings.empty() ? _status.count_unweighted : _first_step_count; + if (_max_weight_fixed || unweighted_count > _config.freeze_max_weight_after) { return; } @@ -686,8 +791,10 @@ void ChannelEventGenerator::set_target_count(std::size_t target_count) { // (e.g. generate()'s round-one seed) would otherwise never be corrected: a // channel that has met its smaller target is not revisited. Tighten the cap // now; unweight_all() re-unweights the whole file against it at the end. - if (target_count < old_target && _max_weight != 0 && - _status.count_unweighted <= _config.freeze_max_weight_after) { + double unweighted_count = + _step_settings.empty() ? _status.count_unweighted : _first_step_count; + if (target_count < old_target && _max_weight != 0 && !_max_weight_fixed && + unweighted_count <= _config.freeze_max_weight_after) { apply_truncation_budget(); } } @@ -708,6 +815,10 @@ void ChannelEventGenerator::apply_truncation_budget() { if (w_sum / w > max_truncation) { if (_max_weight < w) { _status.count_unweighted *= _max_weight / w; + _first_step_count *= _max_weight / w; + for (auto& state : _step_states) { + state.count_accepted *= _max_weight / w; + } _max_weight = w; _unweighted_count = 0; _unweighted_accept_count = 0; @@ -721,6 +832,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 = @@ -735,10 +861,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); @@ -791,6 +915,13 @@ void ChannelEventGenerator::write_events( } } + if (_field_indices.bias != -1) { + auto bias_view = unweighted_events.at(_field_indices.bias).view(); + for (std::size_t i = 0; i < bias_view.size(); ++i) { + event_buffer.event(i).bias() = bias_view[i]; + } + } + if (_field_indices.subprocess_index != -1) { auto subproc_view = unweighted_events.at(_field_indices.subprocess_index).view(); @@ -799,19 +930,32 @@ 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 { std::ofstream f(file_name); + f << to_json(); +} + +std::string ChannelEventGenerator::to_json(bool include_estimates) const { json j; j = *this; - f << j.dump(); + if (include_estimates) { + j["max_weight"] = _max_weight; + auto integral = result_integral(); + double error = std::sqrt(integral.variance_abs / integral.count); + if (integral.count > 1 && std::isfinite(error) && error > 0) { + j["abs_integral"] = {{"mean", integral.mean_abs}, {"error", error}}; + } + if (!_step_states.empty()) { + json step_max_weights = json::array(); + for (auto& state : _step_states) { + step_max_weights.push_back(state.max_weight); + } + j["step_max_weights"] = step_max_weights; + } + } + return j.dump(); } ChannelEventGenerator ChannelEventGenerator::load( @@ -822,7 +966,23 @@ ChannelEventGenerator ChannelEventGenerator::load( const GeneratorConfig& config ) { std::ifstream f(channel_file); - json channel = json::parse(f); + if (!f) { + throw std::runtime_error(std::format("Could not open file '{}'", channel_file)); + } + std::string text( + (std::istreambuf_iterator(f)), std::istreambuf_iterator() + ); + return load_json(text, contexts, event_file, weight_file, config); +} + +ChannelEventGenerator ChannelEventGenerator::load_json( + const std::string& channel_json, + const std::vector& contexts, + const std::string& event_file, + const std::string& weight_file, + const GeneratorConfig& config +) { + json channel = json::parse(channel_json); std::optional hist_function; std::vector histograms; if (!channel.at("histogram_function").is_null()) { @@ -838,7 +998,7 @@ ChannelEventGenerator ChannelEventGenerator::load( } } - return ChannelEventGenerator( + ChannelEventGenerator generator( contexts, channel.at("particle_count").get(), channel.at("integrand_channel_function").get(), @@ -853,6 +1013,441 @@ ChannelEventGenerator ChannelEventGenerator::load( config, histograms ); + if (channel.contains("max_weight") && channel.at("max_weight").get() > 0) { + generator.set_fixed_max_weight(channel.at("max_weight").get()); + } + if (channel.contains("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(), + channel.at("abs_integral").at("error").get() + ); + } + return generator; +} + +void ChannelEventGenerator::set_fixed_max_weight(double max_weight) { + if (!(max_weight > 0)) { + throw std::invalid_argument("The fixed maximum weight must be positive"); + } + _max_weight = max_weight; + _max_weight_fixed = true; + _unweighted_count = 0; + _unweighted_accept_count = 0; +} + +void ChannelEventGenerator::set_abs_integral_prior(double mean, double error) { + if (!(mean > 0) || !(error > 0) || !std::isfinite(mean) || !std::isfinite(error)) { + throw std::invalid_argument("The integral prior must be positive and finite"); + } + _abs_integral_prior = {mean, error}; +} + +double ChannelEventGenerator::abs_integral_estimate(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 = std::sqrt(integral.variance_abs / integral.count); + // too few samples for a variance + if (integral.count < 2 || !std::isfinite(error) || !(error > 0)) { + return prior_mean; + } + double prior_weight = 1. / (prior_error * prior_error); + double weight = 1. / (error * error); + return (prior_mean * prior_weight + mean * weight) / (prior_weight + weight); +} + +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_fixed_step_max_weight( + std::size_t step, double max_weight +) { + if (!(max_weight > 0)) { + throw std::invalid_argument("The fixed maximum weight must be positive"); + } + auto& state = _step_states.at(step - 1); + state.max_weight = max_weight; + state.max_weight_fixed = true; + _unweighted_count = 0; + _unweighted_accept_count = 0; +} + +void ChannelEventGenerator::update_step_max_weight( + std::size_t step, const std::vector& values, std::size_t count_target +) { + auto& state = _step_states.at(step - 1); + auto& settings = _step_settings.at(step - 1); + if (state.max_weight_fixed || + state.count_accepted > settings.freeze_max_weight_after) { + return; + } + double min_nonzero = 0.; + for (double value : values) { + double abs_value = std::abs(value); + if (abs_value != 0 && (min_nonzero == 0 || abs_value < min_nonzero)) { + min_nonzero = abs_value; + } + if (abs_value > state.max_weight) { + state.large_weights.push_back(abs_value); + } + } + if (state.max_weight == 0) { + state.max_weight = min_nonzero; + } + auto& large_weights = state.large_weights; + std::sort(large_weights.begin(), large_weights.end(), std::greater{}); + // same truncation as for the first step, see apply_truncation_budget() + double max_truncation = settings.max_overweight_truncation * + static_cast(std::min(count_target, settings.freeze_max_weight_after)); + double sum = 0; + std::size_t count = 0; + for (double value : large_weights) { + if (value < state.max_weight) { + break; + } + sum += value; + ++count; + if (sum / value > max_truncation) { + if (state.max_weight < value) { + // the events of this and the later steps + double factor = state.max_weight / value; + _status.count_unweighted *= factor; + for (std::size_t i = step - 1; i < _step_states.size(); ++i) { + _step_states.at(i).count_accepted *= factor; + } + state.max_weight = value; + _unweighted_count = 0; + _unweighted_accept_count = 0; + } + 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) { @@ -881,5 +1476,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/context.cpp b/madspace/src/driver/context.cpp index f988b8f731..398fe70cc5 100644 --- a/madspace/src/driver/context.cpp +++ b/madspace/src/driver/context.cpp @@ -1,5 +1,6 @@ #include "madspace/driver/context.hpp" +#include #include #include #include @@ -12,7 +13,18 @@ using namespace madspace; using json = nlohmann::json; namespace { -UmamiStatus umami_key_query_not_implemented(bool const**, int*) { return UMAMI_ERROR_NOT_IMPLEMENTED; } +UmamiStatus umami_key_query_not_implemented(bool const**, int*) { + return UMAMI_ERROR_NOT_IMPLEMENTED; +} + +// Fills `function` with the symbol `name` of `lib`, if it exists. +template +void resolve_symbol(void* lib, const char* name, std::function& function) { + void* address = dlsym(lib, name); + if (address != nullptr) { + function = reinterpret_cast(address); + } +} thread_local std::optional current_caller_stream; } // namespace @@ -34,97 +46,84 @@ MatrixElementApi::MatrixElementApi( ); } - _get_meta = reinterpret_cast( - dlsym(_shared_lib.get(), "umami_get_meta") - ); - if (_get_meta == nullptr) { - throw std::runtime_error( - std::format("Did not find symbol umami_get_meta in shared object {}", file) - ); - } + void* lib = _shared_lib.get(); + resolve_symbol(lib, "umami_get_meta", _functions.get_meta); + resolve_symbol(lib, "umami_supported_inputs", _functions.supported_inputs); + resolve_symbol(lib, "umami_required_inputs", _functions.required_inputs); + resolve_symbol(lib, "umami_supported_outputs", _functions.supported_outputs); + resolve_symbol(lib, "umami_initialize", _functions.initialize); + resolve_symbol(lib, "umami_set_parameter", _functions.set_parameter); + resolve_symbol(lib, "umami_matrix_element", _functions.matrix_element); + resolve_symbol(lib, "umami_free", _functions.free); + init_functions(parameters); + init_instances(thread_pool, device, param_card, parameters); +} - // These symbols are optional: implementations may not report which keys they - // support, in which case queries for that information fail with - // UMAMI_ERROR_NOT_IMPLEMENTED instead of the constructor throwing. - _supported_inputs = reinterpret_cast( - dlsym(_shared_lib.get(), "umami_supported_inputs") - ); - if (_supported_inputs == nullptr) { - _supported_inputs = umami_key_query_not_implemented; - } +MatrixElementApi::MatrixElementApi( + UmamiFunctions functions, + const std::string& name, + const std::string& param_card, + ThreadPool& thread_pool, + DevicePtr device, + std::size_t index, + const std::unordered_map& parameters +) : + _functions(std::move(functions)), _file_name(name), _index(index) { + init_functions(parameters); + init_instances(thread_pool, device, param_card, parameters); +} - _required_inputs = reinterpret_cast( - dlsym(_shared_lib.get(), "umami_required_inputs") - ); - if (_required_inputs == nullptr) { - _required_inputs = umami_key_query_not_implemented; +void MatrixElementApi::init_functions( + const std::unordered_map& parameters +) { + auto require = [&](bool present, const char* name) { + if (!present) { + throw_error(std::format("missing UMAMI function {}", name)); + } + }; + require(bool(_functions.get_meta), "umami_get_meta"); + require(bool(_functions.initialize), "umami_initialize"); + require(bool(_functions.matrix_element), "umami_matrix_element"); + require(bool(_functions.free), "umami_free"); + // Implementations may not report which keys they support; queries for that + // information then fail with UMAMI_ERROR_NOT_IMPLEMENTED. + if (!_functions.supported_inputs) { + _functions.supported_inputs = umami_key_query_not_implemented; } - - _supported_outputs = reinterpret_cast( - dlsym(_shared_lib.get(), "umami_supported_outputs") - ); - if (_supported_outputs == nullptr) { - _supported_outputs = umami_key_query_not_implemented; + if (!_functions.required_inputs) { + _functions.required_inputs = umami_key_query_not_implemented; } - - _initialize = reinterpret_cast( - dlsym(_shared_lib.get(), "umami_initialize") - ); - if (_initialize == nullptr) { - throw std::runtime_error( - std::format( - "Did not find symbol umami_initialize in shared object {}", file - ) - ); + if (!_functions.supported_outputs) { + _functions.supported_outputs = umami_key_query_not_implemented; } - - // Optional too, unless parameters have to be set (checked below). - _set_parameter = reinterpret_cast( - dlsym(_shared_lib.get(), "umami_set_parameter") - ); - if (!parameters.empty() && _set_parameter == nullptr) { - throw std::runtime_error( - std::format( - "Did not find symbol umami_set_parameter in shared object {}", file - ) - ); - } - - _matrix_element = reinterpret_cast( - dlsym(_shared_lib.get(), "umami_matrix_element") - ); - if (_matrix_element == nullptr) { - throw std::runtime_error( - std::format( - "Did not find symbol umami_matrix_element in shared object {}", file - ) - ); - } - - _free = - reinterpret_cast(dlsym(_shared_lib.get(), "umami_free")); - if (_free == nullptr) { - throw std::runtime_error( - std::format("Did not find symbol umami_free in shared object {}", file) - ); + if (!parameters.empty()) { + require(bool(_functions.set_parameter), "umami_set_parameter"); } +} +void MatrixElementApi::init_instances( + ThreadPool& thread_pool, + DevicePtr device, + const std::string& param_card, + const std::unordered_map& parameters +) { _instances = ThreadResource( thread_pool, [this, device, param_card, parameters] { device->activate(); void* instance; - check_umami_status(_initialize(&instance, param_card.c_str())); + check_umami_status(_functions.initialize(&instance, param_card.c_str())); for (const auto& [name, value] : parameters) { - UmamiStatus status = _set_parameter(instance, name.c_str(), value, 0.); + UmamiStatus status = + _functions.set_parameter(instance, name.c_str(), value, 0.); if (status != UMAMI_SUCCESS) { - _free(instance); + _functions.free(instance); throw_error(std::format("could not set parameter {}", name)); } } return InstanceType(instance, [this, device](void* proc) { device->activate(); - _free(proc); + _functions.free(proc); }); }, std::nullopt, @@ -160,8 +159,7 @@ void MatrixElementApi::throw_error(const std::string& message) const { ); } -const MatrixElementApi& -Context::load_matrix_element( +const MatrixElementApi& Context::load_matrix_element( const std::string& file, const std::string& param_card, const std::unordered_map& parameters @@ -180,6 +178,27 @@ Context::load_matrix_element( return *_matrix_elements.back().get(); } +const MatrixElementApi& Context::add_matrix_element( + UmamiFunctions functions, + const std::string& name, + const std::string& param_card, + const std::unordered_map& parameters +) { + _param_card_paths.push_back(param_card); + _matrix_elements.push_back( + std::unique_ptr(new MatrixElementApi( + std::move(functions), + name, + param_card, + *_thread_pool, + _device, + _matrix_elements.size(), + parameters + )) + ); + return *_matrix_elements.back().get(); +} + Tensor Context::define_global( const std::string& name, DataType dtype, const SizeVec& shape, bool requires_grad ) { @@ -310,27 +329,31 @@ const MatrixElementApi& Context::matrix_element(std::size_t index) const { return *_matrix_elements.at(index).get(); } -void Context::save_globals(const std::string& dir) const { - namespace fs = std::filesystem; - fs::path dir_path(dir); - fs::create_directory(dir_path); - +void Context::save_globals(const std::string& file) const { + // sorted, so that the archive does not depend on the hash map order + std::vector names; + names.reserve(_globals.size()); for (auto& [name, tensor_and_grad] : _globals) { - auto& [tensor, requires_grad] = tensor_and_grad; - fs::path tensor_file = dir_path / name; - tensor_file += ".npy"; - save_tensor(tensor_file, tensor); + names.push_back(name); + } + std::sort(names.begin(), names.end()); + + TarWriter tar(file); + for (auto& name : names) { + save_tensor(tar, name + ".npy", _globals.at(name).first); } + tar.close(); } -void Context::load_globals(const std::string& dir) { - namespace fs = std::filesystem; - for (auto& file : fs::directory_iterator(dir)) { - if (file.path().extension() != ".npy") { +void Context::load_globals(const std::string& file) { + constexpr std::string_view extension = ".npy"; + TarReader tar(file); + for (auto& file_name : tar.names()) { + if (!file_name.ends_with(extension)) { continue; } - std::string name = file.path().stem(); - Tensor tensor = load_tensor(file.path()); + std::string name = file_name.substr(0, file_name.size() - extension.size()); + Tensor tensor = load_tensor(tar, file_name); Tensor global_tensor = define_global( name, tensor.dtype(), diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index 39b55d84d4..a4051d8e5e 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{ @@ -55,6 +57,11 @@ EventGenerator::EventGenerator( _channel_unweight_order(channels.size()), _channel_unweight_ready(channels.size()), _context_unweight_queue(contexts.size()), + _pipeline( + groups.empty() + ? nullptr + : std::make_unique(contexts, channels, groups, config) + ), _seed(seed), _status_file(status_file) {} @@ -83,7 +90,11 @@ void EventGenerator::commit_generate_job(GeneratorBatchJob& job) { void EventGenerator::commit_unweight_job(GeneratorBatchJob& job) { auto& channel = _channels.at(job.channel_index); - channel->write_events(job.unweighted_events, job.max_weight); + if (_pipeline && _pipeline->contains(job.channel_index)) { + _pipeline->add_events(job); + } else { + channel->write_events(job.unweighted_events, job.max_weight); + } update_counts(); --_channel_job_counts.at(job.channel_index); --_context_job_counts.at(job.context_index); @@ -106,7 +117,14 @@ void EventGenerator::finish_channel_job(const GeneratorBatchJob& job) { // Jobs run on the pool, but commit strictly in per-channel dispatch order, so // results are a pure function of the seed. -void EventGenerator::generate() { +void EventGenerator::generate() { run_generate(GenerateMode::events); } + +void EventGenerator::optimize() { run_generate(GenerateMode::optimize_only); } + +void EventGenerator::fix_max_weights() { run_generate(GenerateMode::fix_max_weights); } + +void EventGenerator::run_generate(GenerateMode mode) { + _generate_mode = mode; _survey_job = false; reset_start_time(); print_gen_init(); @@ -116,11 +134,22 @@ void EventGenerator::generate() { // so round one truncates at all; _max_weight only rises, and set_target_count() // tightens it when the real (smaller) target arrives. for (auto& channel : _channels) { - if (channel->status().count_opt == 0) { - channel->set_target_count(_config.target_count); + if (channel->status().count_opt == 0 || mode == GenerateMode::fix_max_weights) { + channel->set_target_count( + mode == GenerateMode::fix_max_weights + ? _config.freeze_max_weight_after + : _config.target_count + ); } } + // channels loaded with a known integral can share the events from the start + if (std::any_of(_channels.begin(), _channels.end(), [](auto& channel) { + return channel->has_abs_integral_prior(); + })) { + update_integral_fractions(); + } + while (true) { _abort_check_function(); @@ -128,8 +157,13 @@ void EventGenerator::generate() { ++channel_index) { auto& channel = _channels.at(channel_index); double integral_frac = _channel_integral_fractions.at(channel_index); - if (integral_frac > 0 && - channel->status().count_unweighted >= channel->status().count_target) { + if (mode == GenerateMode::optimize_only) { + if (!channel->needs_optimization()) { + continue; + } + } else if (integral_frac > 0 && + channel->status().count_unweighted >= + channel->status().count_target) { continue; } if (channel->needs_optimization()) { @@ -161,6 +195,16 @@ void EventGenerator::generate() { // 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); @@ -188,7 +232,18 @@ void EventGenerator::generate() { } 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) { + return channel->needs_optimization(); + })) { + break; + } + continue; } update_integral_fractions(); @@ -258,6 +313,15 @@ void EventGenerator::survey(std::size_t survey_pass) { done = true; while (in_flight > 0) { std::size_t job_id = wait_for_result(in_flight); + if (WeightingPipeline::owns_job(job_id)) { + _pipeline->commit_job(job_id); + update_counts(); + std::size_t job_id_refill = _job_id; + std::size_t refill_unweight_dispatched = start_jobs(); + in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched + + start_pipeline_jobs(); + continue; + } auto& job = _running_jobs.at(job_id); if (job.unweighted_events.size() == 0) { // Generate-stage completion, globally ordered via @@ -306,7 +370,13 @@ void EventGenerator::survey(std::size_t survey_pass) { order.pop_front(); auto& uw_job = _running_jobs.at(commit_id); auto& channel = _channels.at(uw_job.channel_index); - channel->write_events(uw_job.unweighted_events, uw_job.max_weight); + if (_pipeline && _pipeline->contains(uw_job.channel_index)) { + _pipeline->add_events(uw_job); + } else { + channel->write_events( + uw_job.unweighted_events, uw_job.max_weight + ); + } update_counts(); auto& channel_job_count = _channel_job_counts.at(uw_job.channel_index); @@ -325,9 +395,11 @@ void EventGenerator::survey(std::size_t survey_pass) { } std::size_t job_id_refill = _job_id; std::size_t refill_unweight_dispatched = start_jobs(); - in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched; + in_flight += (_job_id - job_id_refill) + refill_unweight_dispatched + + start_pipeline_jobs(); print_survey_update(false, done_event_count, total_event_count, iter); } + drain_pipeline(); } print_survey_update(true, done_event_count, total_event_count, iter - 1); } @@ -359,6 +431,9 @@ std::size_t EventGenerator::wait_for_result(std::size_t& in_flight) { void EventGenerator::reset_jobs() { _running_jobs.clear(); + if (_pipeline) { + _pipeline->reset(); + } _ready_jobs.clear(); _ready_job_rr_cursor = 0; _ready_gen.clear(); @@ -492,89 +567,45 @@ 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() { + return _pipeline ? _pipeline->start_jobs(_result_queue, false) : 0; } -void EventGenerator::update_integral_status() { - double total_mean = 0., total_var = 0.; - double total_mean_abs = 0., total_var_abs = 0.; - std::size_t total_count = 0, total_count_opt = 0; - std::size_t total_count_after_cuts = 0, total_count_after_cuts_opt = 0; - std::size_t total_integ_count = 0; - std::size_t iterations = 0; - bool optimized = true; - const ChannelEventGenerator* bad_channel = nullptr; - for (auto& channel : _channels) { - auto& status = channel->status(); - auto& cross_section = channel->cross_section(); - auto& abs_cross_section = channel->abs_cross_section(); - // special case for channels with 0 samples, as they have nan variance - if (bad_channel == nullptr && cross_section.count() > 0 && - (!std::isfinite(cross_section.mean()) || - !std::isfinite(cross_section.variance()) || - !std::isfinite(abs_cross_section.mean()) || - !std::isfinite(abs_cross_section.variance()))) { - bad_channel = channel.get(); - } - total_mean += cross_section.mean(); - total_var += cross_section.variance() / cross_section.count(); - total_mean_abs += abs_cross_section.mean(); - total_var_abs += abs_cross_section.variance() / abs_cross_section.count(); - total_count += status.count; - total_count_opt += status.count_opt; - total_count_after_cuts += status.count_after_cuts; - total_count_after_cuts_opt += status.count_after_cuts_opt; - total_integ_count += cross_section.count(); - iterations = std::max(status.iterations, iterations); - if (!channel->status().optimized) { - optimized = false; - } +// 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(); +} + +void EventGenerator::update_integral() { + update_integral_status(); + update_integral_fractions(false); +} - _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_status() { + sum_channel_status(_channels, _status); } -void EventGenerator::update_integral_fractions() { +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_cross_section().mean() / _status.mean_abs; + integral_fraction = channel->abs_integral_estimate(final_step) / total_estimate; } // Distribute events between channels, ensure sum is exactly target_count @@ -601,7 +632,16 @@ void EventGenerator::update_integral_fractions() { for (std::size_t index : order | std::views::take(leftover)) { ++counts.at(index); } - for (auto [channel, count] : zip(_channels, counts)) { + for (auto [channel, count, fraction] : + zip(_channels, counts, _channel_integral_fractions)) { + if (_generate_mode == GenerateMode::fix_max_weights) { + // the maxima of all steps have to freeze + std::size_t freeze_count = _config.freeze_max_weight_after; + for (auto& step : channel->weighting_steps()) { + freeze_count = std::max(freeze_count, step.freeze_max_weight_after); + } + count = fraction > 0 ? freeze_count : 0; + } channel->set_target_count(count); } } @@ -620,6 +660,13 @@ void EventGenerator::update_counts() { } _status.count_unweighted = total_eff_count; _status.done = done; + if (_generate_mode == GenerateMode::fix_max_weights) { + // every channel runs to its own target, whatever the config says + _status.count_target = 0; + for (auto& channel : _channels) { + _status.count_target += channel->status().count_target; + } + } } void EventGenerator::combine_to_compact_npy( @@ -631,7 +678,7 @@ void EventGenerator::combine_to_compact_npy( _systematics = systematics; _event_histograms = histograms; MixMaxRandom select_rng(DerivedSeed(_seed, DerivedSeed::combine_select)); - auto [channel_data, particle_count, norm_factor] = init_combine(); + auto [channel_data, particle_count, norm_factors] = init_combine(); DataLayout layout = combined_layout(); DataLayout out_layout = output_layout( EventRecord::f_weight | EventRecord::f_subproc_index | @@ -647,7 +694,7 @@ void EventGenerator::combine_to_compact_npy( print_combine_init(); while (true) { _abort_check_function(); - read_and_combine(channel_data, buffer, norm_factor, select_rng); + read_and_combine(channel_data, buffer, norm_factors, select_rng); if (buffer.event_count() == 0) { break; } @@ -680,7 +727,7 @@ void EventGenerator::combine_to_lhe_npy( std::uint64_t seed = _seed; MixMaxRandom select_rng(DerivedSeed(seed, DerivedSeed::combine_select)); MixMaxRandom rand_gen(DerivedSeed(seed, DerivedSeed::lhe_complete)); - auto [channel_data, particle_count, norm_factor] = init_combine(); + auto [channel_data, particle_count, norm_factors] = init_combine(); DataLayout in_layout = combined_layout(); DataLayout out_layout = output_layout(EventRecord::f_lhe_event, ParticleRecord::f_lhe_particle); @@ -696,7 +743,7 @@ void EventGenerator::combine_to_lhe_npy( print_combine_init(); while (true) { _abort_check_function(); - read_and_combine(channel_data, buffer, norm_factor, select_rng); + read_and_combine(channel_data, buffer, norm_factors, select_rng); if (buffer.event_count() == 0) { break; } @@ -707,13 +754,7 @@ void EventGenerator::combine_to_lhe_npy( fill_lhe_event(lhe_completer, lhe_event, buffer, i, rand_gen); buffer_out.event(i).from_lhe_event(lhe_event); copy_event_fields(buffer, i, buffer_out, i, syst_weights); - std::size_t j = 0; - for (; j < lhe_event.particles.size(); ++j) { - buffer_out.particle(i, j).from_lhe_particle(lhe_event.particles[j]); - } - for (; j < lhe_completer.max_particle_count(); ++j) { - buffer_out.particle(i, j).from_lhe_particle(LHEParticle{}); - } + write_lhe_particles(lhe_event, buffer_out, i); } event_file.write(buffer_out); if (event_count - last_update_count > 10000) { @@ -739,7 +780,7 @@ void EventGenerator::combine_to_lhe( std::uint64_t seed = _seed; MixMaxRandom select_rng(DerivedSeed(seed, DerivedSeed::combine_select)); ThreadPool pool(_config.combine_thread_count); - auto [channel_data, particle_count, norm_factor] = init_combine(); + auto [channel_data, particle_count, norm_factors] = init_combine(); struct CombineBuffers { EventBuffer in_buffer; std::string out_buffer; @@ -782,7 +823,7 @@ void EventGenerator::combine_to_lhe( while (idle_buffers.size() > 0 && !done) { std::size_t slot = idle_buffers.back(); auto& [in_buffer, out_buffer, syst_weights] = buffers.at(slot); - read_and_combine(channel_data, in_buffer, norm_factor, select_rng); + read_and_combine(channel_data, in_buffer, norm_factors, select_rng); if (in_buffer.event_count() == 0) { done = true; break; @@ -848,21 +889,16 @@ void EventGenerator::combine_to_lhe( } DataLayout EventGenerator::combined_layout() const { - return DataLayout( - EventRecord::layout( - EventRecord::f_weight | EventRecord::f_subproc_index | - EventRecord::f_event_data | _channels.at(0)->event_layout_extra_flags() - ), - ParticleRecord::layout( - ParticleRecord::f_particle_data | - _channels.at(0)->particle_layout_extra_flags() - ) + return combined_event_layout( + _channels.at(0)->event_layout_extra_flags(), + _channels.at(0)->particle_layout_extra_flags() ); } DataLayout EventGenerator::output_layout(int event_flags, std::size_t particle_flags) const { - int extra_flags = _channels.at(0)->event_layout_extra_flags(); + int extra_flags = + _channels.at(0)->event_layout_extra_flags() & ~EventRecord::f_bias; std::vector syst_ids; if (_systematics) { if (!_systematics->config().write_inputs) { @@ -1071,12 +1107,16 @@ std::vector EventGenerator::histograms() const { return hists; } -std::tuple, std::size_t, double> +std::tuple< + std::vector, + std::size_t, + std::vector> EventGenerator::init_combine() { std::vector channel_data; std::size_t count_sum = 0; std::size_t particle_count = 0; double weight_sum = 0.; + std::vector channel_weight_sums; for (auto& channel : _channels) { particle_count = std::max(particle_count, channel->event_file().particle_count()); @@ -1085,36 +1125,46 @@ 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_subproc_index = - _channels.at(0)->event_layout_extra_flags() & EventRecord::f_subproc_index; + int extra_flags = _channels.at(0)->event_layout_extra_flags(); for (std::size_t event_index = 0; event_index < event_count; ++event_index) { std::size_t random_index = rand_gen.generate_int(channel_data.back().cum_count); @@ -1127,7 +1177,8 @@ void EventGenerator::read_and_combine( std::for_each(sampled_chan, channel_data.end(), [](auto& chan) { --chan.cum_count; }); - auto& channel = _channels.at(sampled_chan - channel_data.begin()); + std::size_t channel_index = sampled_chan - channel_data.begin(); + auto& channel = _channels.at(channel_index); double weight = 0.; while (true) { @@ -1145,93 +1196,19 @@ void EventGenerator::read_and_combine( ++sampled_chan->buffer_index; } - auto event_in = sampled_chan->event_buffer.event(sampled_chan->buffer_index); - auto event_out = buffer.event(event_index); - event_out.weight() = - std::copysign( - std::max(1., std::abs(weight) / channel->max_weight()), weight - ) * - norm_factor; - event_out.subprocess_index() = has_subproc_index - ? event_in.subprocess_index().value() - : static_cast(channel->status().subprocess); - event_out.diagram_index() = event_in.diagram_index(); - event_out.color_index() = event_in.color_index(); - event_out.flavor_index() = event_in.flavor_index(); - event_out.helicity_index() = event_in.helicity_index(); - event_out.ren_scale() = event_in.ren_scale(); - event_out.alpha_qcd() = event_in.alpha_qcd(); - - if (has_beam1) { - event_out.x1() = event_in.x1(); - event_out.fact_scale1() = event_in.fact_scale1(); - } - if (has_beam2) { - event_out.x2() = event_in.x2(); - event_out.fact_scale2() = event_in.fact_scale2(); - } - if (has_partial) { - event_out.partial_weight_product() = event_in.partial_weight_product(); - } - - std::size_t i = 0; - for (; i < sampled_chan->event_buffer.particle_count(); ++i) { - auto particle_in = - sampled_chan->event_buffer.particle(sampled_chan->buffer_index, i); - auto particle_out = buffer.particle(event_index, i); - particle_out.energy() = particle_in.energy(); - particle_out.px() = particle_in.px(); - particle_out.py() = particle_in.py(); - particle_out.pz() = particle_in.pz(); - } - for (; i < buffer.particle_count(); ++i) { - auto particle_out = buffer.particle(event_index, i); - particle_out.energy() = 0.; - particle_out.px() = 0.; - particle_out.py() = 0.; - particle_out.pz() = 0.; - } - ++sampled_chan->buffer_index; - } -} - -void EventGenerator::fill_lhe_event( - LHECompleter& lhe_completer, - LHEEvent& lhe_event, - EventBuffer& buffer, - std::size_t event_index, - MixMaxRandom& rand_gen -) { - EventRecord event_in = buffer.event(event_index); - lhe_event.weight = event_in.weight(); - lhe_event.process_id = 0; - lhe_event.scale = event_in.ren_scale(); - lhe_event.alpha_qed = 0; // TODO: populate this - lhe_event.alpha_qcd = event_in.alpha_qcd(); - lhe_event.particles.clear(); - for (std::size_t i = 0; i < buffer.particle_count(); ++i) { - auto particle_in = buffer.particle(event_index, i); - if (particle_in.energy() == 0.) { - break; - } - lhe_event.particles.push_back( - LHEParticle{ - .px = particle_in.px(), - .py = particle_in.py(), - .pz = particle_in.pz(), - .energy = particle_in.energy(), - } + copy_channel_event( + sampled_chan->event_buffer, + sampled_chan->buffer_index, + buffer, + event_index, + channel->unweighted_weight( + sampled_chan->weight_buffer, sampled_chan->buffer_index + ) * norm_factors.at(channel_index), + static_cast(channel->status().subprocess), + extra_flags ); + ++sampled_chan->buffer_index; } - lhe_completer.complete_event_data( - lhe_event, - event_in.subprocess_index(), - event_in.diagram_index(), - event_in.color_index(), - event_in.flavor_index(), - event_in.helicity_index(), - rand_gen - ); } void EventGenerator::init_status(const std::string& status) { @@ -1372,11 +1349,26 @@ void EventGenerator::print_survey_update_log( } } +std::pair EventGenerator::optimized_channel_count() const { + std::size_t optimizable = 0, optimized = 0; + for (auto& channel : _channels) { + // a channel without an optimizer is never marked optimized + if (channel->status().optimized) { + ++optimizable; + ++optimized; + } else if (channel->needs_optimization()) { + ++optimizable; + } + } + return {optimized, optimizable}; +} + void EventGenerator::print_gen_init() { init_status("generate"); _last_print_time = std::chrono::steady_clock::now(); + bool optimize_only = _generate_mode == GenerateMode::optimize_only; if (_config.verbosity != Verbosity::pretty) { - Logger::info("generating started"); + Logger::info(optimize_only ? "optimizing started" : "generating started"); return; } @@ -1385,33 +1377,109 @@ void EventGenerator::print_gen_init() { _pretty_box_lower = PrettyBox( "Individual channels", _channels.size() < 21 ? _channels.size() + 1 : 22, - {6, 16, 9, 9, 7, 6, 0} + optimize_only ? std::vector{6, 16, 9, 7, 6, 0} + : std::vector{6, 16, 9, 9, 7, 6, 0} ); _pretty_box_lower.set_row( - 0, {"#", "integral ↓", "RSD", "uweff", "N", "opt", "unweighted"} + 0, + optimize_only + ? std::vector< + std::string>{"#", "integral ↓", "RSD", "N", "iter", "patience"} + : std::vector{ + "#", "integral ↓", "RSD", "uweff", "N", "opt", "unweighted" + } ); if (_channels.size() > 20) { _pretty_box_lower.set_cell(21, 0, ".."); } offset = _pretty_box_lower.line_count(); } - _pretty_box_upper = PrettyBox("Integration and unweighting", 7, {19, 0}, offset); - _pretty_box_upper.set_column( - 0, - {"Result:", - "Rel. error:", - "Rel. stddev:", - "Number of events:", - "Unweighting eff.:", - "Unweighted events:", - "Run time:"} - ); + if (optimize_only) { + _pretty_box_upper = + PrettyBox("Integration and optimization", 6, {19, 0}, offset); + _pretty_box_upper.set_column( + 0, + {"Result:", + "Rel. error:", + "Rel. stddev:", + "Number of events:", + "Optimized channels:", + "Run time:"} + ); + } else { + _pretty_box_upper = PrettyBox( + _generate_mode == GenerateMode::fix_max_weights + ? "Integration and max. weight estimation" + : "Integration and unweighting", + 7, + {19, 0}, + offset + ); + _pretty_box_upper.set_column( + 0, + {"Result:", + "Rel. error:", + "Rel. stddev:", + "Number of events:", + "Unweighting eff.:", + "Unweighted events:", + "Run time:"} + ); + 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) { _pretty_box_lower.print_first(); } } +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"); @@ -1431,8 +1499,9 @@ void EventGenerator::print_gen_update_pretty(bool done) { return; } _last_print_time = now; + bool optimize_only = _generate_mode == GenerateMode::optimize_only; - std::string int_str, rel_str, rsd_str, uweff_str, count_str; + 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), @@ -1442,7 +1511,8 @@ void EventGenerator::print_gen_update_pretty(bool done) { if (!std::isnan(_status.error)) { int_str = format_with_error(_status.mean, _status.error); if (std::abs(_status.error) < std::abs(_status.mean)) { - rel_str = std::format("{:.4f} %", _status.error / std::abs(_status.mean) * 100); + rel_str = + std::format("{:.4f} %", _status.error / std::abs(_status.mean) * 100); } rsd_str = std::format("{:.3f}", _status.rel_std_dev); uweff_str = std::format( @@ -1450,29 +1520,65 @@ 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; + double progress_fraction; + if (optimize_only) { + auto [optimized, optimizable] = optimized_channel_count(); + progress_str = std::format("{} / {}", optimized, optimizable); + progress_fraction = + optimizable > 0 ? static_cast(optimized) / optimizable : 1.; + } else { + progress_str = std::format( + "{} / {}", + format_si_prefix(_status.count_unweighted), + format_si_prefix(_status.count_target) + ); + progress_fraction = _status.count_unweighted / _status.count_target; } - std::string unw_str = std::format( - "{} / {}", - format_si_prefix(_status.count_unweighted), - format_si_prefix(_status.count_target) - ); std::string time_str; if (done) { auto [wall_time_sec, cpu_time_sec] = _timing_data.at("generate"); time_str = format_run_time(wall_time_sec, cpu_time_sec); } else { - unw_str = std::format( - "{:<15} {}", - unw_str, - format_progress(_status.count_unweighted / _status.count_target, 52) + progress_str = std::format( + "{:<15} {}", progress_str, format_progress(progress_fraction, 52) ); time_str = std::format( "{:%H:%M:%S}", std::chrono::round(now - _start_time) ); } - _pretty_box_upper.set_column( - 1, {int_str, rel_str, rsd_str, count_str, uweff_str, unw_str, time_str} - ); + if (optimize_only) { + _pretty_box_upper.set_column( + 1, {int_str, rel_str, rsd_str, count_str, progress_str, time_str} + ); + } else 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} + ); + } _pretty_box_upper.print_update(); if (_channels.size() > 1) { @@ -1483,7 +1589,7 @@ void EventGenerator::print_gen_update_pretty(bool done) { for (std::size_t row = 1; auto& channel : channels | std::views::take(20)) { std::string index_str = std::format("{}", channel.name); - std::string int_str, rsd_str, count_str, unw_str, opt_str; + std::string int_str, rsd_str, count_str, unw_str, opt_str, patience_str; if (!std::isnan(channel.error)) { int_str = format_with_error(channel.mean, channel.error); rsd_str = std::format("{:.3f}", channel.rel_std_dev); @@ -1508,11 +1614,37 @@ void EventGenerator::print_gen_update_pretty(bool done) { ); } unw_str = std::format("{:<14} {:<19}", unw_count_str, progress); + std::size_t patience = _config.optimization_patience; + // an optimized channel stays so, even if its counter was reset later + std::size_t waited = + channel.optimized ? patience : channel.iters_without_improvement; + patience_str = std::format( + "{:>3} / {:<3} {}", + waited, + patience, + format_progress( + static_cast(waited) / + std::max(patience, 1), + 19 + ) + ); + } + if (optimize_only) { + _pretty_box_lower.set_row( + row, {index_str, int_str, rsd_str, count_str, opt_str, patience_str} + ); + } else { + _pretty_box_lower.set_row( + row, + {index_str, + int_str, + rsd_str, + uweff_str, + count_str, + opt_str, + unw_str} + ); } - _pretty_box_lower.set_row( - row, - {index_str, int_str, rsd_str, uweff_str, count_str, opt_str, unw_str} - ); ++row; } _pretty_box_lower.print_update(); @@ -1526,33 +1658,68 @@ void EventGenerator::print_gen_update_log(bool done) { return; } _last_print_time = now; + bool optimize_only = _generate_mode == GenerateMode::optimize_only; std::string rel_str = std::abs(_status.error) < std::abs(_status.mean) ? std::format("{:.4f} %", _status.error / std::abs(_status.mean) * 100) : ""; - Logger::info( - std::format( - "generating, events: {} / {}, integral: {}, rel. error: {}, " - "RSD: {:.3f}, samps: {}, samps. after cuts: {}, " - "unw. eff.: {:.5f}, unw. eff. after cuts: {:.5f}, time: {:%H:%M:%S}", - format_si_prefix(_status.count_unweighted), - format_si_prefix(_status.count_target), - format_with_error(_status.mean, _status.error), - rel_str, - _status.rel_std_dev, - format_si_prefix(_status.count), - format_si_prefix(_status.count_after_cuts), - _status.count_unweighted / _status.count_opt, - _status.count_unweighted / _status.count_after_cuts_opt, - std::chrono::round(now - _start_time) - ) - ); + if (optimize_only) { + auto [optimized, optimizable] = optimized_channel_count(); + Logger::info( + std::format( + "optimizing, optimized channels: {} / {}, integral: {}, " + "rel. error: {}, RSD: {:.3f}, samps: {}, samps. after cuts: {}, " + "time: {:%H:%M:%S}", + optimized, + optimizable, + format_with_error(_status.mean, _status.error), + rel_str, + _status.rel_std_dev, + format_si_prefix(_status.count), + format_si_prefix(_status.count_after_cuts), + std::chrono::round(now - _start_time) + ) + ); + } else { + 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}", + format_si_prefix(_status.count_unweighted), + format_si_prefix(_status.count_target), + format_with_error(_status.mean, _status.error), + rel_str, + _status.rel_std_dev, + format_si_prefix(_status.count), + format_si_prefix(_status.count_after_cuts), + _status.count_unweighted / _status.count_opt, + _status.count_unweighted / _status.count_after_cuts_opt, + pass_efficiency_str, + std::chrono::round(now - _start_time) + ) + ); + } if (done) { auto [wall_time_sec, cpu_time_sec] = _timing_data.at("generate"); Logger::info( std::format( - "generating done, {}", format_run_time(wall_time_sec, cpu_time_sec) + "{} done, {}", + optimize_only ? "optimizing" : "generating", + format_run_time(wall_time_sec, cpu_time_sec) ) ); } diff --git a/madspace/src/driver/event_stream.cpp b/madspace/src/driver/event_stream.cpp new file mode 100644 index 0000000000..8779a10a91 --- /dev/null +++ b/madspace/src/driver/event_stream.cpp @@ -0,0 +1,530 @@ +#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() + ); + }()), + _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() + ); + // keep the buffers filled while the batch is assembled + if (state.buffered_events == 0) { + _dispatch_requested = true; + _coordinator_cv.notify_all(); + } + } + _event_count += count; + _dispatch_requested = true; + } + _coordinator_cv.notify_all(); + + std::vector events(count); + for (std::size_t i = 0; i < count; ++i) { + fill_lhe_event(_lhe_completer, events.at(i), buffer, i, _lhe_rng); + } + return events; +} + +GeneratorStatus EventStream::status() const { + std::unique_lock lock(_mutex); + GeneratorStatus status{}; + sum_channel_status(_channels, status); + status.count_unweighted = _event_count; + status.count_target = 0; + status.optimized = true; + status.done = false; + return status; +} + +std::vector EventStream::channel_status() const { + std::unique_lock lock(_mutex); + std::vector ret; + for (auto [channel, state] : zip(_channels, _channel_states)) { + auto status = channel->status(); + status.count_unweighted = state.returned_events; + status.count_target = 0; + ret.push_back(status); + } + return ret; +} + +std::vector EventStream::channel_probabilities() const { + std::unique_lock lock(_mutex); + // const: compute from scratch instead of updating the cache + auto& self = const_cast(*this); + if (self._probabilities_dirty) { + self.update_probabilities(); + } + std::vector probs; + double prev = 0.; + for (double cum_prob : _cum_probabilities) { + probs.push_back(cum_prob - prev); + prev = cum_prob; + } + return probs; +} + +std::size_t EventStream::event_count() const { + std::unique_lock lock(_mutex); + return _event_count; +} diff --git a/madspace/src/driver/event_utils.cpp b/madspace/src/driver/event_utils.cpp new file mode 100644 index 0000000000..8817744039 --- /dev/null +++ b/madspace/src/driver/event_utils.cpp @@ -0,0 +1,200 @@ +#include "madspace/driver/event_utils.hpp" + +#include +#include +#include + +using namespace madspace; + +DataLayout +madspace::combined_event_layout(int extra_event_flags, int extra_particle_flags) { + return DataLayout( + EventRecord::layout( + EventRecord::f_weight | EventRecord::f_subproc_index | + EventRecord::f_event_data | (extra_event_flags & ~EventRecord::f_bias) + ), + ParticleRecord::layout(ParticleRecord::f_particle_data | extra_particle_flags) + ); +} + +void madspace::copy_channel_event( + EventBuffer& in_buffer, + std::size_t in_index, + EventBuffer& out_buffer, + std::size_t out_index, + double weight, + int subprocess_index, + int extra_flags +) { + auto event_in = in_buffer.event(in_index); + auto event_out = out_buffer.event(out_index); + // 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(); + } + for (; i < out_buffer.particle_count(); ++i) { + auto particle_out = out_buffer.particle(out_index, i); + particle_out.energy() = 0.; + particle_out.px() = 0.; + particle_out.py() = 0.; + particle_out.pz() = 0.; + } +} + +void madspace::fill_lhe_event( + LHECompleter& lhe_completer, + LHEEvent& lhe_event, + EventBuffer& buffer, + std::size_t event_index, + MixMaxRandom& rand_gen +) { + EventRecord event_in = buffer.event(event_index); + lhe_event.weight = event_in.weight(); + lhe_event.process_id = 0; + lhe_event.scale = event_in.ren_scale(); + lhe_event.alpha_qed = 0; // TODO: populate this + lhe_event.alpha_qcd = event_in.alpha_qcd(); + lhe_event.particles.clear(); + for (std::size_t i = 0; i < buffer.particle_count(); ++i) { + auto particle_in = buffer.particle(event_index, i); + if (particle_in.energy() == 0.) { + break; + } + lhe_event.particles.push_back( + LHEParticle{ + .px = particle_in.px(), + .py = particle_in.py(), + .pz = particle_in.pz(), + .energy = particle_in.energy(), + } + ); + } + lhe_completer.complete_event_data( + lhe_event, + event_in.subprocess_index(), + event_in.diagram_index(), + event_in.color_index(), + event_in.flavor_index(), + event_in.helicity_index(), + rand_gen + ); +} + +void madspace::write_lhe_particles( + const LHEEvent& lhe_event, EventBuffer& buffer, std::size_t event_index +) { + std::size_t i = 0; + for (; i < lhe_event.particles.size(); ++i) { + buffer.particle(event_index, i).from_lhe_particle(lhe_event.particles[i]); + } + for (; i < buffer.particle_count(); ++i) { + buffer.particle(event_index, i).from_lhe_particle(LHEParticle{}); + } +} + +void madspace::sum_channel_status( + const std::vector>& channels, + GeneratorStatus& status +) { + double total_mean = 0., total_var = 0.; + double total_mean_abs = 0., total_var_abs = 0.; + std::size_t total_count = 0, total_count_opt = 0; + std::size_t total_count_after_cuts = 0, total_count_after_cuts_opt = 0; + std::size_t total_integ_count = 0; + std::size_t iterations = 0; + bool optimized = true; + const ChannelEventGenerator* bad_channel = nullptr; + for (auto& channel : channels) { + auto& chan_status = channel->status(); + auto 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(); + } + 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); + 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.optimized = optimized; +} diff --git a/madspace/src/driver/generator_data.cpp b/madspace/src/driver/generator_data.cpp index 579bfd4f9a..1b3679ed61 100644 --- a/madspace/src/driver/generator_data.cpp +++ b/madspace/src/driver/generator_data.cpp @@ -73,6 +73,7 @@ void madspace::to_json(nlohmann::json& j, const GeneratorStatus& status) { {"count_unweighted", status.count_unweighted}, {"count_target", status.count_target}, {"iterations", status.iterations}, + {"iters_without_improvement", status.iters_without_improvement}, {"optimized", status.optimized}, {"done", status.done}, }; @@ -87,3 +88,27 @@ void madspace::to_json(nlohmann::json& j, const Histogram& hist) { {"bin_errors", hist.bin_errors}, }; } + +StepChainResult madspace::evaluate_step_chain( + double weight0, + double max_weight0, + const double* ratios, + const double* max_weights, + const std::vector& steps +) { + double abs_weight0 = std::abs(weight0); + double acceptance = std::min(1., abs_weight0 / max_weight0); + double overweight = std::max(1., abs_weight0 / max_weight0); + double pending = std::copysign(1., weight0); + for (std::size_t i = 0; auto& step : steps) { + pending *= ratios[i]; + if (step.unweight) { + double unweight_value = overweight * std::abs(pending); + acceptance *= std::min(1., unweight_value / max_weights[i]); + overweight = std::max(1., unweight_value / max_weights[i]); + pending = std::copysign(1., pending); + } + ++i; + } + return {acceptance, overweight * pending}; +} diff --git a/madspace/src/driver/io.cpp b/madspace/src/driver/io.cpp index 3adafb2db9..2c733d9280 100644 --- a/madspace/src/driver/io.cpp +++ b/madspace/src/driver/io.cpp @@ -36,19 +36,20 @@ DataType str_to_dtype(std::string dtype) { } } -json parse_header(std::fstream& file_stream) { - char magic_num[8]; - file_stream.read(magic_num, 8); - if (std::memcmp(magic_num, "\x93NUMPY\x01\x00", 8) != 0) { +constexpr std::size_t npy_preamble_size = 10; + +/// Check the magic number and version of the first 10 bytes of an npy file and +/// return the length of the header that follows. +std::size_t parse_preamble(const char* bytes) { + if (std::memcmp(bytes, "\x93NUMPY\x01\x00", 8) != 0) { throw std::runtime_error("invalid header of npy file"); } - union { - uint16_t size; - char chars[2]; - } header_size; - file_stream.read(&header_size.chars[0], 2); - std::string header(header_size.size, '\0'); - file_stream.read(header.data(), header_size.size); + // little-endian 16 bit length + return static_cast(static_cast(bytes[8])) | + (static_cast(static_cast(bytes[9])) << 8); +} + +json parse_header_text(std::string header) { header.erase( std::remove_if( header.begin(), header.end(), [](char x) { return std::isspace(x); } @@ -80,6 +81,62 @@ json parse_header(std::fstream& file_stream) { return json::parse(header); } +json parse_header(std::fstream& file_stream) { + char preamble[npy_preamble_size]; + file_stream.read(preamble, npy_preamble_size); + if (file_stream.fail()) { + throw std::runtime_error("invalid header of npy file"); + } + std::string header(parse_preamble(preamble), '\0'); + file_stream.read(header.data(), header.size()); + return parse_header_text(std::move(header)); +} + +/// The dtype and shape described by the parsed header of a tensor npy file. +std::pair parse_tensor_header(json header) { + if (!header.is_object()) { + throw std::runtime_error("Invalid header"); + } + json descr = header.at("descr"); + json fortran_order = header.at("fortran_order"); + json header_shape = header.at("shape"); + if (!descr.is_string() || !fortran_order.is_boolean() || + !fortran_order.get() || !header_shape.is_array()) { + throw std::runtime_error("Invalid file header"); + } + DataType dtype = str_to_dtype(descr); + SizeVec shape; + for (auto& size : header_shape) { + if (!size.is_number_unsigned()) { + throw std::runtime_error("Invalid file header"); + } + shape.push_back(size.get()); + } + return {dtype, shape}; +} + +/// The npy header of `tensor`: 128 bytes, padded with spaces. +std::string tensor_header(const Tensor& tensor) { + using namespace std::string_literals; + // magic, version 1.0, little-endian header length of 118 + std::string header = "\x93NUMPY\x01\x00"s; + header.push_back(static_cast(118)); + header.push_back('\0'); + header += std::format( + "{{'descr':'{}','fortran_order':True,'shape':(", dtype_to_str(tensor.dtype()) + ); + for (std::size_t size : tensor.shape()) { + header += std::format("{},", size); + } + header += ")}"; + if (header.size() > 127) { + throw std::runtime_error("Tensor has too many dimensions to save"); + } + header.resize(127, ' '); + header.push_back('\n'); + return header; +} + std::vector> full_descr(std::size_t particle_count, const DataLayout& layout) { std::vector> field_descr; @@ -154,8 +211,13 @@ std::tuple write_event_header( } file_stream.put('\n'); file_stream.seekp(8); - uint16_t header_size_short = header_size - 10; - file_stream.write(reinterpret_cast(&header_size_short), 2); + // little-endian 16 bit length + std::size_t header_length = header_size - npy_preamble_size; + char length_bytes[2] = { + static_cast(header_length & 0xff), + static_cast((header_length >> 8) & 0xff) + }; + file_stream.write(length_bytes, 2); file_stream.seekp(header_size); return {header_size, shape_pos}; } @@ -167,25 +229,7 @@ Tensor madspace::load_tensor(const std::string& file) { if (file_stream.fail()) { throw std::runtime_error(std::format("Could not open file '{}'", file)); } - json header = parse_header(file_stream); - if (!header.is_object()) { - throw std::runtime_error("Invalid header"); - } - json descr = header.at("descr"); - json fortran_order = header.at("fortran_order"); - json header_shape = header.at("shape"); - if (!descr.is_string() || !fortran_order.is_boolean() || - !fortran_order.get() || !header_shape.is_array()) { - throw std::runtime_error("Invalid file header"); - } - DataType dtype = str_to_dtype(descr); - SizeVec shape; - for (auto& size : header_shape) { - if (!size.is_number_unsigned()) { - throw std::runtime_error("Invalid file header"); - } - shape.push_back(size.get()); - } + auto [dtype, shape] = parse_tensor_header(parse_header(file_stream)); Tensor tensor(dtype, shape); file_stream.read(static_cast(tensor.data()), tensor.byte_size()); if (file_stream.fail()) { @@ -195,24 +239,45 @@ Tensor madspace::load_tensor(const std::string& file) { } void madspace::save_tensor(const std::string& file, Tensor tensor) { - using namespace std::string_literals; Tensor cpu_tensor = tensor.cpu().contiguous(); std::ofstream file_stream(file, std::ios::binary); - file_stream << "\x93NUMPY\x01\x00\x76\x00{'descr':'"s - << dtype_to_str(tensor.dtype()) << "','fortran_order':True,'shape':("; - for (std::size_t size : tensor.shape()) { - file_stream << size << ","; - } - file_stream << ")}"; - for (int i = file_stream.tellp(); i < 127; ++i) { - file_stream.put(' '); - } - file_stream.put('\n'); + std::string header = tensor_header(tensor); + file_stream.write(header.data(), header.size()); file_stream.write( static_cast(cpu_tensor.data()), cpu_tensor.byte_size() ); } +Tensor madspace::load_tensor(TarReader& tar, const std::string& name) { + std::vector data = tar.read(name); + if (data.size() < npy_preamble_size) { + throw std::runtime_error(std::format("Invalid npy file '{}'", name)); + } + std::size_t header_length = parse_preamble(data.data()); + std::size_t data_offset = npy_preamble_size + header_length; + if (data.size() < data_offset) { + throw std::runtime_error(std::format("Invalid npy file '{}'", name)); + } + auto [dtype, shape] = parse_tensor_header( + parse_header_text(std::string(data.data() + npy_preamble_size, header_length)) + ); + Tensor tensor(dtype, shape); + if (data.size() - data_offset != tensor.byte_size()) { + throw std::runtime_error(std::format("Wrong data size in npy file '{}'", name)); + } + std::memcpy(tensor.data(), data.data() + data_offset, tensor.byte_size()); + return tensor; +} + +void madspace::save_tensor(TarWriter& tar, const std::string& name, Tensor tensor) { + Tensor cpu_tensor = tensor.cpu().contiguous(); + std::string header = tensor_header(tensor); + std::vector data(header.size() + cpu_tensor.byte_size()); + std::memcpy(data.data(), header.data(), header.size()); + std::memcpy(data.data() + header.size(), cpu_tensor.data(), cpu_tensor.byte_size()); + tar.add(name, data.data(), data.size()); +} + EventFile::EventFile( const std::string& file_name, const DataLayout& layout, diff --git a/madspace/src/driver/tar.cpp b/madspace/src/driver/tar.cpp new file mode 100644 index 0000000000..d64ce7198b --- /dev/null +++ b/madspace/src/driver/tar.cpp @@ -0,0 +1,253 @@ +#include "madspace/driver/tar.hpp" + +#include +#include +#include +#include +#include + +using namespace madspace; + +namespace { + +constexpr std::size_t block_size = 512; +using Block = std::array; + +// ustar header layout +constexpr std::size_t name_pos = 0, name_len = 100; +constexpr std::size_t mode_pos = 100, mode_len = 8; +constexpr std::size_t uid_pos = 108, uid_len = 8; +constexpr std::size_t gid_pos = 116, gid_len = 8; +constexpr std::size_t size_pos = 124, size_len = 12; +constexpr std::size_t mtime_pos = 136, mtime_len = 12; +constexpr std::size_t chksum_pos = 148, chksum_len = 8; +constexpr std::size_t typeflag_pos = 156; +constexpr std::size_t magic_pos = 257, magic_len = 6; +constexpr std::size_t version_pos = 263; +constexpr std::size_t prefix_pos = 345, prefix_len = 155; + +/// Write `value` as zero-padded octal into a field of `len` bytes, ending +/// with a NUL. +void put_octal(Block& block, std::size_t pos, std::size_t len, std::uint64_t value) { + for (std::size_t i = len - 1; i-- > 0;) { + block[pos + i] = static_cast('0' + (value & 7)); + value >>= 3; + } + if (value != 0) { + throw std::runtime_error("Value too large for tar header field"); + } + block[pos + len - 1] = '\0'; +} + +/// Read an octal field, ignoring leading spaces and stopping at the first NUL +/// or space after the digits. +std::uint64_t get_octal(const Block& block, std::size_t pos, std::size_t len) { + if (static_cast(block[pos]) & 0x80) { + throw std::runtime_error("Unsupported binary number in tar header"); + } + std::uint64_t value = 0; + bool digits = false; + for (std::size_t i = 0; i < len; ++i) { + char c = block[pos + i]; + if (c == ' ' && !digits) { + continue; + } + if (c < '0' || c > '7') { + break; + } + value = (value << 3) | static_cast(c - '0'); + digits = true; + } + return value; +} + +std::uint64_t checksum(const Block& block) { + std::uint64_t sum = 0; + for (std::size_t i = 0; i < block_size; ++i) { + bool in_chksum = i >= chksum_pos && i < chksum_pos + chksum_len; + sum += in_chksum ? 32 : static_cast(block[i]); + } + return sum; +} + +bool is_zero_block(const Block& block) { + return std::all_of(block.begin(), block.end(), [](char c) { return c == '\0'; }); +} + +std::string get_string(const Block& block, std::size_t pos, std::size_t len) { + std::size_t end = 0; + while (end < len && block[pos + end] != '\0') { + ++end; + } + return std::string(&block[pos], end); +} + +/// Split `name` into the ustar prefix and name fields. +std::pair split_name(const std::string& name) { + if (name.empty()) { + throw std::invalid_argument("Empty tar entry name"); + } + if (name.size() <= name_len) { + return {"", name}; + } + // the split has to be at a '/', with the last part fitting in the name field + for (std::size_t pos = name.find('/'); pos != std::string::npos; + pos = name.find('/', pos + 1)) { + if (pos <= prefix_len && name.size() - pos - 1 <= name_len && + name.size() - pos - 1 > 0) { + return {name.substr(0, pos), name.substr(pos + 1)}; + } + } + throw std::invalid_argument(std::format("Tar entry name too long: {}", name)); +} + +} // namespace + +TarWriter::TarWriter(const std::string& file_name) : + _file_name(file_name), + _stream(file_name, std::ios::binary | std::ios::trunc), + _closed(false) { + if (!_stream) { + throw std::runtime_error(std::format("Could not open file '{}'", file_name)); + } +} + +TarWriter::~TarWriter() { + try { + close(); + } catch (...) { + } +} + +void TarWriter::add(const std::string& name, const char* data, std::size_t size) { + if (_closed) { + throw std::runtime_error("Tar archive is already closed"); + } + auto [prefix, short_name] = split_name(name); + Block header{}; + std::memcpy(&header[name_pos], short_name.data(), short_name.size()); + std::memcpy(&header[prefix_pos], prefix.data(), prefix.size()); + put_octal(header, mode_pos, mode_len, 0644); + put_octal(header, uid_pos, uid_len, 0); + put_octal(header, gid_pos, gid_len, 0); + put_octal(header, size_pos, size_len, size); + put_octal(header, mtime_pos, mtime_len, 0); + header[typeflag_pos] = '0'; + std::memcpy(&header[magic_pos], "ustar", magic_len); // includes the NUL + header[version_pos] = '0'; + header[version_pos + 1] = '0'; + // six octal digits, NUL, space + put_octal(header, chksum_pos, 7, checksum(header)); + header[chksum_pos + 7] = ' '; + + _stream.write(header.data(), block_size); + _stream.write(data, static_cast(size)); + std::size_t padding = (block_size - size % block_size) % block_size; + static const Block zeros{}; + _stream.write(zeros.data(), static_cast(padding)); + if (!_stream) { + throw std::runtime_error( + std::format("Failed to write to file '{}'", _file_name) + ); + } +} + +void TarWriter::close() { + if (_closed) { + return; + } + _closed = true; + static const Block zeros{}; + _stream.write(zeros.data(), block_size); + _stream.write(zeros.data(), block_size); + _stream.close(); + if (_stream.fail()) { + throw std::runtime_error( + std::format("Failed to write to file '{}'", _file_name) + ); + } +} + +TarReader::TarReader(const std::string& file_name) : + _file_name(file_name), _stream(file_name, std::ios::binary) { + if (!_stream) { + throw std::runtime_error(std::format("Could not open file '{}'", file_name)); + } + Block header; + std::uint64_t pos = 0; + while (true) { + _stream.read(header.data(), block_size); + if (_stream.gcount() == 0 || is_zero_block(header)) { + break; + } + if (_stream.gcount() != static_cast(block_size)) { + throw std::runtime_error("Truncated tar header"); + } + if (get_octal(header, chksum_pos, chksum_len) != checksum(header)) { + throw std::runtime_error( + std::format("Invalid tar header checksum in '{}'", file_name) + ); + } + std::uint64_t size = get_octal(header, size_pos, size_len); + char type = header[typeflag_pos]; + pos += block_size; + if (type == '0' || type == '\0') { + std::string name = get_string(header, name_pos, name_len); + // the prefix field is only defined for ustar archives + if (std::memcmp(&header[magic_pos], "ustar", 5) == 0) { + std::string prefix = get_string(header, prefix_pos, prefix_len); + if (!prefix.empty()) { + name = prefix + "/" + name; + } + } + _entries[name] = {pos, size}; + } + std::uint64_t padded_size = (size + block_size - 1) / block_size * block_size; + pos += padded_size; + _stream.seekg(static_cast(pos)); + } + _stream.clear(); +} + +std::vector TarReader::names() const { + std::vector ret; + ret.reserve(_entries.size()); + for (auto& [name, entry] : _entries) { + ret.push_back(name); + } + return ret; +} + +bool TarReader::contains(const std::string& name) const { + return _entries.contains(name); +} + +const TarReader::Entry& TarReader::entry(const std::string& name) const { + auto search = _entries.find(name); + if (search == _entries.end()) { + throw std::runtime_error( + std::format("File '{}' not found in '{}'", name, _file_name) + ); + } + return search->second; +} + +std::size_t TarReader::size(const std::string& name) const { return entry(name).size; } + +void TarReader::read(const std::string& name, char* data) { + const Entry& e = entry(name); + _stream.clear(); + _stream.seekg(static_cast(e.offset)); + _stream.read(data, static_cast(e.size)); + if (!_stream) { + throw std::runtime_error( + std::format("Failed to read '{}' from '{}'", name, _file_name) + ); + } +} + +std::vector TarReader::read(const std::string& name) { + std::vector data(size(name)); + read(name, data.data()); + return data; +} diff --git a/madspace/src/driver/weighting_pipeline.cpp b/madspace/src/driver/weighting_pipeline.cpp new file mode 100644 index 0000000000..46b61dc0d8 --- /dev/null +++ b/madspace/src/driver/weighting_pipeline.cpp @@ -0,0 +1,697 @@ +#include "madspace/driver/weighting_pipeline.hpp" + +#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) { + NamedVector arg_types; + for (std::size_t part = 1; part <= 2; ++part) { + for (auto [key, type] : zip(types.keys(), types)) { + arg_types.push_back(std::format("arg{}_{}", part, key), type); + } + } + FunctionBuilder fb(arg_types, types); + int count = static_cast(types.size()); + for (int i = 0; i < count; ++i) { + auto [cat, sizes] = fb.batch_cat({fb.input(i), fb.input(i + count)}); + fb.output(i, cat); + } + return fb.function(); +} + +using json = nlohmann::json; + +json matrix_element_to_json(const MatrixElement& matrix_element) { + std::vector me_frame; + for (std::size_t i = 0; auto flag : matrix_element.frame_mask()) { + if (flag) { + me_frame.push_back(static_cast(i + 1)); + } + ++i; + } + // see the MatrixElement constructor: an empty reference mask with a frame + // means that the frame is the one of the incoming particles + auto& reference_mask = matrix_element.reference_mask(); + auto& incoming_mask = + reference_mask.empty() ? matrix_element.frame_mask() : reference_mask; + std::size_t incoming_count = incoming_mask.empty() + ? 2 + : std::count(incoming_mask.begin(), incoming_mask.end(), 1); + json inputs = json::array(), outputs = json::array(); + for (auto input : matrix_element.inputs()) { + inputs.push_back(static_cast(input)); + } + for (auto output : matrix_element.outputs()) { + outputs.push_back(static_cast(output)); + } + return { + {"index", matrix_element.matrix_element_index()}, + {"particle_count", matrix_element.particle_count()}, + {"inputs", inputs}, + {"outputs", outputs}, + {"diagram_count", matrix_element.diagram_count()}, + {"sample_random_inputs", matrix_element.sample_random_inputs()}, + {"me_frame", me_frame}, + {"incoming_count", incoming_count}, + }; +} + +MatrixElement matrix_element_from_json(const json& j) { + std::vector inputs; + std::vector outputs; + for (int input : j.at("inputs")) { + inputs.push_back(static_cast(input)); + } + for (int output : j.at("outputs")) { + outputs.push_back(static_cast(output)); + } + return MatrixElement( + j.at("index").get(), + j.at("particle_count").get(), + inputs, + outputs, + j.at("diagram_count").get(), + j.at("sample_random_inputs").get(), + j.at("me_frame").get>(), + j.at("incoming_count").get() + ); +} + +} // namespace + +std::string +madspace::weighting_groups_to_json(const std::vector& groups) { + json j = json::array(); + for (auto& group : groups) { + json steps = json::array(); + for (auto& step : group.steps) { + json matrix_elements = json::array(), override_outputs = json::array(); + for (auto& matrix_element : step.matrix_elements) { + matrix_elements.push_back(matrix_element_to_json(matrix_element)); + } + for (auto output : step.override_outputs) { + override_outputs.push_back(static_cast(output)); + } + steps.push_back({ + {"matrix_elements", matrix_elements}, + {"mode", + step.mode == WeightingStep::unweight ? "unweight" : "reweight"}, + {"max_overweight_truncation", step.max_overweight_truncation}, + {"freeze_max_weight_after", step.freeze_max_weight_after}, + {"override_outputs", override_outputs}, + }); + } + j.push_back({{"channel_indices", group.channel_indices}, {"steps", steps}}); + } + return j.dump(); +} + +std::vector +madspace::weighting_groups_from_json(const std::string& text) { + std::vector groups; + for (auto& group_json : json::parse(text)) { + WeightingGroup group{ + .channel_indices = + group_json.at("channel_indices").get>() + }; + for (auto& step_json : group_json.at("steps")) { + WeightingStep step{ + .mode = step_json.at("mode").get() == "unweight" + ? WeightingStep::unweight + : WeightingStep::reweight, + .max_overweight_truncation = + step_json.at("max_overweight_truncation").get(), + .freeze_max_weight_after = + step_json.at("freeze_max_weight_after").get(), + }; + for (auto& me_json : step_json.at("matrix_elements")) { + step.matrix_elements.push_back(matrix_element_from_json(me_json)); + } + for (int output : step_json.at("override_outputs")) { + step.override_outputs.push_back( + static_cast(output) + ); + } + group.steps.push_back(std::move(step)); + } + groups.push_back(std::move(group)); + } + return groups; +} + +void WeightingPipeline::Chunk::clear() { + source.reset(); + keep.reset(); + 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()), + _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); + Function concat = concat_function(types); + for (auto& context : contexts) { + ContextRuntimes runtimes{ + .gather = build_runtime(gather, context, false), + .concat = build_runtime(concat, context, false), + }; + for (auto& function : step_functions) { + runtimes.steps.push_back(build_runtime(function, context, false)); + } + group.runtimes.push_back(std::move(runtimes)); + } + group.buffers.assign( + contexts.size(), std::vector>(step_count + 1) + ); + group.buffer_sizes.assign( + contexts.size(), std::vector(step_count + 1) + ); + for (auto channel_index : group_config.channel_indices) { + channels.at(channel_index)->set_weighting_steps(group.settings); + _next_sequences.at(channel_index).assign(step_count + 1, 0); + _ready.at(channel_index).resize(step_count + 1); + } + _groups.push_back(std::move(group)); + ++group_index; + } +} + +std::size_t WeightingPipeline::batch_size(std::size_t context_index) const { + return _contexts.at(context_index)->device()->device_type() == DeviceType::cpu + ? _config.cpu_batch_size + : _config.gpu_batch_size; +} + +void WeightingPipeline::add_events(GeneratorBatchJob& job) { + auto& channel = _channels.at(job.channel_index); + auto chunk = std::make_shared(); + chunk->channel_index = job.channel_index; + chunk->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); + for (std::size_t i = 0; i < count; ++i) { + double previous = chunk.matrix_elements.at(i); + double& value = chunk.step_matrix_elements.at(i); + if (settings.ratio) { + // the implied matrix element is the reference of the next step + ratios.at(i) = value; + value *= previous; + } else { + ratios.at(i) = previous != 0. ? value / previous : 0.; + } + chunk.ratio_products.at(i) *= ratios.at(i); + } + std::vector keep; + 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) { + std::size_t started = 0; + for (std::size_t group_index = 0; auto& group : _groups) { + for (std::size_t context_index = 0; context_index < _contexts.size(); + ++context_index) { + auto& buffers = group.buffers.at(context_index); + auto& sizes = group.buffer_sizes.at(context_index); + for (std::size_t stage = 1; stage <= buffers.size(); ++stage) { + auto& buffer = buffers.at(stage - 1); + auto& size = sizes.at(stage - 1); + if (buffer.empty() || (!flush && size < batch_size(context_index))) { + continue; + } + std::size_t job_id = job_id_flag | _job_counter++; + auto& job = + _jobs + .emplace( + job_id, + Job{ + .group = group_index, + .stage = stage, + .context_index = context_index, + .chunks = std::move(buffer), + } + ) + .first->second; + buffer.clear(); + size = 0; + result_queue.submit( + _contexts.at(context_index)->thread_pool(), job_id, [this, &job] { + run_job(job); + } + ); + ++started; + } + } + ++group_index; + } + return started; +} + +void WeightingPipeline::gather_events(Chunk& chunk, ContextRuntimes& runtimes) { + auto& source = *chunk.source; + std::size_t count = chunk.size(); + if (!chunk.keep && chunk.offset == 0 && source.at(0).size(0) == count) { + chunk.events = source; + } else { + Tensor indices(DataType::dt_int, {count}); + auto view = indices.view(); + for (std::size_t i = 0; i < count; ++i) { + view[i] = static_cast( + chunk.offset + (chunk.keep ? chunk.keep->at(i) : i) + ); + } + auto device = _contexts.at(chunk.context_index)->device(); + if (device != cpu_device()) { + Tensor device_indices(DataType::dt_int, {count}, device); + device_indices.copy_from(indices); + indices = device_indices; + } + TensorVec args{indices}; + args.insert(args.end(), source.begin(), source.end()); + chunk.events = runtimes.gather->run(args); + } + chunk.source.reset(); + chunk.keep.reset(); + chunk.offset = 0; +} + +void WeightingPipeline::run_job(Job& job) { + auto& group = _groups.at(job.group); + auto& runtimes = group.runtimes.at(job.context_index); + for (auto& chunk : job.chunks) { + gather_events(*chunk, runtimes); + } + if (job.stage > group.settings.size()) { + // after the last step, the events go to the CPU to be written + for (auto& chunk : job.chunks) { + for (auto& tensor : chunk->events) { + chunk->cpu_events.push_back(tensor.cpu()); + } + chunk->events.clear(); + } + return; + } + + std::vector parts; + for (auto& chunk : job.chunks) { + parts.push_back(std::move(chunk->events)); + chunk->events.clear(); + } + // concatenate pairwise, so every event is copied log(chunks) times + while (parts.size() > 1) { + std::vector next_parts; + for (std::size_t i = 0; i < parts.size(); i += 2) { + if (i + 1 == parts.size()) { + next_parts.push_back(std::move(parts.at(i))); + } else { + TensorVec args = parts.at(i); + args.insert(args.end(), parts.at(i + 1).begin(), parts.at(i + 1).end()); + next_parts.push_back(runtimes.concat->run(args)); + } + } + parts = std::move(next_parts); + } + TensorVec events = std::move(parts.at(0)); + + TensorVec step_args{ + events.at(group.momenta), events.at(group.alpha_s), events.at(group.flavor) + }; + if (group.multiple_matrix_elements.at(job.stage - 1)) { + step_args.push_back(events.at(group.me_index)); + } + step_args.push_back(events.at(group.color)); + step_args.push_back(events.at(group.helicity)); + step_args.push_back(events.at(group.diagram)); + step_args.push_back(events.at(group.random)); + TensorVec outputs = runtimes.steps.at(job.stage - 1)->run(step_args); + auto matrix_elements = to_vector(outputs.at(0).cpu()); + auto random = to_vector(outputs.at(4).cpu()); + events.at(group.color) = outputs.at(1); + events.at(group.helicity) = outputs.at(2); + events.at(group.diagram) = outputs.at(3); + + auto shared_events = std::make_shared(std::move(events)); + std::size_t offset = 0; + for (auto& chunk : job.chunks) { + std::size_t count = chunk->size(); + chunk->source = shared_events; + chunk->offset = offset; + chunk->step_matrix_elements.assign( + matrix_elements.begin() + offset, matrix_elements.begin() + offset + count + ); + chunk->step_random.assign( + random.begin() + offset, random.begin() + offset + count + ); + offset += count; + } +} + +void WeightingPipeline::commit_job(std::size_t job_id) { + auto node = _jobs.extract(job_id); + auto& job = node.mapped(); + for (auto& chunk : job.chunks) { + set_ready(chunk, job.stage); + } +} + +void WeightingPipeline::reset() { + _jobs.clear(); + for (auto& group : _groups) { + for (auto& buffers : group.buffers) { + for (auto& buffer : buffers) { + buffer.clear(); + } + } + for (auto& sizes : group.buffer_sizes) { + std::fill(sizes.begin(), sizes.end(), 0); + } + } + for (auto& channel_ready : _ready) { + for (auto& ready : channel_ready) { + ready.clear(); + } + } + // the sequences continue, as the dropped chunks are never committed + for (std::size_t channel_index = 0; channel_index < _channels.size(); + ++channel_index) { + for (auto& next_sequence : _next_sequences.at(channel_index)) { + next_sequence = _channel_sequences.at(channel_index); + } + } + _chunk_count = 0; +} diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu index 8e9a06a218..96260efa45 100644 --- a/madspace/src/gpu/runtime.cu +++ b/madspace/src/gpu/runtime.cu @@ -40,8 +40,8 @@ void op_matrix_element( std::size_t output_count = locals[instruction.input_indices[2]].index_value(); TensorVec contiguous_inputs(input_count); std::vector input_keys(input_count); - // one more output: the stream (UMAMI_OUT_GPU_STREAM), so that the matrix element runs - // in stream order with the rest of the function, not on the default stream + // one more output: the stream (UMAMI_OUT_GPU_STREAM), so that the matrix element + // runs in stream order with the rest of the function, not on the default stream std::vector output_keys(output_count + 1); std::vector input_ptrs(input_count), output_ptrs(output_count + 1); for (std::size_t i = 0; i < input_count; ++i) { @@ -76,6 +76,7 @@ void op_matrix_element( // implemented at LUMI hackathon where the coffee was indeed terrible switch (output_keys[i]) { case UMAMI_OUT_MATRIX_ELEMENT: + case UMAMI_OUT_MATRIX_ELEMENT_RATIO: thrust::fill_n( thrust_par.on(device.stream()), thrust::device_pointer_cast(static_cast(output_ptrs[i])), @@ -1408,7 +1409,8 @@ void op_histogram( Tensor square_weights_tmp( DataType::dt_float, {padded_size}, device, AllocHint::temporary ); - // the keys reduce_by_key writes: the n_bins bins, the underflow and the overflow bin + // the keys reduce_by_key writes: the n_bins bins, the underflow and the overflow + // bin Tensor reduce_tmp(DataType::dt_float, {n_bins + 2}, device, AllocHint::temporary); launch_kernel( diff --git a/madspace/src/phasespace/cross_section.cpp b/madspace/src/phasespace/cross_section.cpp index 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/custom_function.cpp b/madspace/src/phasespace/custom_function.cpp new file mode 100644 index 0000000000..44c4f5ce01 --- /dev/null +++ b/madspace/src/phasespace/custom_function.cpp @@ -0,0 +1,104 @@ +#include "madspace/phasespace/custom_function.hpp" + +#include "madspace/phasespace/matrix_element.hpp" + +#include + +using namespace madspace; + +namespace { + +const char* output_name(CustomFunction::CustomFunctionType type) { + switch (type) { + case CustomFunction::cut: + return "cut"; + case CustomFunction::scale: + return "scale"; + case CustomFunction::bias: + return "bias"; + } + throw std::invalid_argument("unknown custom function type"); +} + +MatrixElement::MatrixElementOutput output_key(CustomFunction::CustomFunctionType type) { + switch (type) { + case CustomFunction::cut: + return MatrixElement::cut_out; + case CustomFunction::scale: + return MatrixElement::scale_out; + case CustomFunction::bias: + return MatrixElement::bias_out; + } + throw std::invalid_argument("unknown custom function type"); +} + +} // namespace + +CustomFunction::CustomFunction( + std::size_t function_index, + CustomFunctionType type, + std::size_t particle_count, + const nested_vector2& pdg_id_options +) : + FunctionGenerator( + "CustomFunction", + [&] { + NamedVector arg_types{ + {"momenta", batch_four_vec_array(particle_count)} + }; + if (pdg_id_options.size() > 1) { + arg_types.push_back("flavor_index", batch_int); + } + return arg_types; + }(), + {{output_name(type), batch_float}} + ), + _function_index(function_index), + _type(type), + _particle_count(particle_count), + _pdg_id_options(pdg_id_options) { + for (auto& option : pdg_id_options) { + if (option.size() != particle_count) { + throw std::invalid_argument( + std::format( + "custom function: expected {} PDG ids per option, got {}", + particle_count, + option.size() + ) + ); + } + } +} + +NamedVector CustomFunction::build_function_impl( + FunctionBuilder& fb, const NamedVector& args +) const { + auto momenta = args.at(0); + auto batch_size = fb.batch_size({momenta}); + ValueVec pdg_ids; + for (std::size_t i = 0; i < _particle_count; ++i) { + if (_pdg_id_options.empty()) { + pdg_ids.push_back(fb.full({static_cast(0), batch_size})); + } else if (_pdg_id_options.size() == 1) { + pdg_ids.push_back(fb.full({_pdg_id_options.at(0).at(i), batch_size})); + } else { + std::vector column; + for (auto& option : _pdg_id_options) { + column.push_back(option.at(i)); + } + pdg_ids.push_back(fb.gather_int(args.at(1), column)); + } + } + MatrixElement function( + _function_index, + _particle_count, + {MatrixElement::momenta_in, + MatrixElement::particle_count_in, + MatrixElement::pdg_ids_in}, + {output_key(_type)} + ); + auto particle_count = fb.full({static_cast(_particle_count), batch_size}); + auto value = + function.build_function(fb, {momenta, particle_count, fb.stack(pdg_ids)}).at(0); + return {{output_name(_type), value}}; +} diff --git a/madspace/src/phasespace/cuts.cpp b/madspace/src/phasespace/cuts.cpp index 69795e5814..64a408f4ce 100644 --- a/madspace/src/phasespace/cuts.cpp +++ b/madspace/src/phasespace/cuts.cpp @@ -7,18 +7,42 @@ using namespace madspace; -Cuts::Cuts(const std::vector& cut_data) : +namespace { + +void check_custom_cut(const std::optional& custom_cut) { + if (custom_cut && + (custom_cut->type() != CustomFunction::cut || custom_cut->dynamic_pdg_ids())) { + throw std::invalid_argument( + "custom cut must be of type cut, with a single PDG-id option" + ); + } +} + +} // namespace + +Cuts::Cuts( + const std::vector& cut_data, + const std::optional& custom_cut +) : FunctionGenerator( "Cuts", cut_data.at(0).observable.arg_types(), {{"mask", batch_float}} ), - _cut_data(cut_data) {} + _cut_data(cut_data), + _custom_cut(custom_cut) { + check_custom_cut(custom_cut); +} -Cuts::Cuts(std::size_t particle_count) : +Cuts::Cuts( + std::size_t particle_count, const std::optional& custom_cut +) : FunctionGenerator( "Cuts", {{"momenta", batch_four_vec_array(particle_count)}}, {{"mask", batch_float}} - ) {} + ), + _custom_cut(custom_cut) { + check_custom_cut(custom_cut); +} NamedVector Cuts::build_function_impl(FunctionBuilder& fb, const NamedVector& args) const { @@ -36,6 +60,9 @@ Cuts::build_function_impl(FunctionBuilder& fb, const NamedVector& args) c weights.push_back(fb.cut_any(obs, item.min, item.max)); } } + if (_custom_cut) { + weights.push_back(_custom_cut->build_function(fb, {args.at(0)}).at(0)); + } return {{"mask", fb.product(weights)}}; } @@ -168,9 +195,8 @@ PairList delta_r_pairs(const Observable& o) { } // namespace -std::vector> Cuts::pairwise_min( - const std::function& pairs -) const { +std::vector> +Cuts::pairwise_min(const std::function& pairs) const { std::size_t n = arg_types().at(0).shape.at(0) - 2; std::vector> out(n, std::vector(n, 0.)); for (auto& item : _cut_data) { diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index 379c3bbdae..5e9893e8ef 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -3,6 +3,7 @@ #include "madspace/util.hpp" #include +#include using namespace madspace; @@ -30,6 +31,16 @@ std::size_t final_channel_count( } } +/// Position of `output` in the return values of a differential cross section +/// built on `matrix_element`, or -1 if it is not returned. +int me_output_index( + const MatrixElement& matrix_element, MatrixElement::MatrixElementOutput output +) { + auto& outputs = matrix_element.outputs(); + auto search = std::find(outputs.begin(), outputs.end(), output); + return search == outputs.end() ? -1 : static_cast(search - outputs.begin()); +} + } // namespace static const BatchSize acc_batch_size("acc_batch_size"); @@ -41,7 +52,7 @@ Integrand::Integrand( const AdaptiveDiscrete& discrete_sym, const AdaptiveDiscrete& discrete_flavor, const nested_vector2& pid_options, - const std::optional& pdf_grid, + const PdfSource& pdf_grid, const std::optional& running_coupling, const std::optional& energy_scale, const std::optional& prop_chan_weights, @@ -62,7 +73,8 @@ Integrand::Integrand( const std::vector& flavor_diff_xs_indices, const std::vector& flavor_subproc_indices, const std::vector& flavor_per_subproc_remap, - std::size_t compressed_channel_weight_count + std::size_t compressed_channel_weight_count, + const std::optional& custom_bias ) : FunctionGenerator( "Integrand", @@ -117,7 +129,8 @@ Integrand::Integrand( if (flav_count > 1 && !std::holds_alternative(discrete_flavor)) { ret_types.push_back("discrete_flavor_index", batch_int); - if (pdf_grid && energy_scale) { + if (!std::holds_alternative(pdf_grid) && + energy_scale) { ret_types.push_back("pdf_prior", batch_float_array(flav_count)); } } @@ -148,6 +161,17 @@ Integrand::Integrand( ret_types.push_back("partial_weight_product", batch_float); } } + if (diff_xs_first.returns_matrix_element()) { + ret_types.push_back("matrix_element", batch_float); + 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); + } ret_types.push_back("random", batch_float_array(mapping.random_dim())); if (mapping.channel_count() > 1 && !std::holds_alternative(discrete_sym)) { @@ -198,8 +222,9 @@ Integrand::Integrand( ), _flavor_per_subproc_remap( flavor_per_subproc_remap.begin(), flavor_per_subproc_remap.end() - ) { - if (pdf_grid) { + ), + _custom_bias(custom_bias) { + if (!std::holds_alternative(pdf_grid)) { for (std::size_t i = 0; i < 2; ++i) { std::set pids; for (auto& option : pid_options) { @@ -210,7 +235,24 @@ Integrand::Integrand( std::distance(pids.begin(), pids.find(option.at(i))) ); } - _pdfs.at(i) = PartonDensity(pdf_grid.value(), {pids.begin(), pids.end()}); + std::vector pid_vec(pids.begin(), pids.end()); + if (auto grid = std::get_if(&pdf_grid)) { + _pdfs.at(i) = PartonDensity(*grid, pid_vec); + } else { + _pdfs.at(i) = PartonDensity( + std::get(pdf_grid), pid_vec, static_cast(i) + ); + } + } + } + if (custom_bias) { + if (custom_bias->type() != CustomFunction::bias) { + throw std::invalid_argument("custom bias must be of type bias"); + } + if (custom_bias->dynamic_pdg_ids() != (pid_options.size() > 1)) { + throw std::invalid_argument( + "the PDG-id options of the custom bias must match pid_options" + ); } } @@ -241,6 +283,41 @@ Integrand::Integrand( _flavor_mirror.push_back(mirror ? 2 : 1); _has_mirror |= mirror; } + auto& me_first = diff_xs.at(0).matrix_element(); + for (auto& dxs : diff_xs) { + if (dxs.matrix_element().outputs() != me_first.outputs() || + dxs.matrix_element().external_inputs() != me_first.external_inputs() || + dxs.returns_matrix_element() != diff_xs.at(0).returns_matrix_element()) { + throw std::invalid_argument( + "all matrix elements of an integrand need the same inputs and outputs" + ); + } + } + for (auto input : me_first.external_inputs()) { + if (input != MatrixElement::momenta_in && input != MatrixElement::flavor_in && + input != MatrixElement::alpha_s_in) { + throw std::invalid_argument( + "the integrand only supplies momenta, flavor and alpha_s to the " + "matrix element, the random inputs have to be sampled internally" + ); + } + } + std::size_t channel_count = final_channel_count( + diff_xs, + first_chan_weight_remap, + first_remapped_chan_count, + second_chan_weight_remap, + second_remapped_chan_count, + prop_chan_weights, + subchan_weights + ); + if (channel_count > 1 && !prop_chan_weights && + me_output_index(me_first, MatrixElement::diagram_amp2_out) == -1) { + throw std::invalid_argument( + "diagram-based channel weights need the diagram_amp2 output of the " + "matrix element" + ); + } _channel_part_ret_types = compute_channel_part_ret_types(); } @@ -700,10 +777,18 @@ NamedVector Integrand::build_common_part( _running_coupling.value().build_function(fb, {args.at("ren_scale")}).at(0); // Evaluate differential cross section - ValueVec xs_args{ - momenta_me, - _flavor_remap.size() > 0 ? fb.gather_int(flavor_id, _flavor_remap) : flavor_id, - }; + auto& me_first = _diff_xs.at(0).matrix_element(); + int amp2_index = me_output_index(me_first, MatrixElement::diagram_amp2_out); + Value me_flavor_id = + _flavor_remap.size() > 0 ? fb.gather_int(flavor_id, _flavor_remap) : flavor_id; + ValueVec xs_args; + for (auto input : me_first.external_inputs()) { + if (input == MatrixElement::momenta_in) { + xs_args.push_back(momenta_me); + } else if (input == MatrixElement::flavor_in) { + xs_args.push_back(me_flavor_id); + } + } xs_args.push_back(x1_acc); xs_args.push_back(x2_acc); xs_args.push_back(flavor_id); @@ -717,11 +802,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); + 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, @@ -731,7 +817,7 @@ NamedVector Integrand::build_common_part( } } } else { - Value dxs_index = fb.gather_int(flavor_id, _flavor_diff_xs_indices); + dxs_index = fb.gather_int(flavor_id, _flavor_diff_xs_indices); ps_flavor_id = fb.gather_int(flavor_id, _flavor_per_subproc_remap); subproc_id = fb.gather_int(flavor_id, _flavor_subproc_indices); ValueVec split_indices = @@ -750,7 +836,7 @@ NamedVector Integrand::build_common_part( split_out.push_back(indices); } if (channel_count > 1 && !_prop_chan_weights) { - Value split_cw = outputs.at(1); + Value split_cw = outputs.at(amp2_index); if (_first_chan_weight_remap.size() > 0) { split_cw = fb.collect_channel_weights( split_cw, @@ -764,7 +850,7 @@ NamedVector Integrand::build_common_part( ++i; } for (std::size_t i = 0; auto& split_out : split_outputs) { - if (i == 1) { + if (std::cmp_equal(i, amp2_index)) { dxs_vec.push_back(Value()); } else { dxs_vec.push_back(fb.batch_merge_by_index(split_out)); @@ -780,6 +866,17 @@ NamedVector Integrand::build_common_part( diff_xs_acc = fb.mul(diff_xs_acc, fb.gather(flavor_id, _flavor_factors)); } ValueVec weights_after_cuts{args.at("weight_after_cuts"), diff_xs_acc}; + // the bias enters the weight seen by the adaptation and the unweighting, + // the event generator divides it out again + Value bias_acc; + if (_custom_bias) { + ValueVec bias_args{momenta_me}; + if (_custom_bias->dynamic_pdg_ids()) { + bias_args.push_back(flavor_id); + } + bias_acc = _custom_bias->build_function(fb, bias_args).at(0); + weights_after_cuts.push_back(bias_acc); + } ValueVec extra_weights_after_cuts; if (args.index_map().contains("extra_weight_after_cuts")) { extra_weights_after_cuts.push_back(args.at("extra_weight_after_cuts")); @@ -929,9 +1026,20 @@ NamedVector Integrand::build_common_part( } auto zeros_int = fb.full({static_cast(0), batch_size_val}); auto zeros_float = fb.full({0., batch_size_val}); - outputs.push_back("color_index", scatter_or_drop(zeros_int, dxs_vec.at(2))); - outputs.push_back("helicity_index", scatter_or_drop(zeros_int, dxs_vec.at(3))); - outputs.push_back("diagram_index", scatter_or_drop(zeros_int, dxs_vec.at(4))); + // indices the matrix element does not provide are -1, a later + // reweighting step can fill them in + for (auto [name, me_output] : + {std::pair{"color_index", MatrixElement::color_index_out}, + std::pair{"helicity_index", MatrixElement::helicity_index_out}, + std::pair{"diagram_index", MatrixElement::diagram_index_out}}) { + int index = me_output_index(me_first, me_output); + outputs.push_back( + name, + index == -1 + ? optional_cut(fb.full({static_cast(-1), batch_size_val})) + : scatter_or_drop(zeros_int, dxs_vec.at(index)) + ); + } outputs.push_back("flavor_index", scatter_or_drop(zeros_int, ps_flavor_id)); if (subproc_id) { outputs.push_back( @@ -970,6 +1078,25 @@ NamedVector Integrand::build_common_part( ); } } + if (_diff_xs.at(0).returns_matrix_element()) { + outputs.push_back( + "matrix_element", + scatter_or_drop(zeros_float, dxs_vec.at(me_first.outputs().size())) + ); + outputs.push_back( + "matrix_element_flavor", scatter_or_drop(zeros_int, me_flavor_id) + ); + if (dxs_index) { + outputs.push_back( + "matrix_element_index", scatter_or_drop(zeros_int, dxs_index) + ); + } + } + if (bias_acc) { + outputs.push_back( + "bias", scatter_or_drop(fb.full({1., batch_size_val}), bias_acc) + ); + } outputs.push_back("random", optional_cut(args.at("r"))); if (has_permutations && !std::holds_alternative(_discrete_sym)) { diff --git a/madspace/src/phasespace/matrix_element.cpp b/madspace/src/phasespace/matrix_element.cpp index 851bfb8c09..c0ef852d4c 100644 --- a/madspace/src/phasespace/matrix_element.cpp +++ b/madspace/src/phasespace/matrix_element.cpp @@ -10,19 +10,20 @@ namespace { /// A per-particle 0/1 selector, 1 for every one-based particle number in /// `indices`. `what` names the setting the numbers came from, for the error. std::vector particle_mask( - std::size_t particle_count, - const std::vector& indices, - const char* what + std::size_t particle_count, const std::vector& indices, const char* what ) { std::vector mask(particle_count, 0); for (me_int_t index : indices) { if (index < 1 || std::cmp_greater(index, particle_count)) { - throw std::invalid_argument(std::format( - "{} particle {} out of range, this process has {} external particles", - what, - index, - particle_count - )); + throw std::invalid_argument( + std::format( + "{} particle {} out of range, this process has {} external " + "particles", + what, + index, + particle_count + ) + ); } mask.at(index - 1) = 1; } @@ -45,14 +46,12 @@ MatrixElement::MatrixElement( "MatrixElement", [&] { NamedVector arg_types; - bool found_momenta = false; for (auto input : inputs) { switch (input) { case momenta_in: arg_types.push_back( "momenta", batch_four_vec_array(particle_count) ); - found_momenta = true; break; case alpha_s_in: arg_types.push_back("alpha_s", batch_float); @@ -84,12 +83,30 @@ MatrixElement::MatrixElement( case channel_in: arg_types.push_back("channel", batch_int); break; + case particle_count_in: + arg_types.push_back("particle_count", batch_int); + break; + case pdg_ids_in: + arg_types.push_back("pdg_ids", batch_int_array(particle_count)); + break; + case momentum_fraction_in: + arg_types.push_back("momentum_fraction", batch_float); + break; + case scale_in: + arg_types.push_back("scale", batch_float); + break; + case pdg_id_in: + arg_types.push_back("pdg_id", batch_int); + break; + case beam_index_in: + arg_types.push_back("beam_index", batch_int); + break; default: throw std::invalid_argument("unknown input type"); } } - if (!found_momenta) { - throw std::invalid_argument("missing momentum input"); + if (arg_types.size() == 0) { + throw std::invalid_argument("matrix element needs at least one input"); } return arg_types; }(), @@ -114,6 +131,21 @@ MatrixElement::MatrixElement( case diagram_index_out: ret_types.push_back("diagram_index", batch_int); break; + case cut_out: + ret_types.push_back("cut", batch_float); + break; + case scale_out: + ret_types.push_back("scale", batch_float); + break; + case pdf_out: + ret_types.push_back("pdf", batch_float); + break; + case bias_out: + ret_types.push_back("bias", batch_float); + break; + case matrix_element_ratio_out: + ret_types.push_back("matrix_element_ratio", batch_float); + break; default: throw std::invalid_argument("unknown output type"); } @@ -154,7 +186,17 @@ MatrixElement::MatrixElement( mask.clear(); } return mask; - }()) {} + }()) { + auto has_output = [&](MatrixElementOutput output) { + return std::find(outputs.begin(), outputs.end(), output) != outputs.end(); + }; + if (has_output(matrix_element_out) && has_output(matrix_element_ratio_out)) { + throw std::invalid_argument( + "a matrix element returns either matrix_element or matrix_element_ratio, " + "not both" + ); + } +} NamedVector MatrixElement::build_function_impl( FunctionBuilder& fb, const NamedVector& args @@ -207,6 +249,24 @@ NamedVector MatrixElement::build_function_impl( case channel_in: input_key = UMAMI_IN_CHANNEL_INDEX; break; + case particle_count_in: + input_key = UMAMI_IN_PARTICLE_COUNT; + break; + case pdg_ids_in: + input_key = UMAMI_IN_PDG_IDS; + break; + case momentum_fraction_in: + input_key = UMAMI_IN_MOMENTUM_FRACTION; + break; + case scale_in: + input_key = UMAMI_IN_SCALE; + break; + case pdg_id_in: + input_key = UMAMI_IN_PDG_ID; + break; + case beam_index_in: + input_key = UMAMI_IN_BEAM_INDEX; + break; } matrix_args.push_back(static_cast(input_key)); if (_sample_random_inputs && @@ -261,6 +321,21 @@ NamedVector MatrixElement::build_function_impl( case diagram_index_out: output_key = UMAMI_OUT_DIAGRAM_INDEX; break; + case cut_out: + output_key = UMAMI_OUT_CUT; + break; + case scale_out: + output_key = UMAMI_OUT_SCALE; + break; + case pdf_out: + output_key = UMAMI_OUT_PDF; + break; + case bias_out: + output_key = UMAMI_OUT_BIAS; + break; + case matrix_element_ratio_out: + output_key = UMAMI_OUT_MATRIX_ELEMENT_RATIO; + break; } matrix_args.push_back(static_cast(output_key)); matrix_args.push_back(size_arg); diff --git a/madspace/src/phasespace/pdf.cpp b/madspace/src/phasespace/pdf.cpp index 0a437f341e..fc3a28025c 100644 --- a/madspace/src/phasespace/pdf.cpp +++ b/madspace/src/phasespace/pdf.cpp @@ -1,5 +1,7 @@ #include "madspace/phasespace/pdf.hpp" +#include "madspace/phasespace/matrix_element.hpp" + #include #include #include @@ -384,11 +386,60 @@ PartonDensity::PartonDensity( } } +PartonDensity::PartonDensity( + const CustomPdf& custom_pdf, + const std::vector& pids, + int beam_index, + bool dynamic_pid +) : + FunctionGenerator( + "PartonDensity", + [&] { + NamedVector arg_types{{"x", batch_float}, {"q", batch_float}}; + if (dynamic_pid) { + arg_types.push_back("flavor_index", batch_int); + } + return arg_types; + }(), + {{"pdf", dynamic_pid ? batch_float : batch_float_array(pids.size())}} + ), + _pids(pids.begin(), pids.end()), + _custom_pdf(custom_pdf), + _beam_index(beam_index), + _dynamic_pid(dynamic_pid) {} + NamedVector PartonDensity::build_function_impl( FunctionBuilder& fb, const NamedVector& args ) const { auto x = args.at(0); auto q = args.at(1); + if (_custom_pdf) { + MatrixElement function( + _custom_pdf->function_index, + 0, + {MatrixElement::momentum_fraction_in, + MatrixElement::scale_in, + MatrixElement::pdg_id_in, + MatrixElement::beam_index_in}, + {MatrixElement::pdf_out} + ); + auto batch_size = fb.batch_size({x}); + auto beam_index = fb.full({_beam_index, batch_size}); + if (_dynamic_pid) { + auto pid = fb.gather_int(args.at(2), _pids); + return { + {"pdf", function.build_function(fb, {x, q, pid, beam_index}).at(0)} + }; + } + ValueVec pdfs; + for (me_int_t pid : _pids) { + auto pid_value = fb.full({pid, batch_size}); + pdfs.push_back( + function.build_function(fb, {x, q, pid_value, beam_index}).at(0) + ); + } + return {{"pdf", fb.stack(pdfs)}}; + } auto grid_logx = fb.global( prefixed_name(_prefix, "pdf_logx"), DataType::dt_float, diff --git a/madspace/src/phasespace/scale.cpp b/madspace/src/phasespace/scale.cpp index 46715ca972..ddcc82524d 100644 --- a/madspace/src/phasespace/scale.cpp +++ b/madspace/src/phasespace/scale.cpp @@ -10,7 +10,8 @@ EnergyScale::EnergyScale( double ren_scale, double fact_scale1, double fact_scale2, - double scale_factor + double scale_factor, + const std::optional& custom_scale ) : FunctionGenerator( "EnergyScale", @@ -25,7 +26,16 @@ EnergyScale::EnergyScale( _ren_scale(ren_scale), _fact_scale1(fact_scale1), _fact_scale2(fact_scale2), - _scale_factor(scale_factor) {} + _scale_factor(scale_factor), + _custom_scale(custom_scale) { + if (custom_scale && + (custom_scale->type() != CustomFunction::scale || + custom_scale->dynamic_pdg_ids())) { + throw std::invalid_argument( + "custom scale must be of type scale, with a single PDG-id option" + ); + } +} NamedVector EnergyScale::build_function_impl( FunctionBuilder& fb, const NamedVector& args @@ -40,21 +50,25 @@ NamedVector EnergyScale::build_function_impl( }; } Value scale; - switch (_dynamical_scale_type) { - case transverse_energy: - scale = fb.scale_transverse_energy(momenta); - break; - case transverse_mass: - scale = fb.scale_transverse_mass(momenta); - break; - case half_transverse_mass: - scale = fb.scale_half_transverse_mass(momenta); - break; - case partonic_energy: - scale = fb.scale_partonic_energy(momenta); - break; - default: - throw std::runtime_error("invalid dynamical scale type"); + if (_custom_scale) { + scale = _custom_scale->build_function(fb, {momenta}).at(0); + } else { + switch (_dynamical_scale_type) { + case transverse_energy: + scale = fb.scale_transverse_energy(momenta); + break; + case transverse_mass: + scale = fb.scale_transverse_mass(momenta); + break; + case half_transverse_mass: + scale = fb.scale_half_transverse_mass(momenta); + break; + case partonic_energy: + scale = fb.scale_partonic_energy(momenta); + break; + default: + throw std::runtime_error("invalid dynamical scale type"); + } } // dynamical scale only; a fixed scale is an absolute value if (_scale_factor != 1.) { diff --git a/madspace/src/phasespace/weighting_step.cpp b/madspace/src/phasespace/weighting_step.cpp new file mode 100644 index 0000000000..c0e75f7075 --- /dev/null +++ b/madspace/src/phasespace/weighting_step.cpp @@ -0,0 +1,230 @@ +#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" + ); + } + } + if (matrix_element.external_inputs().size() != matrix_element.inputs().size()) { + throw std::invalid_argument( + "the random inputs of a weighting step matrix element are passed in " + "and must not be sampled internally" + ); + } + } +} + +NamedVector WeightingStepFunction::build_function_impl( + FunctionBuilder& fb, const NamedVector& args +) const { + ValueVec all_random = fb.unstack(args.at("random")); + ValueVec random( + all_random.begin() + random_count * _step_index, + all_random.begin() + random_count * (_step_index + 1) + ); + ValueVec indices{ + args.at("color_index"), args.at("helicity_index"), args.at("diagram_index") + }; + auto me_args = [&](const MatrixElement& matrix_element, auto&& arg) { + ValueVec ret; + for (auto input : matrix_element.inputs()) { + switch (input) { + case MatrixElement::momenta_in: + ret.push_back(arg(args.at("momenta"))); + break; + case MatrixElement::alpha_s_in: + ret.push_back(arg(args.at("alpha_s"))); + break; + case MatrixElement::flavor_in: + ret.push_back(arg(args.at("flavor"))); + break; + case MatrixElement::random_color_in: + ret.push_back(arg(random.at(0))); + break; + case MatrixElement::random_helicity_in: + ret.push_back(arg(random.at(1))); + break; + case MatrixElement::random_diagram_in: + ret.push_back(arg(random.at(2))); + break; + case MatrixElement::helicity_in: + ret.push_back(arg(indices.at(1))); + break; + case MatrixElement::diagram_in: + ret.push_back(arg(indices.at(2))); + break; + default: + break; + } + } + return ret; + }; + // matrix element plus the index outputs, missing ones are empty + auto select_outputs = + [&](const MatrixElement& matrix_element, const ValueVec& outputs) { + 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/function_runtime.cpp b/madspace/src/python/function_runtime.cpp index 161c4e5ec5..16169c420c 100644 --- a/madspace/src/python/function_runtime.cpp +++ b/madspace/src/python/function_runtime.cpp @@ -580,6 +580,8 @@ Tensor madspace_py::dlpack_to_tensor( deleter = [managed_ptr]() { auto ptr = static_cast(managed_ptr); if (ptr->deleter) { + // the producer may be freed on a worker thread without the GIL + PythonGuard guard; ptr->deleter(ptr); } }; @@ -587,6 +589,7 @@ Tensor madspace_py::dlpack_to_tensor( deleter = [managed_ptr]() { auto ptr = static_cast(managed_ptr); if (ptr->deleter) { + PythonGuard guard; ptr->deleter(ptr); } }; @@ -618,6 +621,8 @@ std::vector FunctionRuntime::call(std::vector args) { ProducerStream producer_stream(args); auto [inputs, runtime] = check_and_convert_args(args, *this, &_dlpack_version_cache); + // a Python matrix element called from a worker thread needs the GIL + py::gil_scoped_release release; return runtime->run(inputs); } @@ -628,8 +633,10 @@ FunctionRuntime::call_with_grad( ProducerStream producer_stream(args); auto [inputs, runtime] = check_and_convert_args(args, *this, &_dlpack_version_cache); - auto [outputs, loc_grad, eval_grad] = - runtime->run_with_grad(inputs, input_requires_grad); + auto [outputs, loc_grad, eval_grad] = [&] { + py::gil_scoped_release release; + return runtime->run_with_grad(inputs, input_requires_grad); + }(); std::vector> local_grads; for (auto& grad : loc_grad) { if (grad) { @@ -689,8 +696,10 @@ FunctionRuntime::call_backward( } // TODO: checks here Runtime* runtime = get_runtime(*this, expected_device); - auto [ret_in_grads, ret_glob_grads] = - runtime->run_backward(arg_out, arg_locals, eval_grad); + auto [ret_in_grads, ret_glob_grads] = [&] { + py::gil_scoped_release release; + return runtime->run_backward(arg_out, arg_locals, eval_grad); + }(); std::vector> input_grads; for (auto& grad : ret_in_grads) { if (grad) { diff --git a/madspace/src/python/function_runtime.hpp b/madspace/src/python/function_runtime.hpp index 6d189be7d4..e0ee582443 100644 --- a/madspace/src/python/function_runtime.hpp +++ b/madspace/src/python/function_runtime.hpp @@ -13,6 +13,28 @@ using namespace madspace; namespace madspace_py { +/// Holds the GIL while alive, for code that can run on madspace worker threads. +/// Does nothing once the interpreter is finalized. +class PythonGuard { +public: + PythonGuard() : _active(Py_IsInitialized()) { + if (_active) { + _state = PyGILState_Ensure(); + } + } + ~PythonGuard() { + if (_active) { + PyGILState_Release(_state); + } + } + PythonGuard(const PythonGuard&) = delete; + PythonGuard& operator=(const PythonGuard&) = delete; + +private: + bool _active; + PyGILState_STATE _state; +}; + std::tuple dlpack_device(Tensor tensor); py::object tensor_to_dlpack( Tensor tensor, diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 5cae93612f..42ed9d1c83 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1,4 +1,5 @@ #include +#include #include #include #include @@ -10,6 +11,7 @@ #include "madspace/driver.hpp" #include "madspace/phasespace.hpp" #include "madspace/util.hpp" +#include "python_function.hpp" namespace py = pybind11; using namespace madspace; @@ -141,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) { @@ -357,6 +418,17 @@ PYBIND11_MODULE(_madspace_py, m) { py::return_value_policy::reference_internal, pydoc::doc("Context::load_matrix_element") ) + .def( + "_add_python_function", + &add_python_function, + py::arg("function"), + py::arg("input_keys"), + py::arg("required_input_keys"), + py::arg("output_keys"), + py::arg("framework"), + py::arg("name"), + py::return_value_policy::reference_internal + ) .def( "define_global", &Context::define_global, @@ -445,6 +517,52 @@ PYBIND11_MODULE(_madspace_py, m) { pydoc::doc("set_caller_stream") ); + m.def( + "_tensor_view", + &tensor_view, + py::arg("address"), + py::arg("dtype"), + py::arg("shape"), + py::arg("stride"), + py::arg("offset"), + py::arg("device") + ); + py::dict umami_input_keys, umami_output_keys; + for (auto [name, key] : std::initializer_list>{ + {"momenta", UMAMI_IN_MOMENTA}, + {"alpha_s", UMAMI_IN_ALPHA_S}, + {"flavor_index", UMAMI_IN_FLAVOR_INDEX}, + {"random_color", UMAMI_IN_RANDOM_COLOR}, + {"random_helicity", UMAMI_IN_RANDOM_HELICITY}, + {"random_diagram", UMAMI_IN_RANDOM_DIAGRAM}, + {"helicity_index", UMAMI_IN_HELICITY_INDEX}, + {"diagram_index", UMAMI_IN_DIAGRAM_INDEX}, + {"channel_index", UMAMI_IN_CHANNEL_INDEX}, + {"particle_count", UMAMI_IN_PARTICLE_COUNT}, + {"pdg_ids", UMAMI_IN_PDG_IDS}, + {"momentum_fraction", UMAMI_IN_MOMENTUM_FRACTION}, + {"scale", UMAMI_IN_SCALE}, + {"pdg_id", UMAMI_IN_PDG_ID}, + {"beam_index", UMAMI_IN_BEAM_INDEX}, + }) { + umami_input_keys[name] = key; + } + for (auto [name, key] : std::initializer_list>{ + {"matrix_element", UMAMI_OUT_MATRIX_ELEMENT}, + {"color_index", UMAMI_OUT_COLOR_INDEX}, + {"helicity_index", UMAMI_OUT_HELICITY_INDEX}, + {"diagram_index", UMAMI_OUT_DIAGRAM_INDEX}, + {"cut", UMAMI_OUT_CUT}, + {"scale", UMAMI_OUT_SCALE}, + {"pdf", UMAMI_OUT_PDF}, + {"bias", UMAMI_OUT_BIAS}, + {"matrix_element_ratio", UMAMI_OUT_MATRIX_ELEMENT_RATIO}, + }) { + umami_output_keys[name] = key; + } + m.attr("_UMAMI_INPUT_KEYS") = umami_input_keys; + m.attr("_UMAMI_OUTPUT_KEYS") = umami_output_keys; + py::classh(m, "FunctionRuntime", py::dynamic_attr()) .def(py::init(), py::arg("function")) .def(py::init(), py::arg("function"), py::arg("context")) @@ -955,6 +1073,71 @@ PYBIND11_MODULE(_madspace_py, m) { .def_readonly_static("missing_pids", &Observable::missing_pids) .def_readonly_static("photon_pids", &Observable::photon_pids); + py::classh custom_function( + m, "CustomFunction", pydoc::doc("CustomFunction") + ); + add_enum( + custom_function, + "CustomFunctionType", + { + {"cut", CustomFunction::cut}, + {"scale", CustomFunction::scale}, + {"bias", CustomFunction::bias}, + } + ); + custom_function + .def( + py::init< + std::size_t, + CustomFunction::CustomFunctionType, + std::size_t, + const nested_vector2&>(), + py::arg("function_index"), + py::arg("type"), + py::arg("particle_count"), + py::arg("pdg_id_options") = nested_vector2{}, + pydoc::doc("CustomFunction::CustomFunction") + ) + .def( + py::init< + const MatrixElementApi&, + CustomFunction::CustomFunctionType, + std::size_t, + const nested_vector2&>(), + py::arg("function_api"), + py::arg("type"), + py::arg("particle_count"), + py::arg("pdg_id_options") = nested_vector2{}, + pydoc::doc("CustomFunction::CustomFunction#2") + ) + .def( + "function_index", + &CustomFunction::function_index, + pydoc::doc("CustomFunction::function_index") + ) + .def("type", &CustomFunction::type, pydoc::doc("CustomFunction::type")) + .def( + "dynamic_pdg_ids", + &CustomFunction::dynamic_pdg_ids, + pydoc::doc("CustomFunction::dynamic_pdg_ids") + ); + py::classh(m, "CustomPdf", pydoc::doc("CustomPdf")) + .def( + py::init(), + py::arg("function_index"), + pydoc::doc("CustomPdf::CustomPdf") + ) + .def( + py::init(), + py::arg("function_api"), + pydoc::doc("CustomPdf::CustomPdf#2") + ) + .def_readonly( + "function_index", + &CustomPdf::function_index, + pydoc::doc("CustomPdf::function_index") + ); + auto cuts = py::classh(m, "Cuts", pydoc::doc("Cuts")); add_enum( cuts, @@ -991,13 +1174,17 @@ PYBIND11_MODULE(_madspace_py, m) { "min", &Cuts::PairMassAny::min, pydoc::doc("Cuts::PairMassAny::min") ); cuts.def( - py::init&>(), + py::init< + const std::vector&, + const std::optional&>(), py::arg("cut_data"), + py::arg("custom_cut") = std::nullopt, pydoc::doc("Cuts::Cuts") ) .def( - py::init(), + py::init&>(), py::arg("particle_count"), + py::arg("custom_cut") = std::nullopt, pydoc::doc("Cuts::Cuts#2") ) .def("sqrt_s_min", &Cuts::sqrt_s_min, pydoc::doc("Cuts::sqrt_s_min")) @@ -1313,7 +1500,9 @@ PYBIND11_MODULE(_madspace_py, m) { &PhaseSpaceMapping::channel_count, pydoc::doc("PhaseSpaceMapping::channel_count") ) - .def("empty", &PhaseSpaceMapping::empty, pydoc::doc("PhaseSpaceMapping::empty")); + .def( + "empty", &PhaseSpaceMapping::empty, pydoc::doc("PhaseSpaceMapping::empty") + ); py::classh( m, "MultiChannelFunction", pydoc::doc("MultiChannelFunction") @@ -1341,6 +1530,12 @@ PYBIND11_MODULE(_madspace_py, m) { {"helicity_in", MatrixElement::helicity_in}, {"diagram_in", MatrixElement::diagram_in}, {"channel_in", MatrixElement::channel_in}, + {"particle_count_in", MatrixElement::particle_count_in}, + {"pdg_ids_in", MatrixElement::pdg_ids_in}, + {"momentum_fraction_in", MatrixElement::momentum_fraction_in}, + {"scale_in", MatrixElement::scale_in}, + {"pdg_id_in", MatrixElement::pdg_id_in}, + {"beam_index_in", MatrixElement::beam_index_in}, } ); add_enum( @@ -1352,6 +1547,11 @@ PYBIND11_MODULE(_madspace_py, m) { {"color_index_out", MatrixElement::color_index_out}, {"helicity_index_out", MatrixElement::helicity_index_out}, {"diagram_index_out", MatrixElement::diagram_index_out}, + {"cut_out", MatrixElement::cut_out}, + {"scale_out", MatrixElement::scale_out}, + {"pdf_out", MatrixElement::pdf_out}, + {"bias_out", MatrixElement::bias_out}, + {"matrix_element_ratio_out", MatrixElement::matrix_element_ratio_out}, } ); matrix_element @@ -1782,6 +1982,7 @@ PYBIND11_MODULE(_madspace_py, m) { ); ++i; } + py::gil_scoped_release release; return opt.step(tensors); }, py::arg("inputs"), @@ -1862,6 +2063,14 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("dynamic_pid") = false, py::arg("prefix") = "", pydoc::doc("PartonDensity::PartonDensity") + ) + .def( + py::init&, int, bool>(), + py::arg("custom_pdf"), + py::arg("pids"), + py::arg("beam_index"), + py::arg("dynamic_pid") = false, + pydoc::doc("PartonDensity::PartonDensity#2") ); py::classh(m, "AlphaSGrid", pydoc::doc("AlphaSGrid")) @@ -1951,7 +2160,8 @@ PYBIND11_MODULE(_madspace_py, m) { double, double, double, - double>(), + double, + const std::optional&>(), py::arg("particle_count"), py::arg("dynamical_scale_type"), py::arg("ren_scale_fixed"), @@ -1960,6 +2170,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("fact_scale1"), py::arg("fact_scale2"), py::arg("scale_factor") = 1., + py::arg("custom_scale") = std::nullopt, pydoc::doc("EnergyScale::EnergyScale#4") ); @@ -1993,6 +2204,7 @@ PYBIND11_MODULE(_madspace_py, m) { PdfGrid, DifferentialCrossSection::CachedPdf>&, bool, + bool, bool>(), py::arg("matrix_element"), py::arg("cm_energy"), @@ -2003,6 +2215,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( @@ -2051,7 +2264,7 @@ PYBIND11_MODULE(_madspace_py, m) { const Integrand::AdaptiveDiscrete&, const Integrand::AdaptiveDiscrete&, const nested_vector2&, - const std::optional&, + const Integrand::PdfSource&, const std::optional&, const std::optional&, const std::optional&, @@ -2072,14 +2285,15 @@ PYBIND11_MODULE(_madspace_py, m) { const std::vector&, const std::vector&, const std::vector&, - std::size_t>(), + std::size_t, + const std::optional&>(), py::arg("mapping"), py::arg("diff_xs"), py::arg("adaptive_map") = std::monostate{}, py::arg("discrete_sym") = std::monostate{}, py::arg("discrete_flavor") = std::monostate{}, py::arg("pid_options") = nested_vector2{}, - py::arg("pdf_grid") = std::nullopt, + py::arg("pdf_grid") = std::monostate{}, py::arg("running_coupling") = std::nullopt, py::arg("energy_scale") = std::nullopt, py::arg("prop_chan_weights") = std::nullopt, @@ -2101,6 +2315,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("flavor_subproc_indices") = std::vector{}, py::arg("flavor_per_subproc_remap") = std::vector{}, py::arg("compressed_channel_weight_count") = 50, + py::arg("custom_bias") = std::nullopt, pydoc::doc("Integrand::Integrand") ) .def( @@ -2388,6 +2603,7 @@ PYBIND11_MODULE(_madspace_py, m) { .def( "train_step", &MadnisTraining::train_step, + py::call_guard(), py::arg("batch_index"), pydoc::doc("MadnisTraining::train_step") ) @@ -2445,6 +2661,7 @@ PYBIND11_MODULE(_madspace_py, m) { .def( "train", &MultiMadnisTraining::train, + py::call_guard(), pydoc::doc("MultiMadnisTraining::train") ) .def( @@ -2645,6 +2862,11 @@ PYBIND11_MODULE(_madspace_py, m) { &GeneratorStatus::iterations, pydoc::doc("GeneratorStatus::iterations") ) + .def_readwrite( + "iters_without_improvement", + &GeneratorStatus::iters_without_improvement, + pydoc::doc("GeneratorStatus::iters_without_improvement") + ) .def_readwrite( "optimized", &GeneratorStatus::optimized, @@ -3088,6 +3310,17 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("config"), pydoc::doc("ChannelEventGenerator::load") ) + .def_static( + "load_json", + &ChannelEventGenerator::load_json, + py::arg("channel_json"), + py::arg("contexts"), + py::kw_only(), + py::arg("event_file") = "", + py::arg("weight_file") = "", + py::arg("config"), + pydoc::doc("ChannelEventGenerator::load_json") + ) .def( py::init< const std::vector&, @@ -3118,8 +3351,133 @@ PYBIND11_MODULE(_madspace_py, m) { &ChannelEventGenerator::save, py::arg("save"), pydoc::doc("ChannelEventGenerator::save") + ) + .def( + "to_json", + &ChannelEventGenerator::to_json, + py::arg("include_estimates") = false, + pydoc::doc("ChannelEventGenerator::to_json") + ) + .def( + "set_fixed_max_weight", + &ChannelEventGenerator::set_fixed_max_weight, + py::arg("max_weight"), + pydoc::doc("ChannelEventGenerator::set_fixed_max_weight") + ) + .def( + "max_weight", + &ChannelEventGenerator::max_weight, + pydoc::doc("ChannelEventGenerator::max_weight") + ) + .def( + "step_max_weight", + &ChannelEventGenerator::step_max_weight, + py::arg("step"), + pydoc::doc("ChannelEventGenerator::step_max_weight") + ) + .def( + "set_fixed_step_max_weight", + &ChannelEventGenerator::set_fixed_step_max_weight, + py::arg("step"), + py::arg("max_weight"), + pydoc::doc("ChannelEventGenerator::set_fixed_step_max_weight") ); + py::classh weighting_step( + m, "WeightingStep", pydoc::doc("WeightingStep") + ); + add_enum( + weighting_step, + "Mode", + { + {"unweight", WeightingStep::unweight}, + {"reweight", WeightingStep::reweight}, + } + ); + weighting_step + .def( + py::init([](const std::vector& matrix_elements, + WeightingStep::Mode mode, + double max_overweight_truncation, + std::size_t freeze_max_weight_after, + const std::vector& + override_outputs) { + return WeightingStep{ + .matrix_elements = matrix_elements, + .mode = mode, + .max_overweight_truncation = max_overweight_truncation, + .freeze_max_weight_after = freeze_max_weight_after, + .override_outputs = override_outputs, + }; + }), + py::arg("matrix_elements"), + py::arg("mode") = WeightingStep::unweight, + py::arg("max_overweight_truncation") = 0.01, + py::arg("freeze_max_weight_after") = 10000, + py::arg("override_outputs") = + std::vector{}, + pydoc::doc("WeightingStep") + ) + .def_readwrite( + "matrix_elements", + &WeightingStep::matrix_elements, + pydoc::doc("WeightingStep::matrix_elements") + ) + .def_readwrite("mode", &WeightingStep::mode, pydoc::doc("WeightingStep::mode")) + .def_readwrite( + "max_overweight_truncation", + &WeightingStep::max_overweight_truncation, + pydoc::doc("WeightingStep::max_overweight_truncation") + ) + .def_readwrite( + "freeze_max_weight_after", + &WeightingStep::freeze_max_weight_after, + pydoc::doc("WeightingStep::freeze_max_weight_after") + ) + .def_readwrite( + "override_outputs", + &WeightingStep::override_outputs, + pydoc::doc("WeightingStep::override_outputs") + ) + .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) { + return WeightingGroup{ + .channel_indices = channel_indices, .steps = steps + }; + }), + py::arg("channel_indices"), + py::arg("steps"), + pydoc::doc("WeightingGroup") + ) + .def_readwrite( + "channel_indices", + &WeightingGroup::channel_indices, + pydoc::doc("WeightingGroup::channel_indices") + ) + .def_readwrite( + "steps", &WeightingGroup::steps, pydoc::doc("WeightingGroup::steps") + ); + 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( py::init([](const std::string& set_name, @@ -3602,7 +3960,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"), @@ -3612,22 +3971,38 @@ PYBIND11_MODULE(_madspace_py, m) { EventGenerator::default_config, "EventGenerator.default_config" ), + py::arg("groups") = std::vector{}, pydoc::doc("EventGenerator::EventGenerator") ) .def( "survey", &EventGenerator::survey, + py::call_guard(), py::arg("survey_pass") = 0, pydoc::doc("EventGenerator::survey") ) .def( "generate", &EventGenerator::generate, + py::call_guard(), pydoc::doc("EventGenerator::generate") ) + .def( + "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( "combine_to_compact_npy", &EventGenerator::combine_to_compact_npy, + py::call_guard(), py::arg("file_name"), py::arg("systematics") = nullptr, py::arg("histograms") = nullptr, @@ -3636,6 +4011,7 @@ PYBIND11_MODULE(_madspace_py, m) { .def( "combine_to_lhe_npy", &EventGenerator::combine_to_lhe_npy, + py::call_guard(), py::arg("file_name"), py::arg("lhe_completer"), py::arg("systematics") = nullptr, @@ -3645,6 +4021,7 @@ PYBIND11_MODULE(_madspace_py, m) { .def( "combine_to_lhe", &EventGenerator::combine_to_lhe, + py::call_guard(), py::arg("file_name"), py::arg("lhe_completer"), py::arg_v("meta", LHEMeta{}, "LHEMeta()"), @@ -3674,6 +4051,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, @@ -3745,10 +4206,19 @@ PYBIND11_MODULE(_madspace_py, m) { ); auto abort_check_function = [] { + // the long-running calls release the GIL + py::gil_scoped_acquire acquire; if (PyErr_CheckSignals() != 0) { throw py::error_already_set(); } }; 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/src/python/python_function.cpp b/madspace/src/python/python_function.cpp new file mode 100644 index 0000000000..e2bde8c847 --- /dev/null +++ b/madspace/src/python/python_function.cpp @@ -0,0 +1,223 @@ +#include "python_function.hpp" + +#include +#include +#include + +#include "function_runtime.hpp" + +using namespace madspace; + +namespace madspace_py { + +namespace { + +// Shared by the UMAMI entry points of one Python function. Python calls the +// `_call` method of `function` (a madspace.PythonFunction). +struct PythonFunctionState { + py::object function; + std::string name; + std::string device_name; + UmamiDevice umami_device; + std::array supported_inputs{}; + std::array required_inputs{}; + std::array supported_outputs{}; + + ~PythonFunctionState() { + if (Py_IsInitialized()) { + PythonGuard guard; + function = py::object(); + } else { + function.release(); + } + } +}; + +} // namespace + +Tensor tensor_view( + std::uintptr_t address, + DataType dtype, + const std::vector& shape, + const std::vector& stride, + std::size_t offset, + const std::string& device_name +) { + DevicePtr device; + if (device_name == "cpu") { + device = cpu_device(); + } else if (device_name == "cuda") { + device = cuda_device(); + } else if (device_name == "hip") { + device = hip_device(); + } else { + throw std::invalid_argument(std::format("unknown device {}", device_name)); + } + std::size_t item_size = + dtype == DataType::dt_float ? sizeof(double) : sizeof(me_int_t); + void* data = reinterpret_cast(address + offset * item_size); + return Tensor(dtype, Sizes(shape), Sizes(stride), device, data, [] {}); +} + +const MatrixElementApi& add_python_function( + Context& context, + py::object function, + const std::vector& input_keys, + const std::vector& required_input_keys, + const std::vector& output_keys, + const std::string& framework, + const std::string& name +) { + auto state = std::make_shared(); + state->function = function; + state->name = name; + switch (context.device()->device_type()) { + case DeviceType::cpu: + state->umami_device = UMAMI_DEVICE_CPU; + state->device_name = "cpu"; + break; + case DeviceType::cuda: + state->umami_device = UMAMI_DEVICE_CUDA; + state->device_name = "cuda"; + break; + case DeviceType::hip: + state->umami_device = UMAMI_DEVICE_HIP; + state->device_name = "hip"; + break; + } + if (framework == "numpy" && state->umami_device != UMAMI_DEVICE_CPU) { + throw std::invalid_argument( + std::format( + "Python function {}: numpy functions run on the CPU only, use torch " + "for GPU runs", + name + ) + ); + } + for (int key : input_keys) { + if (key < 0 || key >= UMAMI_INPUT_KEY_COUNT) { + throw std::invalid_argument("invalid UMAMI input key"); + } + state->supported_inputs[key] = true; + } + for (int key : required_input_keys) { + if (key < 0 || key >= UMAMI_INPUT_KEY_COUNT || !state->supported_inputs[key]) { + throw std::invalid_argument("required input key is not supported"); + } + state->required_inputs[key] = true; + } + for (int key : output_keys) { + if (key < 0 || key >= UMAMI_OUTPUT_KEY_COUNT || key == UMAMI_OUT_GPU_STREAM) { + throw std::invalid_argument("invalid UMAMI output key"); + } + state->supported_outputs[key] = true; + } + + UmamiFunctions functions; + functions.get_meta = [state](UmamiMetaKey meta_key, void* result) { + if (meta_key == UMAMI_META_DEVICE) { + *static_cast(result) = state->umami_device; + return UMAMI_SUCCESS; + } + return UMAMI_ERROR_UNSUPPORTED_META; + }; + functions.supported_inputs = [state](bool const** supported, int* count) { + *supported = state->supported_inputs.data(); + *count = UMAMI_INPUT_KEY_COUNT; + return UMAMI_SUCCESS; + }; + functions.required_inputs = [state](bool const** required, int* count) { + *required = state->required_inputs.data(); + *count = UMAMI_INPUT_KEY_COUNT; + return UMAMI_SUCCESS; + }; + functions.supported_outputs = [state](bool const** supported, int* count) { + *supported = state->supported_outputs.data(); + *count = UMAMI_OUTPUT_KEY_COUNT; + return UMAMI_SUCCESS; + }; + functions.initialize = [state](UmamiHandle* handle, char const*) { + *handle = state.get(); + return UMAMI_SUCCESS; + }; + functions.free = [](UmamiHandle) { return UMAMI_SUCCESS; }; + functions.matrix_element = + [state]( + UmamiHandle, + std::size_t count, + std::size_t stride, + std::size_t offset, + std::size_t input_count, + UmamiInputKey const* input_keys, + void const* const* inputs, + std::size_t output_count, + UmamiOutputKey const* output_keys, + void* const* outputs + ) { + std::array has_input{}; + for (std::size_t i = 0; i < input_count; ++i) { + if (input_keys[i] < 0 || input_keys[i] >= UMAMI_INPUT_KEY_COUNT || + !state->supported_inputs[input_keys[i]]) { + return UMAMI_ERROR_UNSUPPORTED_INPUT; + } + has_input[input_keys[i]] = true; + } + for (int key = 0; key < UMAMI_INPUT_KEY_COUNT; ++key) { + if (state->required_inputs[key] && !has_input[key]) { + return UMAMI_ERROR_MISSING_INPUT; + } + } + std::optional stream; + for (std::size_t i = 0; i < output_count; ++i) { + if (output_keys[i] == UMAMI_OUT_GPU_STREAM) { + // the GPU runtime passes the stream handle itself + stream = reinterpret_cast(outputs[i]); + } else if (output_keys[i] < 0 || + output_keys[i] >= UMAMI_OUTPUT_KEY_COUNT || + !state->supported_outputs[output_keys[i]]) { + return UMAMI_ERROR_UNSUPPORTED_OUTPUT; + } + } + if (count == 0) { + return UMAMI_SUCCESS; + } + + std::string error; + { + PythonGuard guard; + try { + py::dict input_dict, output_dict; + for (std::size_t i = 0; i < input_count; ++i) { + input_dict[py::int_(static_cast(input_keys[i]))] = + py::int_(reinterpret_cast(inputs[i])); + } + for (std::size_t i = 0; i < output_count; ++i) { + if (output_keys[i] != UMAMI_OUT_GPU_STREAM) { + output_dict[py::int_(static_cast(output_keys[i]))] = + py::int_(reinterpret_cast(outputs[i])); + } + } + state->function.attr("_call")( + state->device_name, + count, + stride, + offset, + input_dict, + output_dict, + stream ? py::object(py::int_(*stream)) : py::object(py::none()) + ); + } catch (py::error_already_set& e) { + error = e.what(); + } + } + if (!error.empty()) { + throw std::runtime_error( + std::format("Error in Python function {}: {}", state->name, error) + ); + } + return UMAMI_SUCCESS; + }; + return context.add_matrix_element(std::move(functions), name); +} + +} // namespace madspace_py diff --git a/madspace/src/python/python_function.hpp b/madspace/src/python/python_function.hpp new file mode 100644 index 0000000000..fe76343c60 --- /dev/null +++ b/madspace/src/python/python_function.hpp @@ -0,0 +1,37 @@ +#pragma once + +#include +#include +#include +#include + +#include "madspace/driver/context.hpp" + +namespace py = pybind11; + +namespace madspace_py { + +/// Non-owning view of `shape` and element `stride` onto the memory at +/// `address`, starting `offset` elements in. +madspace::Tensor tensor_view( + std::uintptr_t address, + madspace::DataType dtype, + const std::vector& shape, + const std::vector& stride, + std::size_t offset, + const std::string& device_name +); + +/// Add a UMAMI matrix element to `context` that calls `function._call` with +/// the raw UMAMI buffers, see madspace.PythonFunction. +const madspace::MatrixElementApi& add_python_function( + madspace::Context& context, + py::object function, + const std::vector& input_keys, + const std::vector& required_input_keys, + const std::vector& output_keys, + const std::string& framework, + const std::string& name +); + +} // namespace madspace_py diff --git a/madspace/tests/test_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_integrand_me_outputs.py b/madspace/tests/test_integrand_me_outputs.py new file mode 100644 index 0000000000..81b745bdf4 --- /dev/null +++ b/madspace/tests/test_integrand_me_outputs.py @@ -0,0 +1,178 @@ +"""Integrand with matrix elements that only return |M|^2: the channel weights, +helicity, color and diagram index are optional outputs of the matrix element. + +Uses the e+ e- -> d d~ g setup of test_event_stream with PythonFunction matrix +elements, as a stand-in for surrogates that only predict |M|^2. +""" + +import os + +import numpy as np +import pytest +from test_event_stream import ( + CM_ENERGY, + PERMUTATIONS, + PIDS, + TOPOLOGY, + make_config, + make_context, +) + +import madspace as ms + + +def python_matrix_element(context, function, outputs=("matrix_element",)): + api = context.add_matrix_element( + ms.PythonFunction( + function, ["momenta"], list(outputs), particle_count=len(PIDS) + ) + ) + output_enums = { + "matrix_element": ms.MatrixElement.matrix_element_out, + "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], + [output_enums[output] for output in outputs], + ) + + +def flat_matrix_element(): + return ms.MatrixElement( + 0xBADCAFE, + len(PIDS), + ms.Integrand.matrix_element_inputs, + ms.Integrand.matrix_element_outputs, + 1, + True, + ) + + +def make_integrand(alphas_file, matrix_element, return_matrix_element=False, **kwargs): + mapping = ms.PhaseSpaceMapping( + TOPOLOGY, + CM_ENERGY, + permutations=PERMUTATIONS, + leptonic=True, + cuts=ms.Cuts( + [ + ms.CutItem(ms.Observable(PIDS, ms.Observable.obs_pt, [[i]]), min=20.0) + for i in (2, 3, 4) + ] + ), + ) + diff_xs = ms.DifferentialCrossSection( + matrix_element=matrix_element, + cm_energy=CM_ENERGY, + running_coupling=None, + energy_scale=ms.CachedScale(), + return_matrix_element=return_matrix_element, + ) + return ms.Integrand( + mapping, + [diff_xs], + running_coupling=ms.RunningCoupling(ms.AlphaSGrid(alphas_file)), + energy_scale=ms.EnergyScale(len(PIDS), 91.188), + channel_indices=[0], + **kwargs, + ) + + +def generate(run_dir, context, integrand, name): + config = make_config() + config.target_count = 2000 + channel = ms.ChannelEventGenerator( + contexts=[context], + integrand=integrand, + event_file=os.path.join(run_dir, f"events.{name}.npy"), + weight_file=os.path.join(run_dir, f"weights.{name}.npy"), + config=config, + subprocess_index=0, + name=name, + histograms=None, + ) + generator = ms.EventGenerator([context], [channel], seed=7, config=config) + generator.survey() + generator.generate() + output_file = os.path.join(run_dir, f"combined.{name}.npy") + generator.combine_to_compact_npy(output_file) + return generator.status(), np.load(output_file) + + +def test_only_matrix_element_output(tmp_path): + """Without the index outputs the events carry -1, and |M|^2 = 1 reproduces + the integral of the built-in flat matrix element.""" + context, alphas_file = make_context(str(tmp_path), 2) + matrix_element = python_matrix_element( + context, lambda momenta: np.ones(len(momenta)) + ) + status, events = generate( + str(tmp_path), context, make_integrand(alphas_file, matrix_element), "py" + ) + flat_status, flat_events = generate( + str(tmp_path), + context, + make_integrand(alphas_file, flat_matrix_element()), + "flat", + ) + assert len(events) == len(flat_events) == 2000 + for field in ("helicity_index", "color_index", "diagram_index"): + assert (events[field] == -1).all() + assert (flat_events[field] >= 0).all() + error = np.hypot(status.error, flat_status.error) + assert abs(status.mean - flat_status.mean) < 5 * error + + +def test_provided_index_output(tmp_path): + """An index the matrix element returns is written to the events.""" + context, alphas_file = make_context(str(tmp_path), 1) + matrix_element = python_matrix_element( + context, + lambda momenta: (np.ones(len(momenta)), np.full(len(momenta), 1, np.int32)), + ("matrix_element", "helicity_index"), + ) + _, events = generate( + str(tmp_path), context, make_integrand(alphas_file, matrix_element), "hel" + ) + assert (events["helicity_index"] == 1).all() + assert (events["color_index"] == -1).all() + + +def test_raw_matrix_element_output(tmp_path): + """Returning the raw |M|^2 adds an integrand output without changing the + generated events.""" + context, alphas_file = make_context(str(tmp_path), 1) + matrix_element = python_matrix_element( + context, lambda momenta: 1.0 + momenta[:, 2, 0] / CM_ENERGY + ) + results = [ + generate( + str(tmp_path), + context, + make_integrand(alphas_file, matrix_element, raw), + f"raw{raw}", + ) + for raw in (False, True) + ] + (status_a, events_a), (status_b, events_b) = results + assert status_a.mean == status_b.mean + assert np.array_equal(events_a, events_b) + + +def test_diagram_channel_weights_need_amp2(tmp_path): + context, alphas_file = make_context(str(tmp_path), 1) + matrix_element = python_matrix_element( + context, lambda momenta: np.ones(len(momenta)) + ) + with pytest.raises(ValueError, match="diagram_amp2"): + make_integrand( + alphas_file, + matrix_element, + first_chan_weight_remap=[[0, 1]], + first_remapped_chan_count=2, + ) diff --git a/madspace/tests/test_python_function.py b/madspace/tests/test_python_function.py new file mode 100644 index 0000000000..4f290637d5 --- /dev/null +++ b/madspace/tests/test_python_function.py @@ -0,0 +1,160 @@ +import numpy as np +import pytest + +import madspace as ms + +PDG_IDS = [21, 21, 6, -6] + + +@pytest.fixture +def context(): + # several threads, so the function is called from worker threads + return ms.Context(device=ms.cpu_device(), thread_count=4) + + +@pytest.fixture +def momenta(): + rng = np.random.default_rng(1234) + return rng.normal(size=(5000, 4, 4)) * 100.0 + + +def pt(momenta, index): + return np.hypot(momenta[:, index, 1], momenta[:, index, 2]) + + +def run_custom(context, function, kind, momenta, framework="numpy"): + api = context.add_matrix_element( + ms.PythonFunction(function, ["momenta", "pdg_ids"], [kind], framework=framework) + ) + generator = ms.CustomFunction(api, getattr(ms.CustomFunction, kind), 4, [PDG_IDS]) + return ms.FunctionRuntime(generator.function(), context)(momenta) + + +def test_numpy_cut(context, momenta): + def cut(p, pdg_ids): + assert p.shape == (len(p), 4, 4) + assert (pdg_ids == PDG_IDS).all() + return (pt(p, 2) > 100.0).astype(np.float64) + + values = run_custom(context, cut, "cut", momenta) + assert np.array_equal(np.asarray(values), (pt(momenta, 2) > 100.0).astype(float)) + + +def test_numpy_scale(context, momenta): + def scale(p, pdg_ids): + return (pt(p, 2) + pt(p, 3)) / 2 + + values = run_custom(context, scale, "scale", momenta) + assert np.allclose(values, (pt(momenta, 2) + pt(momenta, 3)) / 2) + + +def test_numpy_pdf(context): + def pdf(beam_index, pdg_id, x, q): + return np.where(pdg_id == 21, 3.0, 0.5) * np.sqrt(x) * (1 - x) ** 3 + beam_index + + api = context.add_matrix_element( + ms.PythonFunction( + pdf, ["beam_index", "pdg_id", "momentum_fraction", "scale"], ["pdf"] + ) + ) + x = np.array([0.1, 0.5, 0.9]) + for beam_index in (0, 1): + density = ms.PartonDensity(ms.CustomPdf(api), [21, 2], beam_index) + values = np.asarray( + ms.FunctionRuntime(density.function(), context)(x, np.full(3, 100.0)) + ) + toy = np.sqrt(x) * (1 - x) ** 3 + assert np.allclose(values[:, 0], 3 * toy + beam_index) + assert np.allclose(values[:, 1], 0.5 * toy + beam_index) + + +def test_multiple_outputs_fixed_particle_count(context, momenta): + def matrix_element(p): + return p[:, 2, 0] ** 2, np.full(len(p), 3, dtype=np.int32) + + api = context.add_matrix_element( + ms.PythonFunction( + matrix_element, + ["momenta"], + ["matrix_element", "helicity_index"], + particle_count=4, + ) + ) + generator = ms.MatrixElement( + api.index(), + 4, + [ms.MatrixElement.momenta_in], + [ms.MatrixElement.matrix_element_out, ms.MatrixElement.helicity_index_out], + ) + me, helicity = ms.FunctionRuntime(generator.function(), context)(momenta) + assert np.allclose(me, momenta[:, 2, 0] ** 2) + assert (np.asarray(helicity) == 3).all() + + +@pytest.mark.parametrize( + "result, message", + [ + (lambda p: np.ones(len(p), dtype=np.float32), "must have dtype float64"), + (lambda p: np.ones(3), "must have shape"), + (lambda p: list(np.ones(len(p))), "must be a numpy.ndarray"), + ], +) +def test_result_checks(context, momenta, result, message): + with pytest.raises(RuntimeError, match=message): + run_custom(context, lambda p, pdg_ids: result(p), "cut", momenta) + + +def test_exception_propagates(context, momenta): + def broken(p, pdg_ids): + raise KeyError("something went wrong") + + with pytest.raises( + RuntimeError, match="Error in Python function broken.*something" + ): + run_custom(context, broken, "cut", momenta) + + +def test_invalid_arguments(): + with pytest.raises(ValueError, match="unknown input"): + ms.PythonFunction(lambda: None, ["momentum"], ["cut"]) + with pytest.raises(ValueError, match="unknown output"): + ms.PythonFunction(lambda: None, ["momenta"], ["cuts"]) + with pytest.raises(ValueError, match="framework"): + ms.PythonFunction(lambda: None, ["momenta"], ["cut"], framework="jax") + + +def test_release_inputs_on_worker_threads(context, momenta): + # numpy-backed inputs are freed by madspace, possibly on a worker thread + def cut(p, pdg_ids): + return np.ones(len(p)) + + api = context.add_matrix_element( + ms.PythonFunction(cut, ["momenta", "pdg_ids"], ["cut"]) + ) + generator = ms.CustomFunction(api, ms.CustomFunction.cut, 4, [PDG_IDS]) + runtime = ms.FunctionRuntime(generator.function(), context) + for _ in range(20): + runtime(momenta.copy()) + runtime.release_inputs() + + +def test_torch_cpu(context, momenta): + torch = pytest.importorskip("torch") + + def scale(p, pdg_ids): + assert isinstance(p, torch.Tensor) + return torch.hypot(p[:, 2:, 1], p[:, 2:, 2]).sum(dim=1) / 2 + + expected = (pt(momenta, 2) + pt(momenta, 3)) / 2 + values = run_custom(context, scale, "scale", momenta, framework="torch") + assert np.allclose(values, expected) + values = run_custom( + context, scale, "scale", torch.from_numpy(momenta), framework="torch" + ) + assert torch.allclose(values, torch.from_numpy(expected)) + + +def test_torch_wrong_type(context, momenta): + pytest.importorskip("torch") + with pytest.raises(RuntimeError, match="must be a torch.Tensor"): + run_custom(context, lambda p, ids: np.ones(len(p)), "cut", momenta, "torch") diff --git a/madspace/tests/test_tar.py b/madspace/tests/test_tar.py new file mode 100644 index 0000000000..55cd9091bc --- /dev/null +++ b/madspace/tests/test_tar.py @@ -0,0 +1,105 @@ +import subprocess +import tarfile + +import numpy as np +import pytest + +import madspace as ms + +torch = pytest.importorskip("torch") + + +def make_context(tmp_path): + ctx = ms.Context(device=ms.cpu_device(), thread_count=1) + rng = np.random.default_rng(1) + values = { + "a.float": rng.normal(size=(1, 4, 50)), + "b.float": rng.normal(size=(1, 3, 5)), + "c.int": rng.integers(-100, 100, size=(1, 2, 7)).astype(np.int32), + "d.scalar": rng.normal(size=(1, 1)), + } + for name, value in values.items(): + dtype = ms.DataType.int if value.dtype == np.int32 else ms.DataType.float + glob = ctx.define_global(name, dtype, list(value.shape[1:])) + glob.torch().copy_(torch.from_numpy(value)) + return ctx, values + + +def test_roundtrip(tmp_path): + ctx, values = make_context(tmp_path) + path = str(tmp_path / "globals.tar") + ctx.save_globals(path) + + ctx2 = ms.Context(device=ms.cpu_device(), thread_count=1) + ctx2.load_globals(path) + assert sorted(ctx2.global_names()) == sorted(values) + for name, value in values.items(): + np.testing.assert_array_equal( + ctx2.get_global(name).numpy(), ctx.get_global(name).numpy() + ) + + +def test_standard_tar(tmp_path): + ctx, values = make_context(tmp_path) + path = tmp_path / "globals.tar" + ctx.save_globals(str(path)) + + # system tar and the Python standard library can read it + listing = subprocess.run( + ["tar", "tf", str(path)], capture_output=True, text=True, check=True + ).stdout.split() + assert sorted(listing) == sorted(name + ".npy" for name in values) + with tarfile.open(path) as tar: + assert tar.getnames() == sorted(name + ".npy" for name in values) + for name in values: + tar.extract(name + ".npy", tmp_path / "out", filter="data") + + # and so can numpy, with the same values as the context holds + for name in values: + expected = ctx.get_global(name).numpy() + loaded = np.load(tmp_path / "out" / (name + ".npy")) + np.testing.assert_array_equal(loaded, expected) + + +def test_extract_with_system_tar(tmp_path): + ctx, values = make_context(tmp_path) + path = tmp_path / "globals.tar" + ctx.save_globals(str(path)) + out = tmp_path / "out" + out.mkdir() + subprocess.run(["tar", "xf", str(path), "-C", str(out)], check=True) + assert sorted(p.name for p in out.iterdir()) == sorted( + name + ".npy" for name in values + ) + + +def test_read_tar_from_other_tool(tmp_path): + ctx, values = make_context(tmp_path) + first = tmp_path / "first.tar" + ctx.save_globals(str(first)) + src = tmp_path / "src" + src.mkdir() + with tarfile.open(first) as tar: + tar.extractall(src, filter="data") + + # an archive written by Python's tarfile (GNU format) and by system tar + for fmt, name in [ + (tarfile.GNU_FORMAT, "gnu.tar"), + (tarfile.USTAR_FORMAT, "ustar.tar"), + ]: + other = tmp_path / name + with tarfile.open(other, "w", format=fmt) as tar: + for p in sorted(src.iterdir()): + tar.add(p, arcname=p.name) + ctx2 = ms.Context(device=ms.cpu_device(), thread_count=1) + ctx2.load_globals(str(other)) + for gname in values: + np.testing.assert_array_equal( + ctx2.get_global(gname).numpy(), ctx.get_global(gname).numpy() + ) + + +def test_missing_file(tmp_path): + ctx = ms.Context(device=ms.cpu_device(), thread_count=1) + with pytest.raises(RuntimeError): + ctx.load_globals(str(tmp_path / "nonexistent.tar")) diff --git a/madspace/tests/test_weighting_steps.py b/madspace/tests/test_weighting_steps.py new file mode 100644 index 0000000000..fef4cbcf8b --- /dev/null +++ b/madspace/tests/test_weighting_steps.py @@ -0,0 +1,501 @@ +"""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",), +): + """Survey and generate, returns the status, the combined events and the + channels. step_me and step_outputs can be lists, one entry per step.""" + os.makedirs(run_dir, exist_ok=True) + context, alphas_file = make_context(run_dir, threads) + me0 = python_matrix_element(context, integrand_me, integrand_outputs) + if not isinstance(step_me, list): + step_me = [step_me] * len(step_modes) + step_outputs = [step_outputs] * len(step_modes) + step_mes = [ + python_matrix_element(context, me, outputs) + for me, outputs in zip(step_me, step_outputs) + ] + config = make_config() + config.target_count = target_count + config.cpu_batch_size = 500 + channels = [ + ms.ChannelEventGenerator( + contexts=[context], + integrand=make_integrand(alphas_file, me0, pt_min), + event_file=os.path.join(run_dir, f"events.{i}.npy"), + weight_file=os.path.join(run_dir, f"weights.{i}.npy"), + config=config, + subprocess_index=0, + name=str(i), + histograms=None, + ) + for i, pt_min in enumerate(pt_mins) + ] + steps = [ + ms.WeightingStep([me], mode=mode, override_outputs=list(override_outputs)) + for me, mode in zip(step_mes, step_modes) + ] + groups = [ms.WeightingGroup(list(grouped), steps)] if steps else [] + generator = ms.EventGenerator( + [context], channels, seed=seed, config=config, groups=groups + ) + generator.survey() + generator.generate() + output_file = os.path.join(run_dir, "combined.npy") + generator.combine_to_compact_npy(output_file) + return generator.status(), np.load(output_file), channels + + +@pytest.fixture(scope="module") +def reference(tmp_path_factory): + """The true matrix element in the integrand.""" + return generate( + str(tmp_path_factory.mktemp("reference")), true_me, target_count=20000 + ) + + +def check_compatible(status, events, reference): + ref_status, ref_events, _ = reference + error = np.hypot(status.error, ref_status.error) + assert abs(status.mean - ref_status.mean) < 4 * error + # the events are normalized to the cross section + assert np.mean(events["weight"]) == pytest.approx(status.mean_abs, rel=1e-9) + # mean energy of the first quark, weighted + energies = [ + np.sum(ev["weight"] * ev["part3_energy"]) / np.sum(ev["weight"]) + for ev in (events, ref_events) + ] + assert energies[0] == pytest.approx(energies[1], rel=0.01) + + +def test_unweight_step(tmp_path, reference): + status, events, channels = generate( + str(tmp_path), surrogate_me, ["unweight"], target_count=20000 + ) + assert len(events) == 20000 + check_compatible(status, events, reference) + # unweighted: apart from the overweights, all events have the same weight + weights = events["weight"] + assert np.mean(np.isclose(weights, np.min(weights))) > 0.98 + assert channels[0].step_max_weight(1) > 0 + + +def test_reweight_step(tmp_path, reference): + status, events, _ = generate( + str(tmp_path), surrogate_me, ["reweight"], target_count=20000 + ) + check_compatible(status, events, reference) + # the ratios of the matrix elements stay in the weights + weights = events["weight"] + assert np.std(weights) / np.mean(weights) > 0.1 + + +def test_reweight_then_unweight(tmp_path, reference): + """The unweighting step absorbs the ratio of the reweighting step: the + second step has the same matrix element, so it only unweights.""" + status, events, _ = generate( + str(tmp_path), surrogate_me, ["reweight", "unweight"], target_count=20000 + ) + check_compatible(status, events, reference) + weights = events["weight"] + assert np.mean(np.isclose(weights, np.min(weights))) > 0.98 + + +@pytest.mark.parametrize("step_modes", [["unweight"], ["reweight", "unweight"]]) +def test_reproducible(tmp_path, step_modes): + """Two grouped channels and one without steps, which share the batches + differently with different thread counts.""" + results = [ + generate( + str(tmp_path / str(threads)), + surrogate_me, + step_modes, + threads=threads, + pt_mins=(20.0, 40.0, 60.0), + grouped=(0, 2), + ) + for threads in (1, 4) + ] + (status_a, events_a, _), (status_b, events_b, _) = results + assert len(events_a) == 5000 + assert np.array_equal(events_a, events_b) + assert status_a.mean == status_b.mean + + +def test_index_outputs(tmp_path): + """A step fills in indices the integrand matrix element does not return, + and replaces the others only with override_outputs.""" + + def with_helicity(momenta): + return true_me(momenta), np.full(len(momenta), 1, np.int32) + + outputs = ("matrix_element", "helicity_index") + _, events, _ = generate( + str(tmp_path / "fill"), + surrogate_me, + ["unweight"], + step_me=with_helicity, + step_outputs=outputs, + ) + assert (events["helicity_index"] == 1).all() + assert (events["color_index"] == -1).all() + + def surrogate_with_helicity(momenta): + return surrogate_me(momenta), np.zeros(len(momenta), np.int32) + + for override, helicity in [((), 0), ((ms.MatrixElement.helicity_index_out,), 1)]: + _, events, _ = generate( + str(tmp_path / f"override{helicity}"), + surrogate_with_helicity, + ["unweight"], + step_me=with_helicity, + step_outputs=outputs, + integrand_outputs=outputs, + override_outputs=override, + ) + assert (events["helicity_index"] == helicity).all() + + +def true_over_surrogate(momenta): + return true_me(momenta) / surrogate_me(momenta) + + +@pytest.mark.parametrize("step_modes", [["unweight"], ["reweight"]]) +def test_ratio_step(tmp_path, step_modes): + """A step that returns |M_true|^2 / |M_surrogate|^2 gives the same events + as a step that returns |M_true|^2: the pipeline computes the same ratio.""" + results = [ + generate( + str(tmp_path / name), + surrogate_me, + step_modes, + step_me=step_me, + step_outputs=outputs, + ) + for name, step_me, outputs in [ + ("me", true_me, ("matrix_element",)), + ("ratio", true_over_surrogate, ("matrix_element_ratio",)), + ] + ] + (status_me, events_me, _), (status_ratio, events_ratio, _) = results + assert np.array_equal(events_me, events_ratio) + assert status_me.mean == status_ratio.mean + + +def test_ratio_then_matrix_element(tmp_path, reference): + """After a ratio step, the matrix element it implies is the reference of the + next step, so the ratios still multiply to |M_true|^2 / |M_surrogate|^2.""" + + def intermediate_me(momenta): + return true_me(momenta) * (1.0 + 0.2 * np.cos(momenta[:, 2, 2] / 40.0)) + + status, events, _ = generate( + str(tmp_path), + surrogate_me, + ["reweight", "unweight"], + step_me=[lambda p: intermediate_me(p) / surrogate_me(p), true_me], + step_outputs=[("matrix_element_ratio",), ("matrix_element",)], + target_count=20000, + ) + check_compatible(status, events, reference) + weights = events["weight"] + assert np.mean(np.isclose(weights, np.min(weights))) > 0.98 + + +def test_ratio_matrix_element_rules(tmp_path): + context, _ = make_context(str(tmp_path), 1) + ratio_me = python_matrix_element( + context, true_over_surrogate, ("matrix_element_ratio",) + ) + assert ms.WeightingStep([ratio_me]).returns_ratio() + assert not ms.WeightingStep( + [python_matrix_element(context, true_me)] + ).returns_ratio() + with pytest.raises(ValueError, match="weighting step"): + ms.DifferentialCrossSection( + matrix_element=ratio_me, + cm_energy=CM_ENERGY, + running_coupling=None, + energy_scale=ms.CachedScale(), + ) + with pytest.raises(ValueError, match="not both"): + ms.MatrixElement( + ratio_me.matrix_element_index(), + len(PIDS), + [ms.MatrixElement.momenta_in], + [ + ms.MatrixElement.matrix_element_out, + ms.MatrixElement.matrix_element_ratio_out, + ], + ) + groups = [ms.WeightingGroup([0], [ms.WeightingStep([ratio_me])])] + text = ms.weighting_groups_to_json(groups) + loaded = ms.weighting_groups_from_json(text) + assert loaded[0].steps[0].returns_ratio() + assert ms.weighting_groups_to_json(loaded) == text + + +def test_invalid_groups(tmp_path): + context, alphas_file = make_context(str(tmp_path), 1) + me = python_matrix_element(context, true_me) + 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 + ) 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_custom_functions.py b/tests/acceptance_tests/test_mg7_custom_functions.py new file mode 100644 index 0000000000..c15b31e1b1 --- /dev/null +++ b/tests/acceptance_tests/test_mg7_custom_functions.py @@ -0,0 +1,489 @@ +################################################################################ +# +# Copyright (c) 2026 The MadGraph7 Development team and Contributors +# +# This file is a part of the MadGraph7 project, an application which +# automatically generates Feynman diagrams and matrix elements for arbitrary +# high-energy processes in the Standard Model and beyond. +# +# It is subject to the MadGraph7 license which should accompany this +# distribution. +# +# For more information, visit madgraph.phys.ucl.ac.be and amcatnlo.web.cern.ch +# +################################################################################ +"""MG7 custom cut, scale, PDF and bias functions ([custom_functions]). + +Each custom function is created with bin/init_custom_func in a copy of one +``p p > t t~`` output, given a body that reproduces a built-in setting, and +compared with a run using that built-in setting: + + * cut: pT > 150 GeV for both tops against ``top-pt.min = 150``, and every + event passes the cut; + * scale: H_T/2 against ``dynamical_scale_choice = half_transverse_mass``, + and the LHE scale of every event is H_T/2 of its momenta; + * pdf: the LHAPDF set of the run card called through the LHAPDF C++ API + against the built-in interpolation of the same set; + * bias: b = max((pT / 100 GeV)^2, 1) keeps the cross section and the + weighted pT tail, while putting more events into the tail; + * every function in every language with a toolchain here builds, loads and + returns its default value when evaluated through madspace. + +Run locally with:: + + ./tests/test_manager.py test_.*custom.*_mg7 -pA -t0 + +``MG7_CUSTOM_EVENTS`` sets the events per run (default 5000). +""" + +from __future__ import absolute_import +from __future__ import division + +import glob +import gzip +import json +import math +import os +import re +import shutil +import subprocess +import sys +import tempfile +import unittest + +import madgraph.interface.master_interface as MGCmd +import madgraph.various.misc as misc +from tests.acceptance_tests.test_check_xsec_processes_mg7 import ( + _REFERENCE_PDF, _edit_run_card, _require_mg7_runtime, _tail) + +pjoin = os.path.join + +_PROCESS = 'p p > t t~' +_EVENTS = int(os.environ.get('MG7_CUSTOM_EVENTS', 5000)) +_NSIGMA = 3 +_SEED = 1234 +_TOP_PT_CUT = 150. + +# bodies of the custom functions, in C++ +_CUT_BODY = r''' +#include +extern "C" double custom_cut(int particle_count, const int* pdg_ids, const double momenta[][4]) { + for (int i = 2; i < 4; ++i) { + if (std::hypot(momenta[i][1], momenta[i][2]) < %(cut)r) { + return 0.; + } + } + return 1.; +} +''' % {'cut': _TOP_PT_CUT} + +_SCALE_BODY = r''' +#include +extern "C" double custom_scale(int particle_count, const int* pdg_ids, const double momenta[][4]) { + double ht = 0.; + for (int i = 2; i < particle_count; ++i) { + double e = momenta[i][0], pz = momenta[i][3]; + ht += std::sqrt(std::fmax(e * e - pz * pz, 0.)); + } + return ht / 2.; +} +''' + +_PDF_BODY = r''' +#include +#include "LHAPDF/LHAPDF.h" +extern "C" double custom_pdf(int beam_index, int pdg_id, double x, double q) { + static LHAPDF::PDF* pdf = LHAPDF::mkPDF("%(pdf)s", 0); + static std::mutex mutex; + std::lock_guard lock(mutex); + return pdf->xfxQ(pdg_id, x, q); +} +''' % {'pdf': _REFERENCE_PDF} + +_BIAS_BODY = r''' +#include +extern "C" double custom_bias(int particle_count, const int* pdg_ids, const double momenta[][4]) { + double pt2 = momenta[2][1] * momenta[2][1] + momenta[2][2] * momenta[2][2]; + return std::fmax(pt2 / (100. * 100.), 1.); +} +''' + +# the same functions in Python with numpy, and the cut with torch +_PY_HEADER = 'import numpy as np\nFRAMEWORK = "numpy"\n' + +_CUT_NUMPY = _PY_HEADER + ''' +def custom_cut(momenta, pdg_ids): + pt = np.hypot(momenta[:, 2:4, 1], momenta[:, 2:4, 2]) + return np.where((pt >= %(cut)r).all(axis=1), 1., 0.) +''' % {'cut': _TOP_PT_CUT} + +_CUT_TORCH = '''import torch +FRAMEWORK = "torch" + +def custom_cut(momenta, pdg_ids): + pt = torch.hypot(momenta[:, 2:4, 1], momenta[:, 2:4, 2]) + return (pt >= %(cut)r).all(dim=1).double() +''' % {'cut': _TOP_PT_CUT} + +_SCALE_NUMPY = _PY_HEADER + ''' +def custom_scale(momenta, pdg_ids): + final = momenta[:, 2:] + return np.sqrt(np.maximum(final[..., 0] ** 2 - final[..., 3] ** 2, 0.)).sum(axis=1) / 2. +''' + +_PDF_NUMPY = _PY_HEADER + '''import lhapdf +pdf_set = lhapdf.mkPDF("%(pdf)s", 0) + +def custom_pdf(beam_index, pdg_id, x, q): + return np.array([pdf_set.xfxQ(int(i), a, b) for i, a, b in zip(pdg_id, x, q)]) +''' % {'pdf': _REFERENCE_PDF} + +_BIAS_NUMPY = _PY_HEADER + ''' +def custom_bias(momenta, pdg_ids): + pt2 = momenta[:, 2, 1] ** 2 + momenta[:, 2, 2] ** 2 + return np.maximum(pt2 / 100. ** 2, 1.) +''' + +_EXTENSIONS = {'cpp': 'cpp', 'python-numpy': 'py', 'python-torch': 'py'} + + +def _set(text, key, value): + """Set ``key = value`` (first occurrence at the start of a line).""" + text, count = re.subn(r'(?m)^%s = .*$' % re.escape(key), + '%s = %s' % (key, value), text, count=1) + assert count == 1, 'no %s in the run card' % key + return text + + +def _read_lhe(path): + """(weights, scales, momenta) of an LHE file; momenta[i] lists the + (E, px, py, pz) of every particle of event i.""" + weights, scales, momenta = [], [], [] + with gzip.open(path, 'rt') as f: + lines = iter(f) + for line in lines: + if not line.startswith(' t t~`` output shared by the tests: fixed scales, the + PDF set of the xsec references, LHE output and a fixed seed. The first + call also generates (and compiles) it with a reference run.""" + _require_mg7_runtime(self) + cls = type(self) + if cls.reference is None: + mg = MGCmd.MasterCmd() + mg.no_notification() + mg.exec_cmd('set automatic_html_opening False --no_save') + mg.exec_cmd('generate %s' % _PROCESS) + mg.exec_cmd('output mg7 %s' % self.base) + toml = pjoin(self.base, 'Cards', 'run_card.toml') + _edit_run_card(toml, _EVENTS, disable_jet_cuts=False) + text = open(toml).read() + text = _set(text, 'output_format', '"lhe"') + text = _set(text, 'seed', _SEED) + text = text.replace('[multiparticles]\n', '[multiparticles]\ntop = [6, -6]\n') + open(toml, 'w').write(text) + cls.reference = self.run_process(self.base) + return self.base + + def copy_process(self, name): + target = pjoin(self.path, name) + shutil.copytree(self.base_process(), target, symlinks=True) + # the copy starts without the events of the reference run + shutil.rmtree(pjoin(target, 'Events'), ignore_errors=True) + os.mkdir(pjoin(target, 'Events')) + return target + + def edit_card(self, proc_dir, edit): + toml = pjoin(proc_dir, 'Cards', 'run_card.toml') + text = open(toml).read() + open(toml, 'w').write(edit(text)) + + def init_custom_function(self, proc_dir, kind, body, language='cpp'): + """bin/init_custom_func for ``kind`` in ``language``, with ``body``; + the run builds C++ functions.""" + log = pjoin(proc_dir, 'init_custom_func.log') + with open(log, 'w') as logfh: + ret = subprocess.call( + [sys.executable, pjoin(proc_dir, 'bin', 'init_custom_func'), + '-f', kind, '-l', language, '-y', '--no-build'], + cwd=proc_dir, stdout=logfh, stderr=subprocess.STDOUT) + self.assertEqual(ret, 0, 'bin/init_custom_func failed:\n%s' % _tail(log)) + body_file = 'custom_%s.%s' % (kind, _EXTENSIONS[language]) + with open(pjoin(proc_dir, 'custom_%s' % kind, body_file), 'w') as f: + f.write(body) + + def builtin_reference(self, name, edit): + """The run with the built-in setting ``edit``, shared by the languages.""" + cls = type(self) + if not hasattr(cls, 'builtin_runs'): + cls.builtin_runs = {} + if name not in cls.builtin_runs: + builtin = self.copy_process(name) + self.edit_card(builtin, edit) + cls.builtin_runs[name] = self.run_process(builtin) + return cls.builtin_runs[name] + + def copy_for(self, kind, language): + return self.copy_process('%s_%s' % (kind.upper(), language.upper().replace('-', '_'))) + + def run_process(self, proc_dir): + """Run bin/generate_events: (cross section, error, LHE path).""" + log = pjoin(proc_dir, 'mg7_gen.log') + with open(log, 'w') as logfh: + ret = subprocess.call( + [sys.executable, pjoin(proc_dir, 'bin', 'generate_events'), '-f'], + cwd=proc_dir, stdout=logfh, stderr=subprocess.STDOUT) + self.assertEqual(ret, 0, 'mg7 generate_events failed:\n%s' % _tail(log)) + infos = sorted(glob.glob(pjoin(proc_dir, 'Events', '*', 'info.json'))) + self.assertTrue(infos, 'no info.json produced:\n%s' % _tail(log)) + with open(infos[-1]) as f: + info = json.load(f)['process'] + lhe = pjoin(os.path.dirname(infos[-1]), 'events.lhe.gz') + self.assertTrue(os.path.exists(lhe), 'no LHE file produced') + return float(info['mean']), float(info['error']), lhe + + def assertCompatible(self, got, ref, what): + """Cross sections (mean, error) agree within _NSIGMA.""" + sigma = math.hypot(got[1], ref[1]) + self.assertLessEqual( + abs(got[0] - ref[0]), _NSIGMA * sigma, + '%s: %.5g +- %.3g pb against %.5g +- %.3g pb' + % (what, got[0], got[1], ref[0], ref[1])) + + def check_cut(self, language, body): + ref = self.builtin_reference('CUT_BUILTIN', lambda t: t.replace( + '[cuts]\n', '[cuts]\ntop-pt.min = %r\n' % _TOP_PT_CUT)) + custom = self.copy_for('cut', language) + self.init_custom_function(custom, 'cut', body, language) + got = self.run_process(custom) + self.assertCompatible(got, ref, 'custom cut in %s' % language) + + _, _, momenta = _read_lhe(got[2]) + for event in momenta: + for p in event[2:4]: + # up to the precision of the LHE file + self.assertGreaterEqual(_pt(p), _TOP_PT_CUT * (1 - 1e-6)) + + def test_custom_cut_mg7(self): + self.check_cut('cpp', _CUT_BODY) + + def test_custom_cut_numpy_mg7(self): + self.check_cut('python-numpy', _CUT_NUMPY) + + def test_custom_cut_torch_mg7(self): + check = subprocess.run([sys.executable, '-c', 'import numpy, torch'], + capture_output=True) + if check.returncode != 0: + self.skipTest('torch is unavailable') + self.check_cut('python-torch', _CUT_TORCH) + + def check_scale(self, language, body): + def dynamic(text): + text = _set(text, 'fixed_ren_scale', 'false') + text = _set(text, 'fixed_fact_scale', 'false') + return _set(text, 'dynamical_scale_choice', '"half_transverse_mass"') + + ref = self.builtin_reference('SCALE_BUILTIN', dynamic) + custom = self.copy_for('scale', language) + # partonic_energy differs from H_T/2, so the custom function must win + self.edit_card(custom, lambda t: _set( + dynamic(t), 'dynamical_scale_choice', '"partonic_energy"')) + self.init_custom_function(custom, 'scale', body, language) + got = self.run_process(custom) + self.assertCompatible(got, ref, 'custom scale in %s' % language) + + _, scales, momenta = _read_lhe(got[2]) + for scale, event in zip(scales, momenta): + ht = sum(math.sqrt(max(p[0] ** 2 - p[3] ** 2, 0.)) for p in event[2:]) + self.assertAlmostEqual(scale / (ht / 2), 1., places=5) + + def test_custom_scale_mg7(self): + self.check_scale('cpp', _SCALE_BODY) + + def test_custom_scale_numpy_mg7(self): + self.check_scale('python-numpy', _SCALE_NUMPY) + + def test_custom_pdf_mg7(self): + lhapdf_config = misc.resolve_lhapdf().config + if not lhapdf_config: + self.skipTest('lhapdf-config not found') + flags = {} + for flag in ('cflags', 'libs', 'libdir'): + flags[flag] = subprocess.check_output( + [lhapdf_config, '--%s' % flag], text=True).strip() + + custom = self.copy_process('PDF_CUSTOM') + self.init_custom_function(custom, 'pdf', _PDF_BODY) + build = subprocess.run( + ['make', 'CFLAGS=%s' % flags['cflags'], + 'LIBS=%s -Wl,-rpath,%s' % (flags['libs'], flags['libdir'])], + cwd=pjoin(custom, 'custom_pdf'), capture_output=True, text=True) + self.assertEqual(build.returncode, 0, build.stdout + build.stderr) + got = self.run_process(custom) + self.assertCompatible(got, type(self).reference, 'custom pdf') + + def test_custom_pdf_numpy_mg7(self): + check = subprocess.run([sys.executable, '-c', 'import lhapdf'], + capture_output=True) + if check.returncode != 0: + self.skipTest('the lhapdf Python module is unavailable') + custom = self.copy_for('pdf', 'python-numpy') + self.init_custom_function(custom, 'pdf', _PDF_NUMPY, 'python-numpy') + got = self.run_process(custom) + self.assertCompatible(got, type(self).reference, 'custom pdf in numpy') + + def test_custom_bias_mg7(self): + self.check_bias('cpp', _BIAS_BODY) + + def test_custom_bias_numpy_mg7(self): + self.check_bias('python-numpy', _BIAS_NUMPY) + + def check_bias(self, language, body): + custom = self.copy_for('bias', language) + ref = type(self).reference + self.init_custom_function(custom, 'bias', body, language) + got = self.run_process(custom) + self.assertCompatible(got, ref, 'biased cross section in %s' % language) + + def tail(lhe, threshold=250.): + weights, _, momenta = _read_lhe(lhe) + n = len(weights) + in_tail = [w for w, event in zip(weights, momenta) + if _pt(event[2]) > threshold] + mean = sum(weights) / n + error = math.sqrt(sum((w - mean) ** 2 for w in weights) / n / (n - 1)) + xsec = sum(in_tail) / n + xsec_error = math.sqrt(sum(w ** 2 for w in in_tail)) / n + return weights, (mean, error), (xsec, xsec_error), len(in_tail) / n + + ref_weights, _, ref_tail, ref_fraction = tail(ref[2]) + weights, mean, tail_xsec, fraction = tail(got[2]) + self.assertGreater(max(weights) / min(weights), 2., + 'the biased events should not have a constant weight') + self.assertCompatible(mean, ref[:2], 'mean biased event weight') + self.assertCompatible(tail_xsec, ref_tail, 'pT > 250 GeV cross section') + self.assertGreater(fraction, 2 * ref_fraction, + 'the bias should put more events into the tail') + + def check_defaults(self, ms, context, apis): + """The default functions of bin/init_custom_func evaluated through + madspace: 1 for cut and bias, sqrt(s_hat) for scale, the toy pdf.""" + import numpy as np + momenta = np.array([[[500., 0., 0., 500.], [500., 0., 0., -500.], + [500., 100., 50., 200.], [500., -100., -50., -200.]], + [[300., 0., 0., 300.], [200., 0., 0., -200.], + [250., 30., 0., 100.], [250., -30., 0., 0.]]]) + for kind in ('cut', 'scale', 'bias'): + function = ms.CustomFunction( + apis[kind], getattr(ms.CustomFunction, kind), 4, [[21, 21, 6, -6]]) + values = ms.FunctionRuntime(function.function(), context)(momenta) + for value, event in zip(np.array(values), momenta): + if kind == 'scale': + p = [a + b for a, b in zip(event[0], event[1])] + target = math.sqrt(p[0] ** 2 - p[1] ** 2 - p[2] ** 2 - p[3] ** 2) + else: + target = 1. + self.assertAlmostEqual(value, target, places=9) + pdf = ms.PartonDensity(ms.CustomPdf(apis['pdf']), [21, 2], 0) + values = ms.FunctionRuntime(pdf.function(), context)( + np.array([0.1, 0.5]), np.array([100., 100.])) + for x, row in zip([0.1, 0.5], np.array(values)): + toy = math.sqrt(x) * (1 - x) ** 3 + self.assertAlmostEqual(row[0], 3 * toy, places=12) + self.assertAlmostEqual(row[1], 0.5 * toy, places=12) + + def test_init_custom_func_languages_mg7(self): + ms = type(self).ms + if ms is None or not hasattr(ms, 'CustomFunction'): + self.skipTest('madspace with custom functions unavailable') + compilers = {'c': 'cc', 'cpp': 'c++', 'fortran': 'gfortran', + 'cuda': 'nvcc', 'hip': 'hipcc'} + from madgraph.iolibs.template_files.mg7 import init_custom_func + + tested = [] + for language, compiler in compilers.items(): + if shutil.which(compiler) is None: + continue + me_dir = pjoin(self.path, 'LANG_%s' % language) + os.makedirs(me_dir) + for kind in ('cut', 'scale', 'pdf', 'bias'): + ret = init_custom_func.main(me_dir=me_dir, argv=[ + '-f', kind, '-l', language, '-y', '--build', '--no-set-run-card']) + self.assertEqual(ret, 0, '%s %s did not build' % (language, kind)) + self.assertTrue(os.path.exists(pjoin(me_dir, 'lib', 'custom_%s.so' % kind))) + tested.append(language) + if language in ('cuda', 'hip'): + continue + context = ms.Context(device=ms.cpu_device(), thread_count=1) + self.check_defaults(ms, context, { + kind: context.load_matrix_element( + pjoin(me_dir, 'lib', 'custom_%s.so' % kind), '') + for kind in ('cut', 'scale', 'pdf', 'bias')}) + + from madgraph.iolibs.template_files.mg7.custom_function_loader import ( + load_custom_function) + for language in ('python-numpy', 'python-torch'): + if language == 'python-torch': + try: + import torch + except Exception: + continue + me_dir = pjoin(self.path, 'LANG_%s' % language) + os.makedirs(me_dir) + context = ms.Context(device=ms.cpu_device(), thread_count=2) + apis = {} + for kind in ('cut', 'scale', 'pdf', 'bias'): + ret = init_custom_func.main(me_dir=me_dir, argv=[ + '-f', kind, '-l', language, '-y', '--no-set-run-card']) + self.assertEqual(ret, 0, '%s %s failed' % (language, kind)) + path = pjoin(me_dir, 'custom_%s' % kind, 'custom_%s.py' % kind) + self.assertTrue(os.path.exists(path)) + apis[kind] = load_custom_function(ms, [context], ['cpu'], kind, path)[0] + self.check_defaults(ms, context, apis) + tested.append(language) + if not tested: + self.skipTest('no compiler found') diff --git a/tests/input_files/mg7_run_card.toml b/tests/input_files/mg7_run_card.toml index 958426b0a3..22571b4c9e 100644 --- a/tests/input_files/mg7_run_card.toml +++ b/tests/input_files/mg7_run_card.toml @@ -23,6 +23,10 @@ save_gridpack = false include_source = false include_madspace = true include_madspace_source = false +compress = false +run_mode = "minimal" +temp_output_dir = "" +output_dir = "" [beam] e_cm = 1000.0 diff --git a/tests/unit_tests/interface/test_mg7_launch.py b/tests/unit_tests/interface/test_mg7_launch.py index 9eb5e49aaa..8f7fd39788 100644 --- a/tests/unit_tests/interface/test_mg7_launch.py +++ b/tests/unit_tests/interface/test_mg7_launch.py @@ -727,7 +727,10 @@ def test_status_file_path_is_absolute(self): try: process = self.launch.MadgraphProcess.__new__( self.launch.MadgraphProcess) - process.run_card = {"run": {"run_name": "run"}} + process.run_card = { + "run": {"run_name": "run"}, + "gridpack": {"output_dir": "", "temp_output_dir": ""}, + } process.init_event_dir() finally: os.chdir(cwd)