From 8e0c1b04d4725099e9b44f2fc1d57162b035f39e Mon Sep 17 00:00:00 2001 From: Thomas Madlener Date: Mon, 16 Sep 2024 16:34:26 +0200 Subject: [PATCH] Make sure that files with the complete datamodel can be read (#361) * Move coll creation into type specific functions - Easier for testing if the counter is reset for every type - Easier to read / maintain * Make creation functions importable * Capitalize global variables * Fix missing arguments and returns * Import argparse and sys only when necessary * Make sure to write all vector members of tracks * Add check for complete EDM4hep file * Run tests also for RNTuple files * Also check Link collections * Workaround pythonizations changing mutability of objects * Rename tests to make it easier to define dependencies * Move tests to more appropriate place * Make sure to only run pytest on things it can run on * Move imports to the top * Check Link collections in the same order as in yaml file * Don't start the counter at 0 to avoid hitting empty defaults * Apply suggestions from code review Co-authored-by: Mateusz Jakub Fila <37295697+m-fila@users.noreply.github.com> * Ignore regex matches from Dart printouts --------- Co-authored-by: Mateusz Jakub Fila <37295697+m-fila@users.noreply.github.com> --- scripts/createEDM4hepFile.py | 389 ++++++++++++++++++--------- test/CMakeLists.txt | 34 ++- test/conftest.py | 8 + test/pytest.ini | 2 + test/test_EDM4hepFile.py | 508 +++++++++++++++++++++++++++++++++++ 5 files changed, 804 insertions(+), 137 deletions(-) create mode 100644 test/conftest.py create mode 100644 test/pytest.ini create mode 100644 test/test_EDM4hepFile.py diff --git a/scripts/createEDM4hepFile.py b/scripts/createEDM4hepFile.py index 07fa19220..b532e174c 100755 --- a/scripts/createEDM4hepFile.py +++ b/scripts/createEDM4hepFile.py @@ -3,54 +3,23 @@ # Description: Create a file with EDM4hep data using the EDM4hep python bindings # The purpose of the script is to use all the classes of the EDM4hep library # to create a file with dummy data. -import podio -import edm4hep -from itertools import count + import argparse import sys +from itertools import count -frames = 3 # How many frames or events will be written -vectorsize = 5 # For vector members, each vector member will have this size -counter = count() # next(counter) will return 0, 1, 2, ... - -parser = argparse.ArgumentParser(description="Create a file with EDM4hep data") -parser.add_argument( - "--rntuple", action="store_true", help="Use a ROOT ntuple instead of EDM4hep" -) -parser.add_argument( - "--output-file", type=str, help="Output file name", default="edm4hep.root" -) -args = parser.parse_args() -output_file = args.output_file - -if args.rntuple: - try: - writer = podio.root_io.RNTupleWriter(output_file) - except AttributeError: - print("The RNTuple writer from podio is not available but was requested") - sys.exit(1) -else: - writer = podio.root_io.Writer(output_file) - -for i in range(frames): - print(f"Writing frame {i}") - frame = podio.Frame() +import podio +import edm4hep - # Make covariance matrices needed later +FRAMES = 3 # How many frames or events will be written +VECTORSIZE = 5 # For vector members, each vector member will have this size + +COUNT_START = 42 # Where to the counter from - # With the current version of cppyy (from ROOT 6.30.06) - # it's not possible to initialize an std::array from a list - # In future versions (works with 6.32.02) it will be possible - cov3f = edm4hep.CovMatrix3f() - for i in range(6): - cov3f[i] = next(counter) - cov4f = edm4hep.CovMatrix4f() - for i in range(10): - cov4f[i] = next(counter) - cov6f = edm4hep.CovMatrix6f() - for i in range(21): - cov6f[i] = next(counter) +def create_EventHeaderCollection(vectorsize): + """Create an EventHeaderCollection""" + counter = count(COUNT_START) header = edm4hep.EventHeaderCollection() h = header.create() h.setEventNumber(next(counter)) @@ -59,11 +28,17 @@ h.setWeight(1.0) for j in range(vectorsize): h.addToWeights(next(counter)) - frame.put(header, "EventHeader") + return header + +def create_MCParticleCollection(): + """Create an MCParticleCollection""" + counter = count(COUNT_START) particles = edm4hep.MCParticleCollection() - for i in range(3): + p_list = [] + for _ in range(3): particle = particles.create() + p_list.append(particle) particle.setPDG(next(counter)) particle.setGeneratorStatus(next(counter)) particle.setSimulatorStatus(next(counter)) @@ -85,11 +60,14 @@ particle.setSpin(edm4hep.Vector3f(next(counter), next(counter), next(counter))) particle.setColorFlow(edm4hep.Vector2i(next(counter), next(counter))) - particles[0].addToDaughters(particles[1]) - particles[0].addToParents(particles[2]) - particle = particles[0] - frame.put(particles, "MCParticleCollection") + p_list[0].addToDaughters(p_list[1]) + p_list[0].addToParents(p_list[2]) + return particles + +def create_SimTrackerHitCollection(particle): + """Create a SimTrackerHitCollection""" + counter = count(COUNT_START) hits = edm4hep.SimTrackerHitCollection() hit = hits.create() hit.setCellID(next(counter)) @@ -100,9 +78,12 @@ hit.setPosition(edm4hep.Vector3d(next(counter), next(counter), next(counter))) hit.setMomentum(edm4hep.Vector3f(next(counter), next(counter), next(counter))) hit.setParticle(particle) - simtracker_hit = hit - frame.put(hits, "SimTrackerHitCollection") + return hits + +def create_CaloHitContributionCollection(particle): + """Create a CaloHitContributionCollection""" + counter = count(COUNT_START) hits = edm4hep.CaloHitContributionCollection() hit = hits.create() hit.setPDG(next(counter)) @@ -110,25 +91,35 @@ hit.setTime(next(counter)) hit.setStepPosition(edm4hep.Vector3f(next(counter), next(counter), next(counter))) hit.setParticle(particle) - calohit = hit - frame.put(hits, "CaloHitContributionCollection") + return hits + +def create_SimCalorimeterHitCollection(calo_contrib): + """Create a SimCalorimeterHitCollection""" + counter = count(COUNT_START) hits = edm4hep.SimCalorimeterHitCollection() hit = hits.create() hit.setCellID(next(counter)) hit.setEnergy(next(counter)) hit.setPosition(edm4hep.Vector3f(next(counter), next(counter), next(counter))) - hit.addToContributions(calohit) - simcalo_hit = hit - frame.put(hits, "SimCalorimeterHitCollection") + hit.addToContributions(calo_contrib) + return hits + +def create_RawCalorimeterHitCollection(): + """Crate a RawCalorimeterHitCollection""" + counter = count(COUNT_START) hits = edm4hep.RawCalorimeterHitCollection() hit = hits.create() hit.setCellID(next(counter)) hit.setAmplitude(next(counter)) hit.setTimeStamp(next(counter)) - frame.put(hits, "RawCalorimeterHitCollection") + return hits + +def create_CalorimeterHitCollection(): + """Create a CalorimeterHitCollection""" + counter = count(COUNT_START) hits = edm4hep.CalorimeterHitCollection() hit = hits.create() hit.setCellID(next(counter)) @@ -137,9 +128,28 @@ hit.setTime(next(counter)) hit.setPosition(edm4hep.Vector3f(next(counter), next(counter), next(counter))) hit.setType(next(counter)) - calo_hit = hit - frame.put(hits, "CalorimeterHitCollection") + return hits + + +def create_CovMatrixNf(n_dim): + """Create a covariance matrix of dimension n_dim""" + if n_dim not in (2, 3, 4, 6): + return ValueError( + f"{n_dim} is not a valid dimension for a CovMatrix in EDM4hep. Valid: (2, 3, 4, 6)" + ) + counter = count(COUNT_START) + # With the current version of cppyy (from ROOT 6.30.06) + # it's not possible to initialize an std::array from a list + # In future versions (works with 6.32.02) it will be possible + covNf = getattr(edm4hep, f"CovMatrix{n_dim}f")() + for i in range(n_dim * (n_dim + 1) // 2): + covNf[i] = next(counter) + return covNf + +def create_ParticleIDCollection(vectorsize): + """Create a ParticleIDCollection""" + counter = count(COUNT_START) pids = edm4hep.ParticleIDCollection() pid = pids.create() pid.setType(next(counter)) @@ -148,15 +158,20 @@ pid.setLikelihood(next(counter)) for j in range(vectorsize): pid.addToParameters(next(counter)) - frame.put(pids, "ParticleIDCollection") + return pids, pid + + +def create_ClusterCollection(vectorsize, calo_hit): + """Create a ClusterCollection""" + counter = count(COUNT_START) clusters = edm4hep.ClusterCollection() cluster = clusters.create() cluster.setType(next(counter)) cluster.setEnergy(next(counter)) cluster.setEnergyError(next(counter)) cluster.setPosition(edm4hep.Vector3f(next(counter), next(counter), next(counter))) - cluster.setPositionError(cov3f) + cluster.setPositionError(create_CovMatrixNf(3)) cluster.setITheta(next(counter)) cluster.setPhi(next(counter)) cluster.setDirectionError( @@ -167,8 +182,13 @@ cluster.addToSubdetectorEnergies(next(counter)) cluster.addToClusters(cluster) cluster.addToHits(calo_hit) - frame.put(clusters, "ClusterCollection") + return clusters + + +def create_TrackerHit3DCollection(): + """Create a TrackerHit3DCollection""" + counter = count(COUNT_START) hits = edm4hep.TrackerHit3DCollection() hit = hits.create() hit.setCellID(next(counter)) @@ -178,10 +198,13 @@ hit.setEDep(next(counter)) hit.setEDepError(next(counter)) hit.setPosition(edm4hep.Vector3d(next(counter), next(counter), next(counter))) - hit.setCovMatrix(cov3f) - tracker_hit = hit - frame.put(hits, "TrackerHit3DCollection") + hit.setCovMatrix(create_CovMatrixNf(3)) + return hits + +def create_TrackerHitPlaneCollection(): + """Create a TrackerHitPlaneCollection""" + counter = count(COUNT_START) hits = edm4hep.TrackerHitPlaneCollection() hit = hits.create() hit.setCellID(next(counter)) @@ -195,9 +218,13 @@ hit.setDu(next(counter)) hit.setDv(next(counter)) hit.setPosition(edm4hep.Vector3d(next(counter), next(counter), next(counter))) - hit.setCovMatrix(cov3f) - frame.put(hits, "TrackerHitPlaneCollection") + hit.setCovMatrix(create_CovMatrixNf(3)) + return hits + +def create_RawTimeSeriesCollection(vectorsize): + """Create a RawTimeSeriesCollection""" + counter = count(COUNT_START) series = edm4hep.RawTimeSeriesCollection() serie = series.create() serie.setCellID(next(counter)) @@ -207,8 +234,12 @@ serie.setInterval(next(counter)) for j in range(vectorsize): serie.addToAdcCounts(next(counter)) - frame.put(series, "RawTimeSeriesCollection") + return series + +def create_TrackCollection(vectorsize, tracker_hit): + """Create a TrackCollection""" + counter = count(COUNT_START) tracks = edm4hep.TrackCollection() track = tracks.create() track.setType(next(counter)) @@ -216,6 +247,7 @@ track.setNdf(next(counter)) for j in range(vectorsize): track.addToSubdetectorHitNumbers(next(counter)) + track.addToSubdetectorHoleNumbers(next(counter)) state = edm4hep.TrackState() state.location = next(counter) state.D0 = next(counter) @@ -227,25 +259,34 @@ state.referencePoint = edm4hep.Vector3f( next(counter), next(counter), next(counter) ) - state.covMatrix = cov6f + state.covMatrix = create_CovMatrixNf(6) track.addToTrackStates(state) + track.addToTrackerHits(tracker_hit) track.addToTracks(track) track.setNholes(next(counter)) - frame.put(tracks, "TrackCollection") + return tracks + +def create_VertexCollection(vectorsize): + """Create a VertexCollection""" + counter = count(COUNT_START) vertex = edm4hep.VertexCollection() v = vertex.create() v.setType(next(counter)) v.setChi2(next(counter)) v.setNdf(next(counter)) v.setPosition(edm4hep.Vector3f(next(counter), next(counter), next(counter))) - v.setCovMatrix(cov3f) + v.setCovMatrix(create_CovMatrixNf(3)) v.setAlgorithmType(next(counter)) for j in range(vectorsize): v.addToParameters(next(counter)) - frame.put(vertex, "VertexCollection") + return vertex, v + +def create_ReconstructedParticleCollection(vertex, cluster, track): + """Create a ReconstructedParticleCollection""" + counter = count(COUNT_START) rparticles = edm4hep.ReconstructedParticleCollection() rparticle = rparticles.create() rparticle.setPDG(next(counter)) @@ -257,70 +298,28 @@ rparticle.setCharge(next(counter)) rparticle.setMass(next(counter)) rparticle.setGoodnessOfPID(next(counter)) - rparticle.setCovMatrix(cov4f) - rparticle.setDecayVertex(v) + rparticle.setCovMatrix(create_CovMatrixNf(4)) + rparticle.setDecayVertex(vertex) rparticle.addToClusters(cluster) rparticle.addToTracks(track) rparticle.addToParticles(rparticle) - reco_particle = rparticle - frame.put(rparticles, "ReconstructedParticleCollection") + return rparticles - v.addToParticles(reco_particle) - pid.setParticle(reco_particle) - # TODO: Add when the PID PR is merged - # if not args.use_pre1: - # pid.setParticle(reco_particle) - - links = edm4hep.RecoMCParticleLinkCollection() - link = links.create() - link.setWeight(next(counter)) - link.setFrom(reco_particle) - link.setTo(particle) - frame.put(links, "RecoMCParticleLinkCollection") - - links = edm4hep.CaloHitSimCaloHitLinkCollection() - link = links.create() - link.setWeight(next(counter)) - link.setFrom(calo_hit) - link.setTo(simcalo_hit) - frame.put(links, "CaloHitSimCaloHitLinkCollection") - - links = edm4hep.TrackerHitSimTrackerHitLinkCollection() - link = links.create() - link.setWeight(next(counter)) - link.setFrom(tracker_hit) - link.setTo(simtracker_hit) - frame.put(links, "TrackerHitSimTrackerHitLinkCollection") - - links = edm4hep.CaloHitMCParticleLinkCollection() +def create_LinkCollection(coll_type, from_el, to_el): + """Create a LinkCollection of the given type and add one link to it""" + counter = count(COUNT_START) + links = coll_type() link = links.create() link.setWeight(next(counter)) - link.setFrom(calo_hit) - link.setTo(particle) - frame.put(links, "CaloHitMCParticleLinkCollection") + link.setFrom(from_el) + link.setTo(to_el) + return links - links = edm4hep.ClusterMCParticleLinkCollection() - link = links.create() - link.setWeight(next(counter)) - link.setFrom(cluster) - link.setTo(particle) - frame.put(links, "ClusterMCParticleLinkCollection") - - links = edm4hep.TrackMCParticleLinkCollection() - link = links.create() - link.setWeight(next(counter)) - link.setFrom(track) - link.setTo(particle) - frame.put(links, "TrackMCParticleLinkCollection") - - links = edm4hep.VertexRecoParticleLinkCollection() - link = links.create() - link.setWeight(next(counter)) - link.setTo(reco_particle) - link.setFrom(v) - frame.put(links, "MCVertexRecoParticleLinkCollection") +def create_TimeSeriesCollection(vectorsize): + """Create a TimeSeriesCollection""" + counter = count(COUNT_START) timeseries = edm4hep.TimeSeriesCollection() serie = timeseries.create() serie.setCellID(next(counter)) @@ -328,8 +327,12 @@ serie.setInterval(next(counter)) for j in range(vectorsize): serie.addToAmplitude(next(counter)) - frame.put(timeseries, "TimeSeriesCollection") + return timeseries + +def create_RecDqdxCollection(track): + """Create a RecDqdxCollection""" + counter = count(COUNT_START) recdqdx = edm4hep.RecDqdxCollection() dqdx = recdqdx.create() q = edm4hep.Quantity() @@ -338,8 +341,12 @@ q.error = next(counter) dqdx.setDQdx(q) dqdx.setTrack(track) - frame.put(recdqdx, "RecDqdxCollection") + return recdqdx + +def create_GeneratorEventParametersCollection(vectorsize, particle): + """Create a GeneratorEventParametersCollection""" + counter = count(COUNT_START) gep_coll = edm4hep.GeneratorEventParametersCollection() gep = gep_coll.create() gep.setEventScale(next(counter)) @@ -347,13 +354,16 @@ gep.setAlphaQCD(next(counter)) gep.setSignalProcessId(next(counter)) gep.setSqrts(next(counter)) - frame.put(gep_coll, "GeneratorEventParametersCollection") - for i in range(vectorsize): gep.addToCrossSections(next(counter)) gep.addToCrossSectionErrors(next(counter)) gep.addToSignalVertex(particle) + return gep_coll + +def create_GeneratorPdfInfoCollection(): + """Create a GeneratorPdfInfoCollection""" + counter = count(COUNT_START) gpi_coll = edm4hep.GeneratorPdfInfoCollection() gpi = gpi_coll.create() # Doesn't work with ROOT 6.30.06 @@ -373,6 +383,121 @@ gpi.setXf(0, next(counter)) gpi.setXf(1, next(counter)) gpi.setScale(next(counter)) - frame.put(gpi_coll, "GeneratorPdfInfoCollection") + return gpi_coll + + +def put_link_collection(frame, link_name, from_el, to_el): + """Helper function to put a link collection into the frame""" + coll_name = f"{link_name}Collection" + coll_type = getattr(edm4hep, coll_name) + frame.put(create_LinkCollection(coll_type, from_el, to_el), coll_name) + + +def create_frame(): + """Create a Frame with all types""" + frame = podio.Frame() + + frame.put(create_EventHeaderCollection(VECTORSIZE), "EventHeader") + + particles = frame.put(create_MCParticleCollection(), "MCParticleCollection") + particle = particles[0] + + hits = frame.put( + create_SimTrackerHitCollection(particle), "SimTrackerHitCollection" + ) + simtracker_hit = hits[0] + + calo_contribs = frame.put( + create_CaloHitContributionCollection(particle), "CaloHitContributionCollection" + ) + calo_contrib = calo_contribs[0] + + hits = frame.put( + create_SimCalorimeterHitCollection(calo_contrib), "SimCalorimeterHitCollection" + ) + simcalo_hit = hits[0] + + frame.put(create_RawCalorimeterHitCollection(), "RawCalorimeterHitCollection") - writer.write_frame(frame, "events") + hits = frame.put(create_CalorimeterHitCollection(), "CalorimeterHitCollection") + calo_hit = hits[0] + + clusters = frame.put( + create_ClusterCollection(VECTORSIZE, calo_hit), "ClusterCollection" + ) + cluster = clusters[0] + + hits = frame.put(create_TrackerHit3DCollection(), "TrackerHit3DCollection") + tracker_hit = hits[0] + frame.put(create_TrackerHitPlaneCollection(), "TrackerHitPlaneCollection") + + frame.put(create_RawTimeSeriesCollection(VECTORSIZE), "RawTimeSeriesCollection") + + tracks = frame.put( + create_TrackCollection(VECTORSIZE, tracker_hit), "TrackCollection" + ) + track = tracks[0] + + pids, pid = create_ParticleIDCollection(VECTORSIZE) + + vertices, vertex = create_VertexCollection(VECTORSIZE) + + reco_particles = frame.put( + create_ReconstructedParticleCollection(vertex, cluster, track), + "ReconstructedParticleCollection", + ) + reco_particle = reco_particles[0] + + vertex.addToParticles(reco_particle) + frame.put(vertices, "VertexCollection") + + pid.setParticle(reco_particle) + frame.put(pids, "ParticleIDCollection") + + put_link_collection(frame, "RecoMCParticleLink", reco_particle, particle) + put_link_collection(frame, "CaloHitSimCaloHitLink", calo_hit, simcalo_hit) + put_link_collection( + frame, "TrackerHitSimTrackerHitLink", tracker_hit, simtracker_hit + ) + put_link_collection(frame, "CaloHitMCParticleLink", calo_hit, particle) + put_link_collection(frame, "ClusterMCParticleLink", cluster, particle) + put_link_collection(frame, "TrackMCParticleLink", track, particle) + put_link_collection(frame, "VertexRecoParticleLink", vertex, reco_particle) + + frame.put(create_TimeSeriesCollection(VECTORSIZE), "TimeSeriesCollection") + frame.put(create_RecDqdxCollection(track), "RecDqdxCollection") + + frame.put( + create_GeneratorEventParametersCollection(VECTORSIZE, particle), + "GeneratorEventParametersCollection", + ) + frame.put(create_GeneratorPdfInfoCollection(), "GeneratorPdfInfoCollection") + + return frame + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Create a file with EDM4hep data") + parser.add_argument( + "--rntuple", action="store_true", help="Use a ROOT ntuple instead of EDM4hep" + ) + parser.add_argument( + "--output-file", type=str, help="Output file name", default="edm4hep.root" + ) + args = parser.parse_args() + output_file = args.output_file + + if args.rntuple: + try: + writer = podio.root_io.RNTupleWriter(output_file) + except AttributeError: + print("The RNTuple writer from podio is not available but was requested") + sys.exit(1) + else: + writer = podio.root_io.Writer(output_file) + + for i in range(FRAMES): + frame = create_frame() + print(f"Writing frame {i}") + + writer.write_frame(frame, "events") diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 65c816322..264780c36 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -13,15 +13,39 @@ function(set_test_env _testname) ) endfunction() -add_test(NAME "Create an EDM4hep data file" COMMAND python ${PROJECT_SOURCE_DIR}/scripts/createEDM4hepFile.py --output-file edm4hep_example.root) -set_test_env("Create an EDM4hep data file") +add_test(NAME create_complete_file COMMAND python ${PROJECT_SOURCE_DIR}/scripts/createEDM4hepFile.py --output-file edm4hep_example.root) +set_test_env(create_complete_file) -add_test(NAME "Create an EDM4hep data file (RNTuple)" COMMAND python ${PROJECT_SOURCE_DIR}/scripts/createEDM4hepFile.py --rntuple --output-file edm4hep_example_rntuple.root) -set_test_env("Create an EDM4hep data file (RNTuple)") -set_tests_properties("Create an EDM4hep data file (RNTuple)" PROPERTIES +add_test(NAME create_complete_file_rntuple COMMAND python ${PROJECT_SOURCE_DIR}/scripts/createEDM4hepFile.py --rntuple --output-file edm4hep_example_rntuple.root) +set_test_env(create_complete_file_rntuple) +set_tests_properties(create_complete_file_rntuple PROPERTIES SKIP_REGULAR_EXPRESSION "The RNTuple writer from podio is not available but was requested" ) +add_test(NAME check_complete_file COMMAND pytest --inputfile=${CMAKE_BINARY_DIR}/test/edm4hep_example.root -v) +set_test_env(check_complete_file) +set_tests_properties( + check_complete_file + PROPERTIES + DEPENDS create_complete_file +) +add_test(NAME check_complete_file_rntuple COMMAND pytest --inputfile=${CMAKE_BINARY_DIR}/test/edm4hep_example_rntuple.root -v) +set_test_env(check_complete_file_rntuple) +set_tests_properties( + check_complete_file_rntuple + PROPERTIES + DEPENDS create_complete_file_rntuple +) + +set_tests_properties( + check_complete_file + check_complete_file_rntuple + + PROPERTIES + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} + FAIL_REGULAR_EXPRESSION "" +) + add_executable(write_events write_events.cc) target_include_directories(write_events PUBLIC ${PROJECT_SOURCE_DIR}/edm4hep ) target_link_libraries(write_events edm4hep podio::podioRootIO) diff --git a/test/conftest.py b/test/conftest.py new file mode 100644 index 000000000..eacd06cd0 --- /dev/null +++ b/test/conftest.py @@ -0,0 +1,8 @@ +#!/usr/bin/env python3 +"""pytest config module to make passing of input file name possible""" + + +def pytest_addoption(parser): + """Hook to add an inputfile argument to pytest for checking EDM4hep file + contents""" + parser.addoption("--inputfile", action="store") diff --git a/test/pytest.ini b/test/pytest.ini new file mode 100644 index 000000000..10cf4e47f --- /dev/null +++ b/test/pytest.ini @@ -0,0 +1,2 @@ +[pytest] +addopts = --ignore=test_rdf.py --ignore=tools --ignore=utils diff --git a/test/test_EDM4hepFile.py b/test/test_EDM4hepFile.py new file mode 100644 index 000000000..fc8da2d14 --- /dev/null +++ b/test/test_EDM4hepFile.py @@ -0,0 +1,508 @@ +#!/usr/bin/env python3 +"""Pytest module that can be used to check whether the file that has been +created by scripts/createEDM4hepFile.py has the expected contents +""" + +import podio +import edm4hep +import pytest +from itertools import count + +# For now simply copy these from createEDM4hepFile.py +FRAMES = 3 +VECTORSIZE = 5 +COUNT_START = 42 # Starting point for the counters + + +@pytest.fixture(scope="module") +def inputfile_name(pytestconfig): + """Get the inputfile name that has been passed via commandline""" + return pytestconfig.getoption("inputfile") + + +@pytest.fixture(scope="module") +def reader(inputfile_name): + """Get the reader for the passed filename""" + return podio.reading.get_reader(inputfile_name) + + +@pytest.fixture(scope="module") +def events(reader): + """Get the events from the reader""" + return reader.get("events") + + +@pytest.fixture(scope="module") +def event(events): + """Get the first event from the events. Since all events have the same + contents, we simply assume that everything is alright if we simply check the + first one + """ + return events[0] + + +@pytest.fixture(scope="module") +def particle(event): + """Get the particle that is used in some relations""" + return event.get("MCParticleCollection")[0] + + +@pytest.fixture(scope="module") +def reco_particle(event): + """Get the reconstructed particle that is used in some relations""" + return event.get("ReconstructedParticleCollection")[0] + + +@pytest.fixture(scope="module") +def track(event): + """Get the track that is used in some relations""" + return event.get("TrackCollection")[0] + + +def check_cov_matrix(cov_matrix, n_dim): + """Check the contents of the passed covariance matrix""" + counter = count(COUNT_START) + # for val in cov_matrix: # Doesn't work as expected with root + for i in range(n_dim * (n_dim + 1) // 2): + assert cov_matrix[i] == next(counter) + + +def test_basic_file_contents(events): + """Make sure the basic file contents are OK""" + assert len(events) == FRAMES + + +def test_EventHeaderCollection(event): + """Check an EventHeaderCollection""" + counter = count(COUNT_START) + header = event.get("EventHeader") + assert len(header) == 1 + h = header[0] + assert h.getEventNumber() == next(counter) + assert h.getRunNumber() == next(counter) + assert h.getTimeStamp() == next(counter) + assert h.getWeight() == 1.0 + assert len(h.getWeights()) == VECTORSIZE + for weight in h.getWeights(): + assert weight == next(counter) + + +def test_MCParticleCollection(event): + """Check the MCParticleCollection""" + counter = count(COUNT_START) + particles = event.get("MCParticleCollection") + assert len(particles) == 3 + for i in range(3): + particle = particles[i] + assert particle.getPDG() == next(counter) + assert particle.getGeneratorStatus() == next(counter) + assert particle.getSimulatorStatus() == next(counter) + assert particle.getCharge() == next(counter) + assert particle.getTime() == next(counter) + assert particle.getMass() == next(counter) + assert particle.getVertex() == edm4hep.Vector3d( + next(counter), next(counter), next(counter) + ) + assert particle.getEndpoint() == edm4hep.Vector3d( + next(counter), next(counter), next(counter) + ) + assert particle.getMomentum() == edm4hep.Vector3d( + next(counter), next(counter), next(counter) + ) + assert particle.getMomentumAtEndpoint() == edm4hep.Vector3d( + next(counter), next(counter), next(counter) + ) + assert particle.getSpin() == edm4hep.Vector3f( + next(counter), next(counter), next(counter) + ) + assert particle.getColorFlow() == edm4hep.Vector2i(next(counter), next(counter)) + + assert particles[0].getDaughters()[0] == particles[1] + assert particles[0].getParents()[0] == particles[2] + + +def test_SimTrackerHitCollection(event, particle): + """Check the SimTrackerHitCollection""" + counter = count(COUNT_START) + hits = event.get("SimTrackerHitCollection") + assert len(hits) == 1 + hit = hits[0] + assert hit.getCellID() == next(counter) + assert hit.getEDep() == next(counter) + assert hit.getTime() == next(counter) + assert hit.getPathLength() == next(counter) + assert hit.getQuality() == next(counter) + assert hit.getPosition() == edm4hep.Vector3d( + next(counter), next(counter), next(counter) + ) + assert hit.getMomentum() == edm4hep.Vector3f( + next(counter), next(counter), next(counter) + ) + + assert hit.getParticle() == particle + + +def test_CaloHitContributionCollection(event, particle): + """Check the CaloHitContributionCollection""" + counter = count(COUNT_START) + hits = event.get("CaloHitContributionCollection") + assert len(hits) == 1 + hit = hits[0] + assert hit.getPDG() == next(counter) + assert hit.getEnergy() == next(counter) + assert hit.getTime() == next(counter) + assert hit.getStepPosition() == edm4hep.Vector3f( + next(counter), next(counter), next(counter) + ) + + assert hit.getParticle() == particle + + +def test_SimCalorimeterHitCollection(event): + """Check the SimCalorimeterHitCollection""" + counter = count(COUNT_START) + hits = event.get("SimCalorimeterHitCollection") + assert len(hits) == 1 + hit = hits[0] + assert hit.getCellID() == next(counter) + assert hit.getEnergy() == next(counter) + assert hit.getPosition() == edm4hep.Vector3f( + next(counter), next(counter), next(counter) + ) + + calo_contrib = event.get("CaloHitContributionCollection")[0] + assert len(hit.getContributions()) == 1 + assert hit.getContributions()[0] == calo_contrib + + +def test_RawCalorimeterHitCollection(event): + """Check the RawCalorimeterHitCollection""" + counter = count(COUNT_START) + hits = event.get("RawCalorimeterHitCollection") + assert len(hits) == 1 + hit = hits[0] + assert hit.getCellID() == next(counter) + assert hit.getAmplitude() == next(counter) + assert hit.getTimeStamp() == next(counter) + + +def test_CalorimeterHitCollection(event): + """Check the CalorimeterHitCollection""" + counter = count(COUNT_START) + hits = event.get("CalorimeterHitCollection") + assert len(hits) == 1 + hit = hits[0] + assert hit.getCellID() == next(counter) + assert hit.getEnergy() == next(counter) + assert hit.getEnergyError() == next(counter) + assert hit.getTime() == next(counter) + assert hit.getPosition() == edm4hep.Vector3f( + next(counter), next(counter), next(counter) + ) + assert hit.getType() == next(counter) + + +def test_ParticleIDCollection(event, reco_particle): + """Check the ParticleIDCollection""" + counter = count(COUNT_START) + pids = event.get("ParticleIDCollection") + assert len(pids) == 1 + pid = pids[0] + assert pid.getType() == next(counter) + assert pid.getPDG() == next(counter) + assert pid.getAlgorithmType() == next(counter) + assert pid.getLikelihood() == next(counter) + assert len(pid.getParameters()) == VECTORSIZE + for param in pid.getParameters(): + assert param == next(counter) + + assert pid.getParticle() == reco_particle + + +def test_ClusterCollection(event): + """Check the ClusterCollection""" + counter = count(COUNT_START) + clusters = event.get("ClusterCollection") + assert len(clusters) == 1 + cluster = clusters[0] + assert cluster.getType() == next(counter) + assert cluster.getEnergy() == next(counter) + assert cluster.getEnergyError() == next(counter) + assert cluster.getPosition() == edm4hep.Vector3f( + next(counter), next(counter), next(counter) + ) + check_cov_matrix(cluster.getPositionError(), 3) + assert cluster.getITheta() == next(counter) + assert cluster.getPhi() == next(counter) + assert cluster.getDirectionError() == edm4hep.Vector3f( + next(counter), next(counter), next(counter) + ) + + assert len(cluster.getShapeParameters()) == VECTORSIZE + assert len(cluster.getSubdetectorEnergies()) == VECTORSIZE + for j in range(VECTORSIZE): + cluster.getShapeParameters()[j] == next(counter) + cluster.getSubdetectorEnergies()[j] == next(counter) + + assert len(cluster.getClusters()) == 1 + assert cluster.getClusters()[0] == cluster + + calo_hit = event.get("CalorimeterHitCollection")[0] + assert len(cluster.getHits()) == 1 + assert cluster.getHits()[0] == calo_hit + cluster.addToHits(calo_hit) + + +def test_TrackerHit3DCollection(event): + """Check the TrackerHit3DCollection""" + counter = count(COUNT_START) + hits = event.get("TrackerHit3DCollection") + assert len(hits) == 1 + hit = hits[0] + assert hit.getCellID() == next(counter) + assert hit.getType() == next(counter) + assert hit.getQuality() == next(counter) + assert hit.getTime() == next(counter) + assert hit.getEDep() == next(counter) + assert hit.getEDepError() == next(counter) + assert hit.getPosition() == edm4hep.Vector3d( + next(counter), next(counter), next(counter) + ) + check_cov_matrix(hit.getCovMatrix(), 3) + + +def test_TrackerHitPlaneCollection(event): + """Check the TrackerHitPlaneCollection""" + counter = count(COUNT_START) + hits = event.get("TrackerHitPlaneCollection") + assert len(hits) == 1 + hit = hits[0] + assert hit.getCellID() == next(counter) + assert hit.getType() == next(counter) + assert hit.getQuality() == next(counter) + assert hit.getTime() == next(counter) + assert hit.getEDep() == next(counter) + assert hit.getEDepError() == next(counter) + assert hit.getU() == edm4hep.Vector2f(next(counter), next(counter)) + assert hit.getV() == edm4hep.Vector2f(next(counter), next(counter)) + assert hit.getDu() == next(counter) + assert hit.getDv() == next(counter) + assert hit.getPosition() == edm4hep.Vector3d( + next(counter), next(counter), next(counter) + ) + check_cov_matrix(hit.getCovMatrix(), 3) + + +def test_RawTimeSeriesCollection(event): + """Check a RawTimeSeriesCollection""" + counter = count(COUNT_START) + series = event.get("RawTimeSeriesCollection") + assert len(series) == 1 + serie = series[0] + assert serie.getCellID() == next(counter) + assert serie.getQuality() == next(counter) + assert serie.getTime() == next(counter) + assert serie.getCharge() == next(counter) + assert serie.getInterval() == next(counter) + assert len(serie.getAdcCounts()) == VECTORSIZE + for val in serie.getAdcCounts(): + assert val == next(counter) + + +def test_TrackCollection(event): + """Check the TrackCollection""" + counter = count(COUNT_START) + tracks = event.get("TrackCollection") + assert len(tracks) == 1 + track = tracks[0] + assert track.getType() == next(counter) + assert track.getChi2() == next(counter) + assert track.getNdf() == next(counter) + + subdetHitNumbers = track.getSubdetectorHitNumbers() + assert len(subdetHitNumbers) == VECTORSIZE + subdetHoleNumbers = track.getSubdetectorHoleNumbers() + assert len(subdetHoleNumbers) == VECTORSIZE + trackStates = track.getTrackStates() + assert len(trackStates) == VECTORSIZE + + for j in range(VECTORSIZE): + subdetHitNumbers[j] == next(counter) + subdetHoleNumbers[j] == next(counter) + + state = trackStates[j] + assert state.location == next(counter) + assert state.D0 == next(counter) + assert state.phi == next(counter) + assert state.omega == next(counter) + assert state.Z0 == next(counter) + assert state.tanLambda == next(counter) + assert state.time == next(counter) + assert state.referencePoint == edm4hep.Vector3f( + next(counter), next(counter), next(counter) + ) + check_cov_matrix(state.covMatrix, 6) + + assert track.getNholes() == next(counter) + + tracker_hit = event.get("TrackerHit3DCollection")[0] + assert len(track.getTrackerHits()) == 1 + assert track.getTrackerHits()[0] == tracker_hit + + assert len(track.getTracks()) == 1 + assert track.getTracks()[0] == track + + +def test_VertexCollection(event, reco_particle): + """Check the VertexCollection""" + counter = count(COUNT_START) + vertex = event.get("VertexCollection") + assert len(vertex) == 1 + v = vertex[0] + assert v.getType() == next(counter) + assert v.getChi2() == next(counter) + assert v.getNdf() == next(counter) + assert v.getPosition() == edm4hep.Vector3f( + next(counter), next(counter), next(counter) + ) + check_cov_matrix(v.getCovMatrix(), 3) + assert v.getAlgorithmType() == next(counter) + assert len(v.getParameters()) == VECTORSIZE + for val in v.getParameters(): + assert val == next(counter) + + assert len(v.getParticles()) == 1 + assert v.getParticles()[0] == reco_particle + + +def test_ReconstructedParticleCollection(event, track): + """Check the ReconstructedParticleCollection""" + counter = count(COUNT_START) + rparticles = event.get("ReconstructedParticleCollection") + assert len(rparticles) == 1 + rparticle = rparticles[0] + assert rparticle.getPDG() == next(counter) + assert rparticle.getEnergy() == next(counter) + assert rparticle.getMomentum() == edm4hep.Vector3f( + next(counter), next(counter), next(counter) + ) + assert rparticle.getReferencePoint() == edm4hep.Vector3f( + next(counter), next(counter), next(counter) + ) + assert rparticle.getCharge() == next(counter) + assert rparticle.getMass() == next(counter) + assert rparticle.getGoodnessOfPID() == next(counter) + check_cov_matrix(rparticle.getCovMatrix(), 4) + + vertex = event.get("VertexCollection")[0] + assert rparticle.getDecayVertex() == vertex + + assert len(rparticle.getClusters()) == 1 + cluster = event.get("ClusterCollection")[0] + assert rparticle.getClusters()[0] == cluster + + assert len(rparticle.getTracks()) == 1 + assert rparticle.getTracks()[0] == track + + assert len(rparticle.getParticles()) == 1 + assert rparticle.getParticles()[0] == rparticle + + +def test_TimeSeriesCollection(event): + """Check the TimeSeriesCollection""" + counter = count(COUNT_START) + timeseries = event.get("TimeSeriesCollection") + assert len(timeseries) == 1 + serie = timeseries[0] + assert serie.getCellID() == next(counter) + assert serie.getTime() == next(counter) + assert serie.getInterval() == next(counter) + assert len(serie.getAmplitude()) == VECTORSIZE + for val in serie.getAmplitude(): + assert val == next(counter) + + +def test_RecDqdxCollection(event, track): + """Check the RecDqdxCollection""" + counter = count(COUNT_START) + recdqdx = event.get("RecDqdxCollection") + assert len(recdqdx) == 1 + dqdx = recdqdx[0] + q = dqdx.getDQdx() + assert q.type == next(counter) + assert q.value == next(counter) + assert q.error == next(counter) + assert dqdx.getTrack() == track + + +def test_GeneratorEventParametersCollection(event, particle): + """Check the GeneratorEventParametersCollection""" + counter = count(COUNT_START) + gep_coll = event.get("GeneratorEventParametersCollection") + assert len(gep_coll) == 1 + gep = gep_coll[0] + assert gep.getEventScale() == next(counter) + assert gep.getAlphaQED() == next(counter) + assert gep.getAlphaQCD() == next(counter) + assert gep.getSignalProcessId() == next(counter) + assert gep.getSqrts() == next(counter) + xsecs = gep.getCrossSections() + xsec_errors = gep.getCrossSectionErrors() + assert len(xsecs) == VECTORSIZE + assert len(xsec_errors) == VECTORSIZE + for i in range(VECTORSIZE): + assert xsecs[i] == next(counter) + assert xsec_errors[i] == next(counter) + + assert len(gep.getSignalVertex()) == 1 + assert gep.getSignalVertex()[0] == particle + + +def test_GeneratorPdfInfoCollection(event): + """Check the GeneratorPdfInfoCollection""" + counter = count(COUNT_START) + gpi_coll = event.get("GeneratorPdfInfoCollection") + assert len(gpi_coll) == 1 + gpi = gpi_coll[0] + assert gpi.getPartonId(0) == next(counter) + assert gpi.getPartonId(1) == next(counter) + assert gpi.getLhapdfId(0) == next(counter) + assert gpi.getLhapdfId(1) == next(counter) + assert gpi.getX(0) == next(counter) + assert gpi.getX(1) == next(counter) + assert gpi.getXf(0) == next(counter) + assert gpi.getXf(1) == next(counter) + + assert gpi.getScale() == next(counter) + + +def check_LinkCollection(event, coll_type, from_el, to_el): + """Check a single link collection of a given type""" + counter = count(COUNT_START) + coll_name = f"{coll_type}Collection" + link_coll = event.get(coll_name) + assert len(link_coll) == 1 + link = link_coll[0] + assert link.getWeight() == next(counter) + assert link.getFrom() == from_el + assert link.getTo() == to_el + + +def test_LinkCollections(event, particle, reco_particle, track): + """Check all the link collections""" + calo_hit = event.get("CalorimeterHitCollection")[0] + simcalo_hit = event.get("SimCalorimeterHitCollection")[0] + cluster = event.get("ClusterCollection")[0] + tracker_hit = event.get("TrackerHit3DCollection")[0] + simtracker_hit = event.get("SimTrackerHitCollection")[0] + vertex = event.get("VertexCollection")[0] + + check_LinkCollection(event, "RecoMCParticleLink", reco_particle, particle) + check_LinkCollection(event, "CaloHitSimCaloHitLink", calo_hit, simcalo_hit) + check_LinkCollection( + event, "TrackerHitSimTrackerHitLink", tracker_hit, simtracker_hit + ) + check_LinkCollection(event, "CaloHitMCParticleLink", calo_hit, particle) + check_LinkCollection(event, "ClusterMCParticleLink", cluster, particle) + check_LinkCollection(event, "TrackMCParticleLink", track, particle) + check_LinkCollection(event, "VertexRecoParticleLink", vertex, reco_particle)