PythonFunction

class madspace.PythonFunction(function, inputs, outputs, framework='numpy', particle_count=None, name=None, parameters=None)

Bases: object

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.

Parameters:

  • 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.

  • parameters – Names and default values of parameters set through umami_set_parameter, e.g. the parameters argument of Context.add_matrix_element. With parameters, function gets a namespace of their values as last argument, e.g. params.pt_min.