Repository navigation
madspace: multi-step re- and unweighting (surrogates, color reweighting) - #285
Draft
theoheimel wants to merge 41 commits into
Draft
theoheimel wants to merge 41 commits into
theoheimel wants to merge 41 commits into
Conversation
… channel max weights and integral priors
After splitting the backends code, the subfolder "backend" was added to the process folder. This folder should be copied when the full source gridpack is asked.
…_max_weight gridpacks New madspace EventStream: a coordinator thread keeps generation jobs running, buffers each channel's unweighted events in launch order and picks the channel of every event from the current cross-section estimates. Reproducible from the seed, nothing written to disk. Shared helpers factored out of EventGenerator (event_utils) and gridpack.py (gridpack_setup.py).
…(valid .pyi stub), gridpack section in launch test run card
…t-streaming-event-gen # Conflicts: # madgraph/interface/tutorials/mg7.py # madgraph/iolibs/template_files/mg7/run_card.toml
# Conflicts: # madgraph/iolibs/template_files/mg7/gridpack.py
…and the EventStream API reference
…as functions Add UMAMI input keys for particle count, PDG ids, momentum fraction, scale, PDG id and beam index, and output keys CUT, SCALE, PDF and BIAS. A new CustomFunction generator calls a loaded UMAMI library on the event momenta. Cuts, EnergyScale, PartonDensity (via CustomPdf) and Integrand accept these custom functions.
bin/init_custom_func writes custom_<function>/ with a UMAMI wrapper, a function body for one phase-space point (C, C++, Fortran, CUDA or HIP), the UMAMI header and a Makefile that builds lib/custom_<function>.so. Options are taken from the command line or asked interactively. The run card gets a [custom_functions] section with the library paths.
The libraries are built from custom_<kind>/ if present, and loaded before the matrix elements in every context and gridpack. Custom cuts and scales get the PDG ids of the first flavor, the bias those of the sampled flavor. A custom PDF replaces the PDF grid (alpha_s still comes from the PDF set) and disables the systematics.
…eights The integrand weight includes the bias b, so adaptation, unweighting and the channel allocation see |w| b. The signed running integral, which gives the reported cross section, accumulates w / b, and the event records carry b so that the combined events are written with weight w_unw / b.
test_mg7_custom_functions.py compares a custom cut, scale, PDF and bias with the equivalent built-in settings on p p > t t~, and builds and evaluates every function in every language with a toolchain. A new Usage page describes bin/init_custom_func and the function signatures.
UmamiFunctions holds one std::function per UMAMI entry point. The dlopen path fills it from the library symbols, and Context::add_matrix_element takes it directly, so matrix elements can be implemented without a shared library.
Worker threads can then call back into Python. The abort check and the DLPack deleters of Python-owned tensors take the GIL themselves.
ms.PythonFunction calls a Python function on numpy arrays or torch tensors with selectable UMAMI inputs and outputs, and Context.add_matrix_element adds it to a context. The inputs are zero-copy views of the madspace buffers. The results are checked for type, dtype, shape and device before they are copied back. torch functions run on the stream of the GPU runtime.
_madspace_py_loader imports it, so pybind11-stubgen failed to import the staged package in the ENABLE_DOCS builds of CI.
[custom_functions] takes a .so library or a .py file. custom_function_loader loads both, for bin/generate_events and gridpacks. bin/init_custom_func gains the python-numpy and python-torch languages, which need no build. torch is limited to one thread on the CPU, where madspace already splits the batches.
…eat-multi-step-unweighting The bias of the custom functions is now masked in combined_event_layout and divided out in copy_channel_event, so EventStream handles it as well. Custom functions of a gridpack are loaded in gridpack_setup.load_context_data.
The Integrand looks up the matrix-element outputs by name. Only |M|^2 is required: diagram_amp2 only for diagram-based channel weights, and missing color, helicity and diagram indices are written as -1. With return_matrix_element, DifferentialCrossSection also returns |M|^2 itself, which the Integrand passes on as the matrix_element output.
With return_matrix_element, the Integrand also returns the flavor index passed to the matrix element and, with several differential cross sections, which one was evaluated, so that a later step can evaluate its own matrix element on the same events.
Groups of channels can be given a sequence of WeightingSteps, each with its own matrix element, for example a fast surrogate in the integrand followed by the full matrix element. A step multiplies the weights by the ratio of its matrix element to the previous one. An unweighting step then unweights the ratios since the last unweighting step, times the overweight carried from it, with its own maximum weight, truncation and freeze. A reweighting step keeps the ratio in the weight. The events that pass the first unweighting stay on the device and are collected per group and context until a batch is full, or until the end of a round. The decisions are made per channel in job order, with random numbers drawn per event when the events are first unweighted, so the results do not depend on the batching or the number of threads. Integrals, event counts and the sharing of the events between channels refer to the last step, the importance sampling to the first. The color, helicity and diagram indices are taken from the first step that returns them, or replaced with override_outputs. Without groups the generation is unchanged.
Like survey and generate: with matrix elements or custom functions in Python, the worker threads otherwise wait for the GIL held by the calling thread.
weighting_groups_to_json/from_json store the groups with their matrix elements by index, like the channel functions store theirs, so that a gridpack can restore them. With include_estimates, the channel JSON also holds the maximum weights of the steps, which stay fixed when the channel is loaded again. Tested with fix_max_weights() and a reload into a fresh context.
Destroying the stream waits for its running jobs, which can need the GIL, for example with custom functions or matrix elements in Python. Destroyed from Python with the GIL held, this deadlocked.
The EventStream takes the weighting groups like the EventGenerator and runs the events of their channels through the same pipeline, which now hands the events after the last step to a sink: the EventGenerator writes them to the channel files, the stream buffers them as finished jobs. All maxima have to be fixed, so every decision only depends on the event itself. As for the first step, the integral after the steps only includes the jobs whose events were requested, so the stream stays a pure function of the seed. Also passes the biases of the jobs on to their integration, which the stream dropped before.
The default mode of the WeightingStep constructor made pybind11-stubgen fail the wheel build.
A surrogate of the matrix element replaces it in the integrand, and the full
matrix element unweights or reweights the events in a second pass, as one
weighting step per subprocess. The surrogate is a .so library with the inputs
of the madmatrix libraries, loaded per subprocess with {subprocess} and
{device} in its path, or a Python file defining
custom_surrogate(momenta, pdg_ids, alpha_s). In unweight mode, the
first_pass_* settings apply to the unweighting with the surrogate and those of
[generation] to the second pass. Needs the denominators sde_strategy; no
gridpacks yet.
With weighting steps, the run summary shows the efficiency of the first unweighting, before and after cuts, and of each step relative to the one before, next to the overall one, their product. The channels keep showing the overall efficiency. The accepted counts of the steps are now rescaled like the final count when a maximum of an earlier step rises, and reweighting steps count their events too.
UMAMI 1.2 adds the output UMAMI_OUT_MATRIX_ELEMENT_RATIO, the ratio of the matrix element to a reference defined by the caller. A weighting step whose matrix elements return it (MatrixElement::matrix_element_ratio_out) uses the ratio directly, and the ratio times the previous matrix element is the reference of the next step. Integrands reject such matrix elements.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Multi-step re- and unweighting for the madspace
EventGenerator.Groups of channels can get a sequence of re- or unweighting steps, each with its own matrix element. Main use cases are fast surrogates of the matrix element (NN surrogates) and leading-to-full colour reweighting.
Depends on #264 (custom UMAMI functions) and #253 (event streaming). This branch merges both, so the diff against
mainincludes them. The commits of this PR start at the merge commit 301c66d.Weight definitions
unweight_all) stays rejection-only.Implementation
diagram_amp2is needed only for diagram-based channel weights. Colour, helicity and diagram indices come from the first step that returns them, or from later steps withoverride_outputs.fix_max_weightsandEventStreamsupport steps.Fixes found along the way
copy_channel_eventand in the EventStream integration (merge of Custom cut, scale, PDF and bias functions through UMAMI #264 and Gridpacks: in-memory event streaming API (bin/event_stream.py) #253).EventGenerator.optimize/fix_max_weightsand when anEventStreamis destroyed.Testing
madspace/tests/test_weighting_steps.py:fix_max_weights, then a reload from JSON.EventStreamreproducibility and cross section.madspace/tests/test_integrand_me_outputs.py: optional ME outputs.Not yet done
launch.pyintegration, including requestingdiagram_amp2only for the diagram SDE strategy.