[theoheimel]

Custom cuts, scales, parton densities and event biases for MG7 runs. Each one is either a UMAMI shared library or a Python file using numpy or torch, called by madspace in batches like a matrix element.

All steps are in. CUDA and HIP builds, and torch functions on GPUs, are untested (no GPU was available).

Design

• New UMAMI keys (minor version 1).
  • Inputs: PARTICLE_COUNT, PDG_IDS, MOMENTUM_FRACTION, SCALE, PDG_ID, BEAM_INDEX.
  • Outputs: CUT, SCALE, PDF, BIAS, all doubles.
• Cut: the returned value multiplies the cut mask (0 rejects the event).
• Scale: one value sets the renormalization and both factorization scales, like MG5's user_dynamical_scale. scale_factor and the fixed-scale options still apply.
• PDF: returns x·f(x, Q) for one beam and one parton. alpha_s still comes from the run card's PDF set.
• Bias: works like MG5's bias. Adaptation and unweighting see w·b, the reported cross section uses w, and written events carry w/b.
• Python functions:
  • UmamiFunctions, a table of std::functions, builds a MatrixElementApi without a shared library, through Context::add_matrix_element.
  • ms.PythonFunction implements UMAMI in Python with selectable inputs and outputs. It passes zero-copy numpy or torch views, and checks the type, dtype, shape and device of the results. torch runs on the GPU stream of the runtime. numpy is CPU only.
  • madspace now releases the GIL while runtimes, generators and MadNIS training run, so worker threads can call Python. The abort check and the DLPack deleters of Python-owned tensors take the GIL themselves.
  • With torch on the CPU, MG7 sets torch to one thread, since madspace already splits the batches.

Checklist

• UMAMI keys, madspace CustomFunction and CustomPdf, options on Cuts / EnergyScale / PartonDensity / Integrand, Python bindings

• Bias handling in the event generator (unbiased cross section, unweighted weights divided by the bias)

• bin/init_custom_func and templates (C, C++, Fortran, CUDA, HIP)

• Run card [custom_functions] section and launch.py wiring

• Acceptance tests

• Documentation page under Usage

• UmamiFunctions and Context::add_matrix_element

• GIL release in the madspace bindings

• ms.PythonFunction and its unit tests

• .so or .py in [custom_functions], python-numpy / python-torch in bin/init_custom_func, gridpack support

• Acceptance tests and documentation for Python functions

[oliviermattelaer]

One comment on this:

• Cut and scale see the event before the mirror flip; the bias and the written event see it after (integrand.cpp:586).
• Systematics with a custom scale. resolve_dynamical_scales (launch.py:1024) still treats dynamical_scale_choice as the scale the events were generated with, so it silently drops that variation.
• Gridpack with an absolute path. The file is not copied and data.json keeps the absolute path, with no warning (launch.py:1988). This contradicts the docs.
• need to check with a GPU (claude reports three potential GPU issue)
  • torch gets tensors labelled device 0 even on cuda:1 (python_function.cpp:49).
  • The GPU wrapper uses fixed 32-particle arrays with no bounds check (wrapper_gpu.cu:56).
  • A CPU .so in a GPU run only fails at the first batch, with “Matrix element has incompatible device”, which doesn’t name the custom function. Check the device when the function is loaded.

One limitation, is that all parameter are hardcoded in the cuts/...
Would be nice to be able (like in MG5aMC) to define parameter such that you include the custom function (like a DJR cut function/...) and then specify the cut value in the run_card (such that we split the function from the cut.

Other points that we should consider is to allow to have different mu_f1, mu_f2 and mu_r

One suggestion would be to add tt~j with custom eta cut to check the mirroring part

[theoheimel]

@oliviermattelaer Thanks for the feedback!

• I can address the first point once the mirror-related stuff is merged (it probably solves itself then)
• I'll check the other issues. I indeed have not tested it on GPU so far
• Allowing for custom parameters is a good idea! Since the custom functions work through UMAMI, all the infrastructure for setting parameters is there already.

[oliviermattelaer]

Excellent