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