numpy_call_wrapper

class coffea.ml_tools.numpy_call_wrapper[source]

Bases: ABC

Wrapper for awkward.to_numpy evaluations for dask_awkward array inputs.

For tools outside the coffea package (like for ML inference), the inputs typically expect a numpy-like input. This class wraps up the user-level awkward->numpy data mangling and the underling numpy evaluation calls to recognizable to dask.

For the class to be fully functional, the user must overload these methods:

  • numpy_call: How the evaluation using all numpy tool be performed

  • prepare_awkward: How awkward arrays should be translated to the a numpy format that is compatible with the numpy_call

Additionally, the following helper functions can be omitted, but will help assist with either code debugging or for data mangling niceties.

  • validate_numpy_input: makes sure the computation routine understand the input.

  • numpy_to_awkward: Additional translation to convert numpy outputs to awkward (defaults to a simple awkward.from_numpy conversion)

Methods Summary

__call__(*args, **kwargs)

Highest level abstraction to be directly called by the user.

get_awkward_lib(*args, **kwargs)

numpy_call(*args, **kwargs)

Underlying numpy-like evaluation.

postprocess_awkward(return_array, *args, ...)

Additional conversion from the numpy_call output back to awkward arrays.

prepare_awkward(*args, **kwargs)

Converting awkward-array like inputs into be numpy-compatible awkward-arrays compatible with the numpy_call method.

validate_numpy_input(*args, **kwargs)

Validating that the numpy-like input arguments are compatible with the underlying evaluation calls.

Methods Documentation

__call__(*args, **kwargs)[source]

Highest level abstraction to be directly called by the user. Checks whether the inputs has any awkward arrays or dask_awkward arrays, and call the corresponding function if they are found. If no dask awkward or awkward arrays are found, calling the underlying _call_numpy method.

get_awkward_lib(*args, **kwargs)[source]
abstract numpy_call(*args, **kwargs)[source]

Underlying numpy-like evaluation. This method should be reimplemented by the user or by tool-specialized classes.

postprocess_awkward(return_array, *args, **kwargs)[source]

Additional conversion from the numpy_call output back to awkward arrays. This method does not need to need to be overloaded, but can make the data-mangling that occurs outside the class cleaner (ex: additional awkward unflatten calls). To ensure that the data mangling can occur, the unformatted awkward-like inputs are also passed to this function.

For the base method, we will simply iterate over the containers and call the default awkward.from_numpy conversion

abstract prepare_awkward(*args, **kwargs) Tuple[source]

Converting awkward-array like inputs into be numpy-compatible awkward-arrays compatible with the numpy_call method. The actual conversion to numpy is handled automatically. The return value should be (*args, **kwargs) pair that is compatible with the numpy_call.

Consult the following documentation to find the awkward operations needed. https://awkward-array.org/doc/main/user-guide/how-to-restructure-pad.html

validate_numpy_input(*args, **kwargs) None[source]

Validating that the numpy-like input arguments are compatible with the underlying evaluation calls. This function should raise an exception if invalid input values are found. The base method performs no checks but raises a warning that no checks were performed.