NGTScouting/test/readNano.py

0001 import uproot
0002 import awkward as ak
0003 import numpy as np
0004 import argparse
0005 import sys
0006
0007 def main():
0008     # Argument parsing
0009     parser = argparse.ArgumentParser(description="Process a NanoAOD ROOT file.")
0010     parser.add_argument("filename", help="Path to the NanoAOD ROOT file")
0011     args = parser.parse_args()
0012
0013     file_path = args.filename
0014     tree_name = "Events"
0015
0016     # Open the ROOT file and load the TTree
0017     try:
0018         with uproot.open(file_path) as file:
0019             tree = file[tree_name]
0020             print(tree.keys())
0021
0022             events = tree.arrays(library="ak")
0023             print("Fields:", "\n\t".join(events[0].fields))
0024     except Exception as e:
0025         print(f"Error opening file or reading tree: {e}")
0026         sys.exit(1)
0027
0028     # Event loop
0029     for i, event in enumerate(events):
0030         print(f"Processing event {i}")
0031         print("Found {} tracksters".format(event.ntracksters))
0032         for t_idx in range(event.ntracksters):
0033             offset = event.tracksters_overtices[t_idx]
0034             count = event.tracksters_nvertices[t_idx]
0035             vertices = event.vertices_vertices[offset : offset + count]
0036             vertex_multiplicity = event.vertices_vertex_mult[offset : offset + count]
0037             print(
0038                 t_idx,
0039                 list(zip(vertices, vertex_multiplicity)),
0040                 event.tracksters_raw_energy[t_idx],
0041             )
0042
0043         print("Exploring connections, scores, and sharedEnergy")
0044         print("Connections for {} objects".format(event.nSimTS2TSMergeByHits))
0045
0046         try:  # Offset pattern
0047             offset = 0
0048             for obj_idx in range(event.nSimTS2TSMergeByHits - 1):
0049                 next_offset = event.SimTS2TSMergeByHits_oSimTS2TSMergeByHitsLinks[
0050                     obj_idx + 1
0051                 ]
0052                 elements = event.SimTS2TSMergeByHitsLinks_index[offset:next_offset]
0053                 scores = event.SimTS2TSMergeByHitsLinks_score[offset:next_offset]
0054                 sharedEnergy = event.SimTS2TSMergeByHitsLinks_shardEnergy[
0055                     offset:next_offset
0056                 ]
0057                 if len(elements) > 0:
0058                     print("Offset ", obj_idx, elements, scores, sharedEnergy)
0059                 offset = next_offset
0060         except AttributeError as e:
0061             print(f"An AttributeError occurred (Offset): {e}")
0062
0063         try:  # Count pattern
0064             offset = 0
0065             for obj_idx in range(event.nSimTS2TSMergeByHits):
0066                 count = event.SimTS2TSMergeByHits_nSimTS2TSMergeByHitsLinks[obj_idx]
0067                 elements = event.SimTS2TSMergeByHitsLinks_index[offset : offset + count]
0068                 scores = event.SimTS2TSMergeByHitsLinks_score[offset : offset + count]
0069                 sharedEnergy = event.SimTS2TSMergeByHitsLinks_sharedEnergy[
0070                     offset : offset + count
0071                 ]
0072                 if len(elements) > 0:
0073                     print("Count ", obj_idx, elements, scores, sharedEnergy)
0074                 offset += count
0075         except AttributeError as e:
0076             print(f"An AttributeError occurred (Count): {e}")
0077
0078         try:
0079             for i in range(event.nSimCl2CPWithFraction):
0080                 print(
0081                     "SimCl {} is linked to CP {} with fraction {}".format(
0082                         i,
0083                         event.SimCl2CPWithFraction_index[i],
0084                         event.SimCl2CPWithFraction_fraction[i],
0085                     )
0086                 )
0087         except AttributeError as e:
0088             print(f"An AttributeError occurred (SimCl2CPWithFraction): {e}")
0089
0090 if __name__ == "__main__":
0091     main()