tests/unit_tests/test_mc_particles.py

0001 # test_mc_particles.py
0002
0003 import math
0004 import os
0005
0006 import uproot
0007
0008 from pyrobird.dex import validate_dex
0009 from pyrobird.edm4eic import edm4eic_to_dex_dict
0010 from pyrobird.edm4hep import edm4hep_to_dex_dict
0011 from pyrobird.mc_particles import (
0012     interpolate_line_points,
0013     mc_particles_to_trajectories,
0014     C_LIGHT,
0015     DEFAULT_MC_MAX_POINTS,
0016 )
0017
0018 TEST_EDM4HEP_FILE = os.path.join(os.path.dirname(__file__), 'data', 'k_lambda_10x100_2evt.edm4hep.root')
0019 TEST_EDM4EIC_FILE = os.path.join(os.path.dirname(__file__), 'data', 'reco_2024-09_craterlake_2evt.edm4eic.root')
0020
0021
0022 def open_events_tree(path):
0023     return uproot.open(path)['events']
0024
0025
0026 def test_interpolate_line_points_grid():
0027     # beta=1 over 3 x c_light mm takes 3 ns; step 1 ns => vertex + 2 interior + endpoint
0028     points = interpolate_line_points(0, 0, 0, 3 * C_LIGHT, 0, 0,
0029                                      start_time=10.0, beta=1.0, step_time=1.0)
0030     assert len(points) == 4
0031     assert points[0] == [0, 0, 0, 10.0]
0032     assert points[-1][0] == 3 * C_LIGHT
0033     assert math.isclose(points[-1][3], 13.0)
0034     # Interior points sit on the fixed grid with uniform spacing
0035     assert math.isclose(points[1][3], 11.0)
0036     assert math.isclose(points[2][3], 12.0)
0037     assert math.isclose(points[1][0], C_LIGHT)
0038
0039
0040 def test_interpolate_line_points_max_points_cap():
0041     # A flight needing thousands of grid steps is coarsened, not truncated:
0042     # the endpoint stays exact and the subdivision stays uniform
0043     points = interpolate_line_points(0, 0, 0, 100000, 0, 0,
0044                                      start_time=0.0, beta=1.0, step_time=0.001, max_points=16)
0045     assert len(points) == 16
0046     assert points[-1][0] == 100000
0047     steps = [points[i + 1][3] - points[i][3] for i in range(len(points) - 1)]
0048     assert all(math.isclose(step, steps[0]) for step in steps)
0049
0050
0051 def test_interpolate_line_points_at_rest():
0052     # No momentum => no motion; both end points carry the creation time
0053     points = interpolate_line_points(1, 2, 3, 1, 2, 3, start_time=5.0, beta=0.0)
0054     assert points == [[1, 2, 3, 5.0], [1, 2, 3, 5.0]]
0055
0056
0057 def test_mc_particles_to_trajectories_edm4hep():
0058     tree = open_events_tree(TEST_EDM4HEP_FILE)
0059     piece = mc_particles_to_trajectories(tree, 0)
0060
0061     assert piece['name'] == 'MCParticles'
0062     assert piece['type'] == 'PointTrajectory'
0063     assert piece['pointColumns'] == ['x', 'y', 'z', 't']
0064     # Every MCParticle becomes a line: count equals the collection size and
0065     # trajectory id equals the particle index
0066     mc_count = len(tree['MCParticles/MCParticles.PDG'].array(entry_start=0, entry_stop=1)[0])
0067     assert piece['count'] == mc_count
0068     assert len(piece['points']) == mc_count
0069     for column in ('pdg', 'charge', 'gen_status', 'px', 'py', 'pz', 'p', 'mass'):
0070         assert len(piece['columns'][column]) == mc_count
0071
0072     for points in piece['points']:
0073         assert 2 <= len(points) <= DEFAULT_MC_MAX_POINTS
0074         times = [point[3] for point in points]
0075         assert times == sorted(times)
0076
0077
0078 def test_mc_particles_to_trajectories_edm4eic():
0079     # eicrecon output carries the same MCParticles collection
0080     tree = open_events_tree(TEST_EDM4EIC_FILE)
0081     piece = mc_particles_to_trajectories(tree, 0)
0082     assert piece['count'] > 0
0083     assert piece['count'] == len(piece['points'])
0084
0085
0086 def test_mc_particles_in_default_conversions():
0087     # Both model paths include the mc_particles group by default and the
0088     # documents stay structurally valid
0089     hep_dex = edm4hep_to_dex_dict(open_events_tree(TEST_EDM4HEP_FILE), 0)
0090     validate_dex(hep_dex)
0091     assert any(p['name'] == 'MCParticles' for p in hep_dex['events'][0]['pieces'])
0092
0093     eic_dex = edm4eic_to_dex_dict(open_events_tree(TEST_EDM4EIC_FILE), 0)
0094     validate_dex(eic_dex)
0095     assert any(p['name'] == 'MCParticles' for p in eic_dex['events'][0]['pieces'])
0096
0097
0098 def test_mc_particles_collections_filter():
0099     # Selecting only mc_particles yields exactly the one piece
0100     dex = edm4hep_to_dex_dict(open_events_tree(TEST_EDM4HEP_FILE), 0, collections=['mc_particles'])
0101     pieces = dex['events'][0]['pieces']
0102     assert len(pieces) == 1
0103     assert pieces[0]['name'] == 'MCParticles'
0104
0105     # And leaving it out keeps it out
0106     dex = edm4eic_to_dex_dict(open_events_tree(TEST_EDM4EIC_FILE), 0, collections=['tracker_hits'])
0107     assert not any(p['name'] == 'MCParticles' for p in dex['events'][0]['pieces'])