Skip to content

Commit

Permalink
Format
Browse files Browse the repository at this point in the history
  • Loading branch information
@jmcarcell
jmcarcell committed Feb 21, 2024
1 parent ad50b91 commit 655e3d2
Showing 1 changed file with 118 additions and 66 deletions.
184 changes: 118 additions & 66 deletions scripts/createEDM4hepFile.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,29 +8,33 @@
import argparse
import sys

frames = 10 # 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
frames = 10 # 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('--use-pre1', action='store_true', help='Use a pre 1.0 version of EDM4hep')
parser.add_argument('--rntuple', action='store_true', help='Use a ROOT ntuple instead of EDM4hep')
parser = argparse.ArgumentParser(description="Create a file with EDM4hep data")
parser.add_argument(
"--use-pre1", action="store_true", help="Use a pre 1.0 version of EDM4hep"
)
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')
print("The RNTuple writer from podio is not available but was requested")
sys.exit(1)


writer = podio.root_io.Writer(output_file)

for i in range(frames):
print(f'Writing frame {i}')
print(f"Writing frame {i}")
frame = podio.Frame()

header = edm4hep.EventHeaderCollection()
Expand All @@ -52,15 +56,29 @@
parent_particle.setCharge(next(counter))
parent_particle.setTime(next(counter))
parent_particle.setMass(next(counter))
parent_particle.setVertex(edm4hep.Vector3d(next(counter), next(counter), next(counter)))
parent_particle.setEndpoint(edm4hep.Vector3d(next(counter), next(counter), next(counter)))
parent_particle.setVertex(
edm4hep.Vector3d(next(counter), next(counter), next(counter))
)
parent_particle.setEndpoint(
edm4hep.Vector3d(next(counter), next(counter), next(counter))
)
if not args.use_pre1:
parent_particle.setMomentum(edm4hep.Vector3d(next(counter), next(counter), next(counter)))
parent_particle.setMomentumAtEndpoint(edm4hep.Vector3d(next(counter), next(counter), next(counter)))
parent_particle.setMomentum(
edm4hep.Vector3d(next(counter), next(counter), next(counter))
)
parent_particle.setMomentumAtEndpoint(
edm4hep.Vector3d(next(counter), next(counter), next(counter))
)
else:
parent_particle.setMomentum(edm4hep.Vector3f(next(counter), next(counter), next(counter)))
parent_particle.setMomentumAtEndpoint(edm4hep.Vector3f(next(counter), next(counter), next(counter)))
parent_particle.setSpin(edm4hep.Vector3f(next(counter), next(counter), next(counter)))
parent_particle.setMomentum(
edm4hep.Vector3f(next(counter), next(counter), next(counter))
)
parent_particle.setMomentumAtEndpoint(
edm4hep.Vector3f(next(counter), next(counter), next(counter))
)
parent_particle.setSpin(
edm4hep.Vector3f(next(counter), next(counter), next(counter))
)
parent_particle.setColorFlow(edm4hep.Vector2i(next(counter), next(counter)))

daughter_particle = particles.create()
Expand All @@ -70,19 +88,37 @@
daughter_particle.setCharge(next(counter))
daughter_particle.setTime(next(counter))
daughter_particle.setMass(next(counter))
daughter_particle.setVertex(edm4hep.Vector3d(next(counter), next(counter), next(counter)))
daughter_particle.setEndpoint(edm4hep.Vector3d(next(counter), next(counter), next(counter)))
daughter_particle.setVertex(
edm4hep.Vector3d(next(counter), next(counter), next(counter))
)
daughter_particle.setEndpoint(
edm4hep.Vector3d(next(counter), next(counter), next(counter))
)
if not args.use_pre1:
daughter_particle.setMomentum(edm4hep.Vector3d(next(counter), next(counter), next(counter)))
daughter_particle.setMomentumAtEndpoint(edm4hep.Vector3d(next(counter), next(counter), next(counter)))
daughter_particle.setMomentum(
edm4hep.Vector3d(next(counter), next(counter), next(counter))
)
daughter_particle.setMomentumAtEndpoint(
edm4hep.Vector3d(next(counter), next(counter), next(counter))
)
else:
daughter_particle.setMomentum(edm4hep.Vector3f(next(counter), next(counter), next(counter)))
daughter_particle.setMomentumAtEndpoint(edm4hep.Vector3f(next(counter), next(counter), next(counter)))
daughter_particle.setSpin(edm4hep.Vector3f(next(counter), next(counter), next(counter)))
daughter_particle.setMomentum(
edm4hep.Vector3f(next(counter), next(counter), next(counter))
)
daughter_particle.setMomentumAtEndpoint(
edm4hep.Vector3f(next(counter), next(counter), next(counter))
)
daughter_particle.setSpin(
edm4hep.Vector3f(next(counter), next(counter), next(counter))
)
daughter_particle.setColorFlow(edm4hep.Vector2i(next(counter), next(counter)))

