tests/unit_tests/test_edm4hep.py

0001 # test_edm4hep.py
0002
0003 import math
0004 import os
0005
0006 import uproot
0007
0008 from pyrobird.dex import validate_dex
0009 from pyrobird.edm4hep import (
0010     detect_file_type,
0011     edm4hep_entry_to_dict,
0012     edm4hep_to_dex_dict,
0013     get_sim_hit_branches,
0014     sim_hits_to_trajectories,
0015     sim_tracker_hits_to_box_hits,
0016     DEFAULT_HIT_BOX_SIZE,
0017     DEFAULT_TRAJECTORY_EXCLUDED_COLLECTIONS,
0018 )
0019
0020 # Path to the test ROOT files
0021 TEST_EDM4HEP_FILE = os.path.join(os.path.dirname(__file__), 'data', 'k_lambda_10x100_2evt.edm4hep.root')
0022 TEST_EDM4EIC_FILE = os.path.join(os.path.dirname(__file__), 'data', 'reco_2024-09_craterlake_2evt.edm4eic.root')
0023
0024
0025 def open_events_tree(path=TEST_EDM4HEP_FILE):
0026     return uproot.open(path)['events']
0027
0028
0029 def test_detect_file_type():
0030     assert detect_file_type(open_events_tree(TEST_EDM4HEP_FILE)) == 'edm4hep'
0031     # eicrecon/edm4eic files also contain sim hits, reconstruction should win
0032     assert detect_file_type(open_events_tree(TEST_EDM4EIC_FILE)) == 'edm4eic'
0033
0034
0035 def test_get_sim_hit_branches():
0036     tree = open_events_tree()
0037     branches = get_sim_hit_branches(tree)
0038     assert 'SiBarrelHits' in branches
0039     assert 'VertexBarrelHits' in branches
0040     assert 'DIRCBarHits' in branches
0041     assert 'MCParticles' not in branches
0042
0043
0044 def test_sim_tracker_hits_to_box_hits():
0045     tree = open_events_tree()
0046     piece = sim_tracker_hits_to_box_hits(tree, 'SiBarrelHits', entry_start=0)
0047
0048     assert piece['name'] == 'SiBarrelHits'
0049     assert piece['type'] == 'BoxHit'
0050     assert piece['origin']['type'] == 'edm4hep::SimTrackerHitData'
0051     assert piece['count'] > 0
0052
0053     columns = piece['columns']
0054     assert len(columns['pos']) == 3 * piece['count']
0055     # Sim hits have no position error => fixed box size on every axis
0056     assert columns['dim'] == [DEFAULT_HIT_BOX_SIZE] * (3 * piece['count'])
0057     assert len(columns['time']) == piece['count']
0058     assert all(edep > 0 for edep in columns['edep'])
0059     # Sim data has no errors: the error columns are omitted, not zero-filled
0060     assert 'timeError' not in columns
0061     assert 'edepError' not in columns
0062
0063
0064 def test_sim_tracker_hits_box_size_override():
0065     tree = open_events_tree()
0066     piece = sim_tracker_hits_to_box_hits(tree, 'SiBarrelHits', entry_start=0, box_size=5.0)
0067     assert piece['columns']['dim'][:3] == [5.0, 5.0, 5.0]
0068
0069
0070 def test_sim_hits_to_trajectories():
0071     tree = open_events_tree()
0072     hit_branches = [b for b in get_sim_hit_branches(tree)
0073                     if b not in DEFAULT_TRAJECTORY_EXCLUDED_COLLECTIONS]
0074     piece = sim_hits_to_trajectories(tree, hit_branches, entry_start=0)
0075
0076     assert piece['name'] == 'McHitTrajectories'
0077     assert piece['type'] == 'PointTrajectory'
0078     assert sorted(piece['columns'].keys()) == sorted(['pdg', 'charge', 'px', 'py', 'pz', 'p', 'mc_index'])
0079     assert piece['pointColumns'] == ['x', 'y', 'z', 't', 'dx', 'dy', 'dz', 'dt']
0080     assert piece['count'] > 0
0081     assert len(piece['points']) == piece['count']
0082
0083     columns = piece['columns']
0084     for trajectory_i in range(piece['count']):
0085         points = piece['points'][trajectory_i]
0086         # min_hits=2 + prepended vertex point
0087         assert len(points) >= 3
0088         # |p| column matches momentum components (same index in every column)
0089         px, py, pz = columns['px'][trajectory_i], columns['py'][trajectory_i], columns['pz'][trajectory_i]
0090         assert math.isclose(columns['p'][trajectory_i], math.sqrt(px ** 2 + py ** 2 + pz ** 2), rel_tol=1e-9)
0091         # points are time-ordered and each has all point columns
0092         times = [point[3] for point in points]
0093         assert times == sorted(times)
0094         for point in points:
0095             assert len(point) == len(piece['pointColumns'])
0096             # sim data has no per-point errors
0097             assert point[4:] == [0, 0, 0, 0]
0098
0099     # Event 0 of the sample: with Cherenkov collections excluded, exactly 4 particles
0100     # leave >=2 tracker hits: the scattered electron (2), K+ (3) and Lambda decay
0101     # products pi- (5) and proton (6)
0102     assert sorted(columns['mc_index']) == [2, 3, 5, 6]
0103
0104     # The scattered electron: 9 tracker hits + 1 vertex point
0105     electron_i = columns['mc_index'].index(2)
0106     assert columns['pdg'][electron_i] == 11
0107     assert len(piece['points'][electron_i]) == 10
0108
0109
0110 def test_sim_hits_to_trajectories_vertex_and_endpoint_options():
0111     tree = open_events_tree()
0112     hit_branches = [b for b in get_sim_hit_branches(tree)
0113                     if b not in DEFAULT_TRAJECTORY_EXCLUDED_COLLECTIONS]
0114
0115     default = sim_hits_to_trajectories(tree, hit_branches, entry_start=0)
0116     no_vertex = sim_hits_to_trajectories(tree, hit_branches, entry_start=0, prepend_vertex=False)
0117     with_endpoint = sim_hits_to_trajectories(tree, hit_branches, entry_start=0, append_endpoint=True)
0118
0119     for default_points, no_vertex_points, endpoint_points in zip(
0120             default['points'], no_vertex['points'], with_endpoint['points']):
0121         assert len(no_vertex_points) == len(default_points) - 1
0122         assert len(endpoint_points) == len(default_points) + 1
0123         # appended endpoint reuses the last hit time (edm4hep stores no endpoint time)
0124         assert endpoint_points[-1][3] == endpoint_points[-2][3]
0125
0126
0127 def test_edm4hep_entry_to_dict():
0128     tree = open_events_tree()
0129     event = edm4hep_entry_to_dict(tree, entry_index=0)
0130
0131     assert event['id'] == 0
0132     piece_names = [piece['name'] for piece in event['pieces']]
0133
0134     # Hit pieces + one merged trajectory piece
0135     assert 'SiBarrelHits' in piece_names
0136     assert 'McHitTrajectories' in piece_names
0137
0138     # Collections empty in this event are skipped
0139     for piece in event['pieces']:
0140         if piece['type'] == 'BoxHit':
0141             assert piece['count'] > 0
0142
0143     # Cherenkov collections are excluded from trajectories but still present as hits
0144     assert 'DIRCBarHits' in piece_names
0145     trajectory_piece = next(p for p in event['pieces'] if p['name'] == 'McHitTrajectories')
0146     assert 'DIRCBarHits' not in trajectory_piece['origin']['collections']
0147
0148
0149 def test_edm4hep_entry_to_dict_collections_filter():
0150     tree = open_events_tree()
0151
0152     hits_only = edm4hep_entry_to_dict(tree, entry_index=0, collections=['tracker_hits'])
0153     assert all(piece['type'] == 'BoxHit' for piece in hits_only['pieces'])
0154
0155     trajectories_only = edm4hep_entry_to_dict(tree, entry_index=0, collections=['mc_trajectories'])
0156     assert [piece['type'] for piece in trajectories_only['pieces']] == ['PointTrajectory']
0157
0158
0159 def test_edm4hep_to_dex_dict():
0160     tree = open_events_tree()
0161     result = edm4hep_to_dex_dict(tree, [0, 1], origin_info={'file': 'test'})
0162
0163     assert result['type'] == 'firebird-dex-json'
0164     assert result['version'] == '1.0'
0165     assert result['origin'] == {'file': 'test'}
0166     assert len(result['events']) == 2
0167     assert result['events'][0]['id'] == 0
0168     assert result['events'][1]['id'] == 1
0169     for event in result['events']:
0170         assert len(event['pieces']) > 0
0171     # raises on structural problems (column lengths, refs, ragged sizes)
0172     validate_dex(result)