Add a script to create an EDM4hep file making use of the complete datamodel

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")

test/CMakeLists.txt

@@ -6,12 +6,16 @@ function(set_test_env _testname)
   set_property(TEST ${_testname} APPEND PROPERTY ENVIRONMENT
       LD_LIBRARY_PATH=$<TARGET_FILE_DIR:edm4hep>:$<TARGET_FILE_DIR:podio::podio>:$<$<TARGET_EXISTS:edm4hepRDF>:$<TARGET_FILE_DIR:edm4hepRDF>:>$<TARGET_FILE_DIR:ROOT::Core>:$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)