parent_particle.addToDaughters(cppyy.ll.static_cast['edm4hep::MCParticle'](daughter_particle))
daughter_particle.addToParents(cppyy.ll.static_cast['edm4hep::MCParticle'](parent_particle))

parent_particle.addToDaughters(
cppyy.ll.static_cast["edm4hep::MCParticle"](daughter_particle)
)
daughter_particle.addToParents(
cppyy.ll.static_cast["edm4hep::MCParticle"](parent_particle)
)
frame.put(particles, "MCParticleCollection")

hits = edm4hep.SimTrackerHitCollection()
Expand All @@ -94,7 +130,7 @@
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(cppyy.ll.static_cast['edm4hep::MCParticle'](daughter_particle))
hit.setParticle(cppyy.ll.static_cast["edm4hep::MCParticle"](daughter_particle))
simtracker_hit = hit
frame.put(hits, "SimTrackerHitCollection")

Expand All @@ -104,7 +140,7 @@
hit.setEnergy(next(counter))
hit.setTime(next(counter))
hit.setStepPosition(edm4hep.Vector3f(next(counter), next(counter), next(counter)))
hit.setParticle(cppyy.ll.static_cast['edm4hep::MCParticle'](daughter_particle))
hit.setParticle(cppyy.ll.static_cast["edm4hep::MCParticle"](daughter_particle))
calohit = hit
frame.put(hits, "CaloHitContributionCollection")

Expand All @@ -113,7 +149,9 @@
hit.setCellID(next(counter))
hit.setEnergy(next(counter))
hit.setPosition(edm4hep.Vector3f(next(counter), next(counter), next(counter)))
hit.addToContributions(cppyy.ll.static_cast['edm4hep::CaloHitContribution'](calohit))
hit.addToContributions(
cppyy.ll.static_cast["edm4hep::CaloHitContribution"](calohit)
)
simcalo_hit = hit
frame.put(hits, "SimCalorimeterHitCollection")

Expand Down Expand Up @@ -155,15 +193,17 @@
cluster.setPositionError(j, next(counter))
cluster.setITheta(next(counter))
cluster.setPhi(next(counter))
cluster.setDirectionError(edm4hep.Vector3f(next(counter), next(counter), next(counter)))
cluster.setDirectionError(
edm4hep.Vector3f(next(counter), next(counter), next(counter))
)
for j in range(vectorsize):
cluster.addToShapeParameters(next(counter))
for j in range(vectorsize):
cluster.addToSubdetectorEnergies(next(counter))
cluster.addToClusters(cppyy.ll.static_cast['edm4hep::Cluster'](cluster))
cluster.addToHits(cppyy.ll.static_cast['edm4hep::CalorimeterHit'](calo_hit))
cluster.addToClusters(cppyy.ll.static_cast["edm4hep::Cluster"](cluster))
cluster.addToHits(cppyy.ll.static_cast["edm4hep::CalorimeterHit"](calo_hit))
if args.use_pre1:
cluster.addToParticleIDs(cppyy.ll.static_cast['edm4hep::ParticleID'](pid))
cluster.addToParticleIDs(cppyy.ll.static_cast["edm4hep::ParticleID"](pid))
frame.put(clusters, "ClusterCollection")

if args.use_pre1:
Expand Down Expand Up @@ -241,13 +281,15 @@
state.z0 = next(counter)
state.tanLambda = next(counter)
state.time = next(counter)
state.referencePoint = edm4hep.Vector3f(next(counter), next(counter), next(counter))
state.referencePoint = edm4hep.Vector3f(
next(counter), next(counter), next(counter)
)
# TODO
# for k in range(21):
# state.setCovMatrix(k, next(counter))
track.addToTrackStates(state)
track.addToTrackerHits(cppyy.ll.static_cast['edm4hep::TrackerHit3D'](tracker_hit))
track.addToTracks(cppyy.ll.static_cast['edm4hep::Track'](track))
track.addToTrackerHits(cppyy.ll.static_cast["edm4hep::TrackerHit3D"](tracker_hit))
track.addToTracks(cppyy.ll.static_cast["edm4hep::Track"](track))
frame.put(tracks, "TrackCollection")

