MF Onnx and cleaner output predictions
Description
Getting MF to work for onnx export:
- Modifies the main MF code slightly to allow for times when we have 0 tracks
- Fix some of the lines that were causing other issues for onnx
- Implements onnx outputs, currently we write the 'leading' object as scalers to decorate the jet. Here leading is defined as the object with the highest predicted regression[0], which is pT for the vertexing case
- Modified the mask_to_indices function to return indices in the same format as GN2 vertexing. That is, first find any tracks which the model defines as a vertex. Then, for all remaining non-padded tracks, use a consecutively increasing index. E.g 1 mask
[False, True, True, False, False] -> [1, 0, 0, 2, 3]
as we have vertex0
atidx=1,2
and all remaining are not part of a vertex. - Also modify the prediction writer to output in the same format.
This MR relates to #53. New issue added - need to be eventually able to write all regression values, rather than just leading regression. But this will require some work.
Review checklist:
-
CI Passing -
Comments addressed -
Source branch is up to date with target
Edited by Nikita Ivvan Pond