diff --git a/scripts/createEDM4hepFile.py b/scripts/createEDM4hepFile.py new file mode 100644 index 000000000..f838adeed --- /dev/null +++ b/scripts/createEDM4hepFile.py @@ -0,0 +1,373 @@ +# 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 + +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, ... +# used to generate the dummy data +output_file = "output.root" + +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" +) +args = parser.parse_args() + +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() + + # Make covariance matrices needed later + + # 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) + + header = edm4hep.EventHeaderCollection() + h = header.create() + h.setEventNumber(next(counter)) + h.setRunNumber(next(counter)) + h.setTimeStamp(next(counter)) + h.setWeight(1.0) + for j in range(vectorsize): + h.addToWeights(next(counter)) + frame.put(header, "EventHeader") + + particles = edm4hep.MCParticleCollection() + for i in range(3): + particle = particles.create() + particle.setPDG(next(counter)) + particle.setGeneratorStatus(next(counter)) + particle.setSimulatorStatus(next(counter)) + particle.setCharge(next(counter)) + particle.setTime(next(counter)) + particle.setMass(next(counter)) + particle.setVertex( + edm4hep.Vector3d(next(counter), next(counter), next(counter)) + ) + particle.setEndpoint( + edm4hep.Vector3d(next(counter), next(counter), next(counter)) + ) + particle.setMomentum( + edm4hep.Vector3d(next(counter), next(counter), next(counter)) + ) + particle.setMomentumAtEndpoint( + edm4hep.Vector3d(next(counter), next(counter), next(counter)) + ) + 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") + + hits = edm4hep.SimTrackerHitCollection() + hit = hits.create() + hit.setCellID(next(counter)) + hit.setEDep(next(counter)) + hit.setTime(next(counter)) + hit.setPathLength(next(counter)) + hit.setQuality(next(counter)) + 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") + + hits = edm4hep.CaloHitContributionCollection() + hit = hits.create() + hit.setPDG(next(counter)) + hit.setEnergy(next(counter)) + hit.setTime(next(counter)) + hit.setStepPosition(edm4hep.Vector3f(next(counter), next(counter), next(counter))) + hit.setParticle(particle) + calohit = hit + frame.put(hits, "CaloHitContributionCollection") + + 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") + + hits = edm4hep.RawCalorimeterHitCollection() + hit = hits.create() + hit.setCellID(next(counter)) + hit.setAmplitude(next(counter)) + hit.setTimeStamp(next(counter)) + frame.put(hits, "RawCalorimeterHitCollection") + + hits = edm4hep.CalorimeterHitCollection() + hit = hits.create() + hit.setCellID(next(counter)) + hit.setEnergy(next(counter)) + hit.setEnergyError(next(counter)) + 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") + + pids = edm4hep.ParticleIDCollection() + pid = pids.create() + pid.setType(next(counter)) + pid.setPDG(next(counter)) + pid.setAlgorithmType(next(counter)) + pid.setLikelihood(next(counter)) + for j in range(vectorsize): + pid.addToParameters(next(counter)) + frame.put(pids, "ParticleIDCollection") + + 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.setITheta(next(counter)) + cluster.setPhi(next(counter)) + cluster.setDirectionError( + edm4hep.Vector3f(next(counter), next(counter), next(counter)) + ) + for j in range(vectorsize): + cluster.addToShapeParameters(next(counter)) + cluster.addToSubdetectorEnergies(next(counter)) + cluster.addToClusters(cluster) + cluster.addToHits(calo_hit) + frame.put(clusters, "ClusterCollection") + + hits = edm4hep.TrackerHit3DCollection() + hit = hits.create() + hit.setCellID(next(counter)) + hit.setType(next(counter)) + hit.setQuality(next(counter)) + hit.setTime(next(counter)) + 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") + + hits = edm4hep.TrackerHitPlaneCollection() + hit = hits.create() + hit.setCellID(next(counter)) + hit.setType(next(counter)) + hit.setQuality(next(counter)) + hit.setTime(next(counter)) + hit.setEDep(next(counter)) + hit.setEDepError(next(counter)) + hit.setU(edm4hep.Vector2f(next(counter), next(counter))) + hit.setV(edm4hep.Vector2f(next(counter), next(counter))) + 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") + + series = edm4hep.RawTimeSeriesCollection() + serie = series.create() + serie.setCellID(next(counter)) + serie.setQuality(next(counter)) + serie.setTime(next(counter)) + serie.setCharge(next(counter)) + serie.setInterval(next(counter)) + for j in range(vectorsize): + serie.addToAdcCounts(next(counter)) + frame.put(series, "RawTimeSeriesCollection") + + tracks = edm4hep.TrackCollection() + track = tracks.create() + track.setType(next(counter)) + track.setChi2(next(counter)) + track.setNdf(next(counter)) + for j in range(vectorsize): + track.addToSubdetectorHitNumbers(next(counter)) + state = edm4hep.TrackState() + state.location = next(counter) + state.d0 = next(counter) + state.phi = next(counter) + state.omega = next(counter) + state.z0 = next(counter) + state.tanLambda = next(counter) + state.time = next(counter) + state.referencePoint = edm4hep.Vector3f( + next(counter), next(counter), next(counter) + ) + state.CovMatrix = cov6f + track.addToTrackStates(state) + track.addToTrackerHits(tracker_hit) + track.addToTracks(track) + frame.put(tracks, "TrackCollection") + + 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.setAlgorithmType(next(counter)) + for j in range(vectorsize): + v.addToParameters(next(counter)) + frame.put(vertex, "VertexCollection") + + rparticles = edm4hep.ReconstructedParticleCollection() + rparticle = rparticles.create() + rparticle.setPDG(next(counter)) + rparticle.setEnergy(next(counter)) + rparticle.setMomentum(edm4hep.Vector3f(next(counter), next(counter), next(counter))) + rparticle.setReferencePoint( + edm4hep.Vector3f(next(counter), next(counter), next(counter)) + ) + rparticle.setCharge(next(counter)) + rparticle.setMass(next(counter)) + rparticle.setGoodnessOfPID(next(counter)) + rparticle.setCovMatrix(cov4f) + rparticle.setDecayVertex(v) + rparticle.addToClusters(cluster) + rparticle.addToTracks(track) + rparticle.addToParticles(rparticle) + reco_particle = rparticle + frame.put(rparticles, "ReconstructedParticleCollection") + + 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) + + assocs = edm4hep.MCRecoParticleAssociationCollection() + assoc = assocs.create() + assoc.setWeight(next(counter)) + assoc.setRec(reco_particle) + assoc.setSim(particle) + frame.put(assocs, "MCRecoParticleAssociationCollection") + + assocs = edm4hep.MCRecoCaloAssociationCollection() + assoc = assocs.create() + assoc.setWeight(next(counter)) + assoc.setRec(calo_hit) + assoc.setSim(simcalo_hit) + frame.put(assocs, "MCRecoCaloAssociationCollection") + + assocs = edm4hep.MCRecoTrackerAssociationCollection() + assoc = assocs.create() + assoc.setWeight(next(counter)) + assoc.setRec(tracker_hit) + assoc.setSim(simtracker_hit) + frame.put(assocs, "MCRecoTrackerAssociationCollection") + + assocs = edm4hep.MCRecoCaloParticleAssociationCollection() + assoc = assocs.create() + assoc.setWeight(next(counter)) + assoc.setRec(calo_hit) + assoc.setSim(particle) + frame.put(assocs, "MCRecoCaloParticleAssociationCollection") + + assocs = edm4hep.MCRecoClusterParticleAssociationCollection() + assoc = assocs.create() + assoc.setWeight(next(counter)) + assoc.setRec(cluster) + assoc.setSim(particle) + frame.put(assocs, "MCRecoClusterParticleAssociationCollection") + + assocs = edm4hep.MCRecoTrackParticleAssociationCollection() + assoc = assocs.create() + assoc.setWeight(next(counter)) + assoc.setRec(track) + assoc.setSim(particle) + frame.put(assocs, "MCRecoTrackParticleAssociationCollection") + + assocs = edm4hep.RecoParticleVertexAssociationCollection() + assoc = assocs.create() + assoc.setWeight(next(counter)) + assoc.setRec(reco_particle) + assoc.setVertex(v) + frame.put(assocs, "MCRecoParticleVertexAssociationCollection") + + timeseries = edm4hep.TimeSeriesCollection() + serie = timeseries.create() + serie.setCellID(next(counter)) + serie.setTime(next(counter)) + serie.setInterval(next(counter)) + for j in range(vectorsize): + serie.addToAmplitude(next(counter)) + frame.put(timeseries, "TimeSeriesCollection") + + recdqdx = edm4hep.RecDqdxCollection() + dqdx = recdqdx.create() + q = edm4hep.Quantity() + q.type = next(counter) + q.value = next(counter) + q.error = next(counter) + dqdx.setDQdx(q) + dqdx.setTrack(track) + frame.put(recdqdx, "RecDqdxCollection") + + gep_coll = edm4hep.GeneratorEventParametersCollection() + gep = gep_coll.create() + gep.setEventScale(next(counter)) + gep.setAlphaQED(next(counter)) + gep.setAlphaQCD(next(counter)) + gep.setSignalProcessId(next(counter)) + gep.setSqrts(next(counter)) + + for i in range(vectorsize): + gep.addToCrossSections(next(counter)) + gep.addToCrossSectionErrors(next(counter)) + gep.addToSignalVertex(particle) + + gpi_coll = edm4hep.GeneratorPdfInfoCollection() + gpi = gpi_coll.create() + # Doesn't work with ROOT 6.30.06 + # gpi.setPartonId([next(counter), next(counter)]) + gpi.setPartonId(0, next(counter)) + gpi.setPartonId(1, next(counter)) + # Doesn't work with ROOT 6.30.06 + # gpi.setLhapdfId([next(counter), next(counter)]) + gpi.setLhapdfId(0, next(counter)) + gpi.setLhapdfId(1, next(counter)) + # Doesn't work with ROOT 6.30.06 + # gpi.setX([next(counter), next(counter)]) + gpi.setX(0, next(counter)) + gpi.setX(1, next(counter)) + # Doesn't work with ROOT 6.30.06 + # gpi.setXf([next(counter), next(counter)]) + gpi.setXf(0, next(counter)) + gpi.setXf(1, next(counter)) + gpi.setScale(next(counter)) + + writer.write_frame(frame, "events") diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 3e238e52e..ee6306bdb 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -6,12 +6,16 @@ function(set_test_env _testname) set_property(TEST ${_testname} APPEND PROPERTY ENVIRONMENT LD_LIBRARY_PATH=$:$:$<$:$:>$:$ENV{LD_LIBRARY_PATH} ROOT_INCLUDE_PATH=${PROJECT_SOURCE_DIR}/edm4hep:${PROJECT_SOURCE_DIR}/utils/include:$ENV{ROOT_INCLUDE_PATH} + PYTHONPATH=${PROJECT_SOURCE_DIR}/python:$ENV{PYTHONPATH} ) set_tests_properties(${_testname} PROPERTIES FAIL_REGULAR_EXPRESSION "[^a-z]Error;ERROR;error;Failed" ) endfunction() +add_test(NAME "Create an EDM4hep data file" COMMAND python ${PROJECT_SOURCE_DIR}/scripts/createEDM4hepFile.py) +set_test_env("Create an EDM4hep data file") + 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)