vertex = edm4hep.VertexCollection()
Expand All @@ -268,75 +310,83 @@
particle.setType(next(counter))
particle.setEnergy(next(counter))
particle.setMomentum(edm4hep.Vector3f(next(counter), next(counter), next(counter)))
particle.setReferencePoint(edm4hep.Vector3f(next(counter), next(counter), next(counter)))
particle.setReferencePoint(
edm4hep.Vector3f(next(counter), next(counter), next(counter))
)
particle.setCharge(next(counter))
particle.setMass(next(counter))
particle.setGoodnessOfPID(next(counter))
for j in range(10):
particle.setCovMatrix(j, next(counter))
particle.setStartVertex(cppyy.ll.static_cast['edm4hep::Vertex'](v))
particle.setStartVertex(cppyy.ll.static_cast["edm4hep::Vertex"](v))
if args.use_pre1:
particle.setParticleIDUsed(cppyy.ll.static_cast['edm4hep::ParticleID'](pid))
particle.addToClusters(cppyy.ll.static_cast['edm4hep::Cluster'](cluster))
particle.addToTracks(cppyy.ll.static_cast['edm4hep::Track'](track))
particle.addToParticles(cppyy.ll.static_cast['edm4hep::ReconstructedParticle'](particle))
particle.setParticleIDUsed(cppyy.ll.static_cast["edm4hep::ParticleID"](pid))
particle.addToClusters(cppyy.ll.static_cast["edm4hep::Cluster"](cluster))
particle.addToTracks(cppyy.ll.static_cast["edm4hep::Track"](track))
particle.addToParticles(
cppyy.ll.static_cast["edm4hep::ReconstructedParticle"](particle)
)
if args.use_pre1:
particle.addToParticleIDs(cppyy.ll.static_cast['edm4hep::ParticleID'](pid))
particle.addToParticleIDs(cppyy.ll.static_cast["edm4hep::ParticleID"](pid))
reco_particle = particle
frame.put(particles, "ReconstructedParticleCollection")

v.setAssociatedParticle(cppyy.ll.static_cast['edm4hep::ReconstructedParticle'](reco_particle))
v.setAssociatedParticle(
cppyy.ll.static_cast["edm4hep::ReconstructedParticle"](reco_particle)
)

if not args.use_pre1:
pid.setParticle(cppyy.ll.static_cast['edm4hep::ReconstructedParticle'](reco_particle))
pid.setParticle(
cppyy.ll.static_cast["edm4hep::ReconstructedParticle"](reco_particle)
)

assocs = edm4hep.MCRecoParticleAssociationCollection()
assoc = assocs.create()
assoc.setWeight(next(counter))
assoc.setRec(cppyy.ll.static_cast['edm4hep::ReconstructedParticle'](reco_particle))
assoc.setSim(cppyy.ll.static_cast['edm4hep::MCParticle'](daughter_particle))
assoc.setRec(cppyy.ll.static_cast["edm4hep::ReconstructedParticle"](reco_particle))
assoc.setSim(cppyy.ll.static_cast["edm4hep::MCParticle"](daughter_particle))
frame.put(assocs, "MCRecoParticleAssociationCollection")

assocs = edm4hep.MCRecoCaloAssociationCollection()
assoc = assocs.create()
assoc.setWeight(next(counter))
assoc.setRec(cppyy.ll.static_cast['edm4hep::CalorimeterHit'](calo_hit))
assoc.setSim(cppyy.ll.static_cast['edm4hep::SimCalorimeterHit'](simcalo_hit))
assoc.setRec(cppyy.ll.static_cast["edm4hep::CalorimeterHit"](calo_hit))
assoc.setSim(cppyy.ll.static_cast["edm4hep::SimCalorimeterHit"](simcalo_hit))
frame.put(assocs, "MCRecoCaloAssociationCollection")

assocs = edm4hep.MCRecoTrackerAssociationCollection()
assoc = assocs.create()
assoc.setWeight(next(counter))
assoc.setRec(cppyy.ll.static_cast['edm4hep::TrackerHit3D'](tracker_hit))
assoc.setSim(cppyy.ll.static_cast['edm4hep::SimTrackerHit'](simtracker_hit))
assoc.setRec(cppyy.ll.static_cast["edm4hep::TrackerHit3D"](tracker_hit))
assoc.setSim(cppyy.ll.static_cast["edm4hep::SimTrackerHit"](simtracker_hit))
frame.put(assocs, "MCRecoTrackerAssociationCollection")

