AmpliCol matrix elements in mg7 ([amplicol] processes)

.github/workflows/docs_examples.yml

@@ -3,13 +3,17 @@ name: MadSpace docs examples
 # The other MG7 acceptance workflows all carry `paths-ignore: ['docs/**']`, so a
 # docs-only commit currently triggers none of them. This workflow does the
 # opposite: it triggers ONLY on the paths that can affect the documentation
-# examples, so editing docs/source/madspace/examples/*.rst is actually tested.
+# examples, so editing docs/source/madspace/examples/*.rst (or the event
+# streaming example in docs/source/event_stream.rst) is actually tested.
 on:
   push:
     branches:
       - '**'
     paths:
       - 'docs/source/madspace/examples/**'
+      - 'docs/source/event_stream.rst'
+      - 'madgraph/iolibs/template_files/mg7/**'
+      - 'madgraph/mg7/**'
       - 'madspace/**'
       - 'tests/acceptance_tests/test_madspace_examples.py'
       - '.github/workflows/docs_examples.yml'
@@ -18,6 +22,9 @@ on:
     types: [opened, synchronize, reopened]
     paths:
       - 'docs/source/madspace/examples/**'
+      - 'docs/source/event_stream.rst'
+      - 'madgraph/iolibs/template_files/mg7/**'
+      - 'madgraph/mg7/**'
       - 'madspace/**'
       - 'tests/acceptance_tests/test_madspace_examples.py'
       - '.github/workflows/docs_examples.yml'
@@ -55,4 +62,4 @@ jobs:
       - name: run the madspace doc examples
         run: |
             cd $GITHUB_WORKSPACE
-            ./tests/test_manager.py test_madspace_example_simple_mappings test_madspace_example_cuts test_madspace_example_diagram_mapping test_madspace_example_integration_order test_madspace_example_matrix_element test_madspace_example_pdf test_madspace_example_integrator test_madspace_example_integrator_madnis test_madspace_example_flow_training -pA -t0 -l INFO
+            ./tests/test_manager.py test_madspace_example_simple_mappings test_madspace_example_cuts test_madspace_example_diagram_mapping test_madspace_example_integration_order test_madspace_example_matrix_element test_madspace_example_pdf test_madspace_example_integrator test_madspace_example_integrator_madnis test_madspace_example_flow_training test_gridpack_example_event_stream -pA -t0 -l INFO

.github/workflows/pyamplicol.yml

@@ -73,7 +73,7 @@ jobs:
           if [[ -z "$base" || "$base" =~ ^0+$ ]] || ! git cat-file -e "$base^{commit}" 2>/dev/null; then
             changed=true
           elif git diff --name-only "$base" HEAD \
-              | grep -qE '^(pyamplicol/|\.github/workflows/pyamplicol|\.github/actions/(pyamplicol|build_pyamplicol|install_pyamplicol)|VERSION$)'; then
+              | grep -qE '^(pyamplicol/|\.github/workflows/pyamplicol|\.github/actions/(pyamplicol|build_pyamplicol|install_pyamplicol)|madspace/include/madspace/umami\.h$|madgraph/mg7/bootstrap\.py$|VERSION$)'; then
             changed=true
           else
             changed=false
@@ -173,7 +173,7 @@ jobs:
           python - <<'EOF'
           import sys
           sys.path.insert(0, ".")
-          from madgraph.iolibs.template_files.mg7 import bootstrap
+          from madgraph.mg7 import bootstrap
           bootstrap.ensure_pyamplicol(interactive=False)
           bootstrap.check_pyamplicol_license()
           assert bootstrap.check_pyamplicol_source_hash()

.github/workflows/release.yml

@@ -227,12 +227,16 @@ jobs:
           done
           test -d madspace/install/madspace
 
-          # Also checks that Symbolica accepts pyamplicol's library key.
+          # Also checks that Symbolica accepts pyamplicol's library key; only the
+          # wheel download is retried, so a rejected key fails the step.
           for attempt in $(seq 10); do
-            python3 pyamplicol/install.py --bin -y && break
+            if python3 pyamplicol/install.py --bin -y; then status=0; break; else status=$?; fi
+            # The wheel installed, so the license check failed: retrying cannot help.
+            test -d pyamplicol/install/pyamplicol && break
             echo "attempt $attempt: pyamplicol==${{ needs.check_version.outputs.version }} not installable yet, retrying..."
             sleep 15
           done
