pyrobird/tests/test_edm4eic.py

0001 # test_edm4eic.py
0002
0003 import pytest
0004 import os
0005 import json
0006 import uproot
0007 import numpy as np
0008 import awkward as ak
0009 from pyrobird.edm4eic import edm4eic_entry_to_dict
0010 # Helper function to import tracker_hits_to_box_hits
0011 from pyrobird.edm4eic import tracker_hits_to_box_hits
0012
0013 import pytest
0014 from pyrobird.edm4eic import parse_entry_numbers
0015
0016
0017 # Path to the test ROOT file
0018 TEST_ROOT_FILE = os.path.join(os.path.dirname(__file__), 'data', 'reco_2024-09_craterlake_2evt.edm4eic.root')
0019
0020 def test_edm4eic_to_dict_structure():
0021     # Open the ROOT file
0022     file = uproot.open(TEST_ROOT_FILE)
0023     tree = file['events']
0024
0025     # Convert the first event
0026     event = edm4eic_entry_to_dict(tree, entry_index=0)
0027
0028     # Check that the event is a dictionary
0029     assert isinstance(event, dict), "Event should be a dictionary"
0030
0031     # Check that the event has 'id' and "groups" keys
0032     assert 'id' in event, "'id' key missing in event dictionary"
0033     assert 'groups' in event, "'groups' key missing in event dictionary"
0034
0035     # Check that 'data' is a list
0036     assert isinstance(event['groups'], list), "'groups' should be a list"
0037
0038     # Check that each item in 'data' is a dictionary with expected keys
0039     for group in event['groups']:
0040         assert isinstance(group, dict), "Each group in 'groups' should be a dictionary"
0041         assert 'name' in group, "'name' key missing in group"
0042         assert 'type' in group, "'type' key missing in group"
0043         assert 'origin' in group, "'origin' key missing in group"
0044
0045         if group["type"] == "BoxHit":
0046             assert 'hits' in group, "'hits' key missing in group"
0047             assert isinstance(group['hits'], list), "'hits' should be a list"
0048
0049         if group["type"] == "PointTrajectory":
0050             assert "paramColumns" in group
0051             assert "pointColumns" in group
0052             assert "trajectories" in group
0053
0054         # Optionally, check the first hit for expected structure
0055         if 'hits' in group and  len(group['hits']) > 0:
0056             hit = group['hits'][0]
0057             assert isinstance(hit, dict), "Each hit should be a dictionary"
0058             assert 'pos' in hit, "'pos' key missing in hit"
0059             assert 'dim' in hit, "'dim' key missing in hit"
0060             assert 't' in hit, "'t' key missing in hit"
0061             assert 'ed' in hit, "'ed' key missing in hit"
0062
0063             # Check that 'pos' is a list of three floats
0064             assert isinstance(hit['pos'], list) and len(hit['pos']) == 3, "'pos' should be a list of three elements"
0065             assert all(isinstance(x, (float, int)) for x in hit['pos']), "'pos' elements should be numbers"
0066
0067             # Check that 'dim' is a list of three floats
0068             assert isinstance(hit['dim'], list) and len(hit['dim']) == 3, "'dim' should be a list of three elements"
0069             assert all(isinstance(x, (float, int)) for x in hit['dim']), "'dim' elements should be numbers"
0070
0071             # Check that 't' is a list of two floats
0072             assert isinstance(hit['t'], list) and len(hit['t']) == 2, "'t' should be a list of two elements"
0073             assert all(isinstance(x, (float, int)) for x in hit['t']), "'t' elements should be numbers"
0074
0075             # Check that 'ed' is a list of two floats
0076             assert isinstance(hit['ed'], list) and len(hit['ed']) == 2, "'ed' should be a list of two elements"
0077             assert all(isinstance(x, (float, int)) for x in hit['ed']), "'ed' elements should be numbers"
0078
0079 def test_edm4eic_to_dict_values():
0080     # Open the ROOT file
0081     file = uproot.open(TEST_ROOT_FILE)
0082     tree = file['events']
0083
0084     # Convert the first event
0085     event = edm4eic_entry_to_dict(tree, entry_index=0)
0086
0087     # Ensure there is at least one group with hits
0088     assert len(event["groups"]) > 0, "No data groups found in event"
0089
0090     # Get the first group
0091     group = event["groups"][0]
0092     hits = group['hits']
0093     assert len(hits) > 0, "No hits found in the first group"
0094
0095     # Get the first hit
0096     hit = hits[0]
0097
0098     # Uncomment and set the expected values based on your data
0099     assert abs(hit['pos'][0]) > 0
0100     assert len(hit['pos']) == 3
0101     assert abs(hit['dim'][0]) > 0
0102     assert len(hit['dim']) == 3
0103     assert hit['t'][0] > 0
0104     assert len(hit['t']) == 2
0105     assert hit['ed'][0] > 0
0106     assert len(hit['ed']) == 2
0107
0108
0109 def test_edm4eic_to_dict_multiple_entries():
0110     # Open the ROOT file
0111     file = uproot.open(TEST_ROOT_FILE)
0112     tree = file['events']
0113
0114     # Loop over events in the ROOT file (assuming there are at least 2 events)
0115     for entry in range(2):
0116         event = edm4eic_entry_to_dict(tree, entry_index=entry)
0117
0118         # Check that the event number matches
0119         assert event['id'] == entry, f"Event name does not match entry number: {event['id']} != {entry}"
0120
0121         # Perform the same checks as in the previous test
0122         assert "groups" in event, "groups key missing in event dictionary"
0123         assert isinstance(event["groups"], list), "'data' should be a list"
0124         assert len(event["groups"]) > 0, "No data groups found in event"
0125
0126         # Optionally, perform additional checks for each event
0127
0128
0129 def test_tracker_hits_to_box_hits():
0130     # Open the ROOT file
0131     file = uproot.open(TEST_ROOT_FILE)
0132     tree = file['events']
0133
0134     # Get the list of tracker branches
0135     tracker_branches = tree.typenames(recursive=False, full_paths=True, filter_typename="vector<edm4eic::TrackerHitData>")
0136
0137     # Check that there are tracker branches
0138     assert len(tracker_branches) > 0, "No tracker branches of type vector<edm4eic::TrackerHitData> found"
0139
0140     # Test the conversion for each tracker branch
0141     for branch_name in tracker_branches.keys():
0142         group = tracker_hits_to_box_hits(tree, branch_name, entry_start=0)
0143
0144         # Check that the group has expected keys
0145         assert isinstance(group, dict), "Group should be a dictionary"
0146         assert 'name' in group, "'name' key missing in group"
0147         assert 'type' in group, "'type' key missing in group"
0148         assert 'origin' in group, "'origin' key missing in group"
0149         assert 'hits' in group, "'hits' key missing in group"
0150
0151         # Check that 'hits' is a list
0152         assert isinstance(group['hits'], list), "'hits' should be a list"
0153
0154         # Optionally, check the first hit
0155         if len(group['hits']) > 0:
0156             hit = group['hits'][0]
0157             assert isinstance(hit, dict), "Each hit should be a dictionary"
0158             # Perform the same checks as before
0159
0160
0161 @pytest.mark.parametrize("input_value, expected", [
0162     ('3', [3]),
0163     ('1-5', [1, 2, 3, 4, 5]),
0164     ('1,2,3', [1, 2, 3]),
0165     ('1,2-5,8', [1, 2, 3, 4, 5, 8]),
0166     ([1, 2, 3], [1, 2, 3]),
0167     ((1, 2, 3), [1, 2, 3]),
0168     ({1, 2, 3}, [1, 2, 3]),
0169 ])
0170 def test_parse_entry_numbers_valid_inputs(input_value, expected):
0171     assert parse_entry_numbers(input_value) == expected
0172
0173
0174 @pytest.mark.parametrize("input_value", [
0175     ('5-1'),
0176     ('abc'),
0177     ('1,a,3'),
0178 ])
0179 def test_parse_entry_numbers_invalid_inputs(input_value):
0180     with pytest.raises(ValueError):
0181         parse_entry_numbers(input_value)