benchmarks/lowq2_reconstruction/LoadData.py

0001 import uproot
0002 import awkward as ak
0003
0004 def create_arrays(dataFiles,featureName="_TaggerTrackerFeatureTensor_floatData",targetName="_TaggerTrackerTargetTensor_floatData", entries=None, treeName="events"):
0005
0006     # List of branches to load
0007     branches = [featureName,targetName]
0008
0009     # Load data from concatenated list of files
0010     data = uproot.concatenate([f"{file}:{treeName}" for file in dataFiles], branches, entry_stop=entries, library="ak")
0011
0012     input_data  = data[featureName]
0013     target_data = data[targetName]
0014
0015     return input_data, target_data