+          test "$status" -eq 0
           test -d pyamplicol/install/pyamplicol
 
       - name: Sign the tarball with Sigstore

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:

docs/source/custom_functions.rst

@@ -290,6 +290,27 @@ matrix products. If your function relies on it, set ``OMP_NUM_THREADS=1``.
 The Python API behind these functions is ``madspace.PythonFunction``. It can also be used
 directly with madspace.
 
+.. _custom-function-surrogate:
+
+Python surrogates
+-----------------
+
+A Python surrogate of the matrix element (see :ref:`run-card-surrogate`) uses the same
+mechanism. Its file defines ``FRAMEWORK`` and ``custom_surrogate(momenta, pdg_ids,
+alpha_s)``, which returns :math:`|M|^2` as a float64 array of shape ``(batch,)``. It takes
+no parameters. ``momenta`` and ``pdg_ids`` have the shapes and types given above, with the
+momenta that the full matrix element gets, and ``alpha_s`` is a float64 array of shape
+``(batch,)``. The function is called for all subprocesses, which ``pdg_ids`` identifies:
+
+.. code-block:: python
+
+    import numpy as np
+
+    FRAMEWORK = "numpy"
+
+    def custom_surrogate(momenta, pdg_ids, alpha_s):
+        return np.ones(len(momenta))
+
 Examples
 --------
 

docs/source/event_stream.rst