assocs = edm4hep.MCRecoCaloParticleAssociationCollection()
assoc = assocs.create()
assoc.setWeight(next(counter))
assoc.setRec(cppyy.ll.static_cast['edm4hep::CalorimeterHit'](calo_hit))
assoc.setSim(cppyy.ll.static_cast['edm4hep::MCParticle'](daughter_particle))
assoc.setRec(cppyy.ll.static_cast["edm4hep::CalorimeterHit"](calo_hit))
assoc.setSim(cppyy.ll.static_cast["edm4hep::MCParticle"](daughter_particle))
frame.put(assocs, "MCRecoCaloParticleAssociationCollection")

assocs = edm4hep.MCRecoClusterParticleAssociationCollection()
assoc = assocs.create()
assoc.setWeight(next(counter))
assoc.setRec(cppyy.ll.static_cast['edm4hep::Cluster'](cluster))
assoc.setSim(cppyy.ll.static_cast['edm4hep::MCParticle'](daughter_particle))
assoc.setRec(cppyy.ll.static_cast["edm4hep::Cluster"](cluster))
assoc.setSim(cppyy.ll.static_cast["edm4hep::MCParticle"](daughter_particle))
frame.put(assocs, "MCRecoClusterParticleAssociationCollection")

assocs = edm4hep.MCRecoTrackParticleAssociationCollection()
assoc = assocs.create()
assoc.setWeight(next(counter))
assoc.setRec(cppyy.ll.static_cast['edm4hep::Track'](track))
assoc.setSim(cppyy.ll.static_cast['edm4hep::MCParticle'](daughter_particle))
assoc.setRec(cppyy.ll.static_cast["edm4hep::Track"](track))
assoc.setSim(cppyy.ll.static_cast["edm4hep::MCParticle"](daughter_particle))
frame.put(assocs, "MCRecoTrackParticleAssociationCollection")

assocs = edm4hep.RecoParticleVertexAssociationCollection()
assoc = assocs.create()
assoc.setWeight(next(counter))
assoc.setRec(cppyy.ll.static_cast['edm4hep::ReconstructedParticle'](reco_particle))
assoc.setVertex(cppyy.ll.static_cast['edm4hep::Vertex'](v))
assoc.setRec(cppyy.ll.static_cast["edm4hep::ReconstructedParticle"](reco_particle))
assoc.setVertex(cppyy.ll.static_cast["edm4hep::Vertex"](v))
frame.put(assocs, "MCRecoParticleVertexAssociationCollection")

simiocluster = edm4hep.SimPrimaryIonizationClusterCollection()
Expand All @@ -348,10 +398,12 @@
for j in range(vectorsize):
cluster.addToElectronCellID(next(counter))
cluster.addToElectronTime(next(counter))
cluster.addToElectronPosition(edm4hep.Vector3d(next(counter), next(counter), next(counter)))
cluster.addToElectronPosition(
edm4hep.Vector3d(next(counter), next(counter), next(counter))
)
cluster.addToPulseTime(next(counter))
cluster.addToPulseAmplitude(next(counter))
cluster.setParticle(cppyy.ll.static_cast['edm4hep::MCParticle'](daughter_particle))
cluster.setParticle(cppyy.ll.static_cast["edm4hep::MCParticle"](daughter_particle))
frame.put(simiocluster, "SimPrimaryIonizationClusterCollection")

timeseries = edm4hep.TimeSeriesCollection()
Expand All @@ -371,15 +423,15 @@
pulse.setQuality(next(counter))
for j in range(3):
pulse.setCovMatrix(j, next(counter))
pulse.setTimeSeries(cppyy.ll.static_cast['edm4hep::TimeSeries'](serie))
pulse.setTimeSeries(cppyy.ll.static_cast["edm4hep::TimeSeries"](serie))
frame.put(trackerpulse, "TrackerPulseCollection")

recioncluster = edm4hep.RecIonizationClusterCollection()
cluster = recioncluster.create()
cluster.setCellID(next(counter))
cluster.setSignificance(next(counter))
cluster.setType(next(counter))
cluster.addToTrackerPulse(cppyy.ll.static_cast['edm4hep::TrackerPulse'](pulse))
cluster.addToTrackerPulse(cppyy.ll.static_cast["edm4hep::TrackerPulse"](pulse))
frame.put(recioncluster, "RecIonizationClusterCollection")

recdqdx = edm4hep.RecDqdxCollection()
Expand All @@ -399,7 +451,7 @@
hd.eDep = next(counter)
hd.pathLength = next(counter)
dqdx.addToHitData(hd)
dqdx.setTrack(cppyy.ll.static_cast['edm4hep::Track'](track))
dqdx.setTrack(cppyy.ll.static_cast["edm4hep::Track"](track))
frame.put(recdqdx, "RecDqdxCollection")

writer.write_frame(frame, "events")

0 comments on commit 655e3d2

Please sign in to comment.