@@ -0,0 +1,159 @@
+Streaming events from Python
+============================
+
+A gridpack can generate events directly into a Python program, without writing any files.
+This is useful when the events are consumed right away, for example by an analysis written
+in Python or as training data for a neural network. The ``EventStream`` class in the
+gridpack's ``bin/event_stream.py`` returns the events either as LHE event objects or as
+batches of NumPy arrays or PyTorch tensors.
+
+Event streaming needs a gridpack created with ``run_mode = "fix_max_weight"`` (see
+:ref:`gridpack-run-mode`). In this mode, the maximum weight of every channel is fixed when
+the gridpack is created. An event is final as soon as it is unweighted, so it can be returned
+right away without a survey or a combination step.
+
+Example: the e+ e- invariant mass
+---------------------------------
+
+This example fills a histogram of the invariant mass of the electron pair in
+``p p > e+ e-``. First, create the gridpack at the MadGraph7 prompt:
+
+.. code-block:: text
+    :name: event-stream-mg5-commands
+
+    generate p p > e+ e-
+    output mg7 PROC_ppee
+    launch PROC_ppee
+    set save_gridpack true
+    set run_mode fix_max_weight
+    done
+
+The gridpack is written to ``PROC_ppee/Events/run_01/gridpack``. Its ``bin`` directory
+contains ``event_stream.py``, so add it to the Python path and import ``EventStream``. The
+script can run from any directory.
+
+.. code-block:: python
+
+    import sys
+
+    import numpy as np
+
+    sys.path.insert(0, "PROC_ppee/Events/run_01/gridpack/bin")
+    from event_stream import EventStream
+
+Creating an ``EventStream`` starts the event generation in background threads. Its
+arguments, like the device, the seed and the batch sizes, default to the values in the
+gridpack's ``Cards/grid_run_card.toml``. The stream is endless, so ``max_events`` limits
+the number of events. Each batch is an :py:class:`EventBatch <event_stream.EventBatch>`
+with one array per LHE field. The particle fields have the shape
+``(batch_size, max_particle_count)``.
+
+.. code-block:: python
+
+    bins = np.linspace(60.0, 120.0, 31)
+    hist = np.zeros(len(bins) - 1)
+    weight_sum = 0.0
+
+    with EventStream(seed=42) as stream:
+        for batch in stream.batches(10000, max_events=100000):
+            # four-momentum of the e+ e- pair: sum over the outgoing electrons
+            is_electron = (np.abs(batch.pdg_id) == 11) & (batch.status_code == 1)
+            e, px, py, pz = (
+                np.sum(np.where(is_electron, p, 0.0), axis=1)
+                for p in (batch.energy, batch.px, batch.py, batch.pz)
+            )
+            m_ee = np.sqrt(np.maximum(e**2 - px**2 - py**2 - pz**2, 0.0))
+            hist += np.histogram(m_ee, bins, weights=batch.weight)[0]
+            weight_sum += np.sum(batch.weight)
+        cross_section = stream.integral
+        print(f"cross section: {cross_section:.1f} +- {stream.integral_error:.1f} pb")
+
+The event weights are not normalized to the cross section (see :ref:`event-stream-weights`).
+Scaling them to the cross section gives the differential cross section in pb/GeV:
+
+.. code-block:: python
+
+    dsigma = hist * cross_section / weight_sum / np.diff(bins)
+    for low, value in zip(bins[:-1], dsigma):
+        print(f"{low:5.0f} GeV {value:8.3f} pb/GeV " + "#" * int(50 * value / dsigma.max()))
+
+    peak = np.argmax(dsigma)
+    assert bins[peak] < 91.2 < bins[peak + 1], "the Z peak is missing"
+
+The output shows the Z peak (shortened)::
+
+    cross section: 880.2 +- 0.6 pb
+    ...
+       84 GeV    7.569 pb/GeV ##
+       86 GeV   16.543 pb/GeV #####
+       88 GeV   51.287 pb/GeV ################
+       90 GeV  158.376 pb/GeV ##################################################
+       92 GeV   66.743 pb/GeV #####################
+       94 GeV   18.457 pb/GeV #####
+    ...
+
+Output formats
+--------------
+
+There are three ways to get events from a stream:
+
+- Iterating over the stream yields one ``madspace.LHEEvent`` at a time. The events are
+  completed in chunks of 1000 internally. If the loop is left early, the rest of the chunk
+  is returned first by the next call, so no event is skipped.
+- :py:meth:`next_events <event_stream.EventStream.next_events>` returns a list of a given
+  number of ``LHEEvent`` objects.
+- :py:meth:`batches <event_stream.EventStream.batches>` and
+  :py:meth:`next_batch <event_stream.EventStream.next_batch>` return
+  :py:class:`EventBatch <event_stream.EventBatch>` objects. With ``output_type="torch"``,
+  their fields are PyTorch tensors on the CPU instead of NumPy arrays.
+
+Events have different numbers of particles if the process has several subprocesses or
+includes decays. The particle fields are then padded with zeros. ``particle_count`` holds
+the number of particles of every event and ``particle_mask`` marks the valid entries.
+
+.. _event-stream-weights:
+
+Weights and cross section
+-------------------------
+
+After the unweighting, every event weight is ``1``, or ``-1`` for an event with a negative
+weight. The exception are overweight events, whose weight exceeded the fixed maximum weight
+of their channel. They keep the ratio of their weight to the maximum weight, so their weight
+is larger than ``1``. Treating them as ``1`` would bias the distributions in the phase-space
+regions where they occur.
+
+The events of ``[amplicol]`` subprocesses are then reweighted to full colour in the default
+``mode = "reweight"``: their weight is multiplied by the ratio of the full-colour to the
+leading-colour matrix element, so it is not ``1``. With ``mode = "unweight"``, they are
+unweighted a second time and have unit weights apart from overweights.
+
+In all cases, to normalize a sample to the cross section, scale its weights by
+``stream.integral / sum(weights)``. The cross section estimate improves while events are
+generated: it combines the integrals stored in the gridpack with the weighted events
+generated by the stream so far. Read it after generating the sample, as in the example
+above.
+
+How it works
+------------
+
+Each channel of the gridpack keeps a buffer of unweighted events, which background threads
+refill whenever it runs low. The channel of every event is drawn at random, with
+probabilities given by the current cross section estimates of the channels. Generating
+the phase-space points and evaluating the matrix elements happen in separate steps, as in a
+normal gridpack run.
+
+The events are reproducible: for a given seed, the stream always returns the same sequence of
+events. It does not depend on the number of threads or on how the events are split into
+batches.
+
+API reference
+-------------
+
+.. autoclass:: event_stream.EventStream
+    :members:
+    :special-members: __iter__
+    :member-order: bysource
+
+.. autoclass:: event_stream.EventBatch
+    :members:
+    :member-order: bysource

docs/source/gridpacks.rst

@@ -70,6 +70,7 @@ ends:
     longer, especially for channels with a low unweighting efficiency, but the gridpack
     runs do not have to estimate the maximum weights first. This helps if each run only
     generates a few events.
+    The gridpack also gets ``bin/event_stream.py``, see :ref:`gridpack-event-stream`.
 
 ``regular``
     The events of the run are generated as usual, then the gridpack is stored.
@@ -96,6 +97,15 @@ Make sure every job uses a different seed. The default seed ``-1`` draws a fresh
 seed for each call, which is fine for independent jobs. The seed that was used is recorded
 in the output.
 
+.. _gridpack-event-stream:
+
+Streaming events from Python
+----------------------------
+
+A gridpack created with ``run_mode = "fix_max_weight"`` also contains
+``bin/event_stream.py``. Its ``EventStream`` class generates events in memory and returns
+them to a Python program, without writing any files. See :doc:`event_stream`.
+
 Run card settings
 -----------------
 
@@ -147,6 +157,9 @@ gridpack is created.
     * -
       - ``cut_efficiency_threshold``, ``max_cut_repetitions``
       - Card only.
+    * - ``[amplicol]``
+      - ``first_pass_max_overweight_truncation``, ``first_pass_freeze_max_weight_after``
+      - ``--first_pass_max_overweight_truncation``, ``--first_pass_freeze_max_weight_after``
     * - ``[gridpack]``
       - ``output_dir``, ``temp_output_dir``
       - ``--output_dir``, ``--temp_output_dir``
@@ -160,6 +173,14 @@ the main run card, see :doc:`run_card`. ``./bin/generate_events --help`` lists a
 Relative ``output_dir`` and ``temp_output_dir`` paths in the card refer to the gridpack
 directory. Both are reset to ``""`` when the gridpack is created.
 
+With ``[amplicol]`` processes unweighted to full colour, ``max_overweight_truncation`` and
+``freeze_max_weight_after`` set the second pass, which gives the events, and the
+``first_pass_`` options the leading-colour unweighting of these processes, as in the main
+run. The other processes are unweighted once, as set by ``[generation]``. The colour
+accuracy and the mode are fixed when the gridpack is created. The pyamplicol artifacts
+the libraries read are copied to ``SubProcesses/``, and ``include_source = true`` is not
+supported for these processes.
+
 The card names the ``cpu_mode`` that the gridpack was built with, which can differ from the
 ``"auto"`` of the original run card. The matrix element libraries carry this name. If you
 change ``cpu_mode`` or ``precision``, you have to rebuild them.

docs/source/index.rst

@@ -30,6 +30,7 @@ If you need a stable release, use the
    custom_functions
    diagram_filters
    gridpacks
+   event_stream
    madevent_to_mg7
 
 .. toctree::

docs/source/madspace/cpp-api.rst

@@ -70,6 +70,7 @@ Builders that assemble compute-graph functions for channel weights, energy scale
    cpp/SubchannelWeights
    cpp/Unweighter
    cpp/VegasHistogram
+   cpp/WeightingStepFunction
 
 Phase-space building blocks
 ---------------------------
@@ -85,6 +86,7 @@ Supporting types shared by the mappings: topologies, cuts, observables and densi
    cpp/PdfGrid
    cpp/Propagator
    cpp/Topology
+   cpp/WeightingStep
    cpp/initialize_vegas_grid
 
 Compute graph
@@ -127,6 +129,7 @@ The high-level driver plus the tensor and device layer it runs on: integration,
    cpp/EventHistogramSpec
    cpp/EventHistograms
    cpp/EventRecord
+   cpp/EventStream
    cpp/FieldLayout
    cpp/GeneratorBatchJob
    cpp/GeneratorConfig
@@ -157,6 +160,10 @@ The high-level driver plus the tensor and device layer it runs on: integration,
    cpp/Runtime
    cpp/Sizes
    cpp/StatusFile
+   cpp/StepChainResult
+   cpp/StepEvents
+   cpp/StepIntegral
+   cpp/StepSettings
    cpp/SubprocessObservables
    cpp/SubprocessSystArgs
    cpp/SystematicsCalculator
@@ -171,6 +178,8 @@ The high-level driver plus the tensor and device layer it runs on: integration,
    cpp/UnalignedRef
    cpp/Variation
    cpp/VegasGridOptimizer
+   cpp/WeightingGroup
+   cpp/WeightingPipeline
    cpp/available_backends
    cpp/compute_generation_batch_event_count
    cpp/cpu_device
@@ -186,6 +195,8 @@ The high-level driver plus the tensor and device layer it runs on: integration,
    cpp/set_simd_mode
    cpp/supported_simd_modes
    cpp/set_lib_path
+   cpp/weighting_groups_from_json
+   cpp/weighting_groups_to_json
 
 Utilities
 ---------