Tracking/material_map/material_recording_ePIC.py

0001 #!/usr/bin/env python3
0002
0003 import os
0004 import warnings
0005 from pathlib import Path
0006 import argparse
0007
0008 import acts
0009 from acts.examples import (
0010     GaussianVertexGenerator,
0011     ParametricParticleGenerator,
0012     FixedMultiplicityGenerator,
0013     EventGenerator,
0014     RandomNumbers,
0015 )
0016
0017 import acts.examples.dd4hep
0018 import acts.examples.geant4
0019 import acts.examples.geant4.dd4hep
0020 from acts.examples.odd import getOpenDataDetector
0021 from acts.examples.geant4 import GdmlDetectorConstructionFactory
0022
0023 u = acts.UnitConstants
0024
0025 import build_geometry_ePIC as geo
0026
0027 _material_recording_executed = False
0028
0029
0030 def runMaterialRecording(
0031     detectorConstructionFactory,
0032     outputDir,
0033     outputName="geant4_material_tracks.root",
0034     tracksPerEvent=1000,
0035     s=None,
0036     etaRange=(-8, 8),
0037 ):
0038     global _material_recording_executed
0039     if _material_recording_executed:
0040         warnings.warn("Material recording already ran in this process. Expect crashes")
0041     _material_recording_executed = True
0042
0043     rnd = RandomNumbers(seed=228)
0044
0045     evGen = EventGenerator(
0046         level=acts.logging.INFO,
0047         generators=[
0048             EventGenerator.Generator(
0049                 multiplicity=FixedMultiplicityGenerator(n=1),
0050                 vertex=GaussianVertexGenerator(
0051                     stddev=acts.Vector4(0, 0, 0, 0),
0052                     mean=acts.Vector4(0, 0, 0, 0),
0053                 ),
0054                 particles=ParametricParticleGenerator(
0055                     pdg=acts.PdgParticle.eInvalid,
0056                     charge=0,
0057                     randomizeCharge=False,
0058                     mass=0,
0059                     p=(1 * u.GeV, 10 * u.GeV),
0060                     eta=etaRange,
0061                     numParticles=tracksPerEvent,
0062                     etaUniform=True,
0063                 ),
0064             )
0065         ],
0066         outputParticles="particles_initial",
0067         randomNumbers=rnd,
0068     )
0069
0070     s.addReader(evGen)
0071
0072     g4Alg = acts.examples.geant4.Geant4MaterialRecording(
0073         level=acts.logging.INFO,
0074         detectorConstructionFactory=detectorConstructionFactory,
0075         randomNumbers=rnd,
0076         inputParticles=evGen.config.outputParticles,
0077         outputMaterialTracks="material-tracks",
0078     )
0079
0080     s.addAlgorithm(g4Alg)
0081
0082     outpath=os.path.join(outputDir, outputName)
0083     s.addWriter(
0084         acts.examples.RootMaterialTrackWriter(
0085             prePostStep=True,
0086             recalculateTotals=True,
0087             collection="material-tracks",
0088             filePath=outpath,
0089             level=acts.logging.INFO,
0090         )
0091     )
0092     print("Done! Recorded steps in "+outpath)
0093
0094     return s
0095
0096 # python material_recording_ePIC.py -i epic_craterlake_matmap.xml -n 1000 -t 1000
0097 if "__main__" == __name__:
0098
0099     p = argparse.ArgumentParser(description="Record the Geant4 materials with geantino")
0100     p.add_argument(
0101         "-n", "--nevents", type=int, default=1000, help="Number of events to generate"
0102     )
0103     p.add_argument(
0104         "-t", "--ntracks", type=int, default=1000, help="Particle tracks per event"
0105     )
0106     p.add_argument(
0107         "-i", "--xmlFile", type=str, default=os.environ.get("DETECTOR_PATH", "")+"epic_craterlake.xml", help="DD4hep input file"
0108     )
0109     p.add_argument(
0110         "-o", "--outputName", type=str, default="geant4_material_tracks.root", help="Name of the output rootfile"
0111     )
0112     p.add_argument(
0113         "--eta_min",
0114         type=float,
0115         default=-8.0,
0116         help="eta min (optional)",
0117     )
0118     p.add_argument(
0119         "--eta_max",
0120         type=float,
0121         default=8.0,
0122         help="eta max (optional)",
0123     )
0124     args = p.parse_args()
0125
0126     detectorConstructionFactory = None
0127
0128     detector, trackingGeometry, decorators = geo.buildePICGeometry(
0129         args.xmlFile)
0130
0131     detectorConstructionFactory = (
0132             acts.examples.geant4.dd4hep.DDG4DetectorConstructionFactory(detector)
0133         )
0134
0135     runMaterialRecording(
0136         detectorConstructionFactory=detectorConstructionFactory,
0137         tracksPerEvent=args.ntracks,
0138         outputDir=os.getcwd(),
0139         outputName=args.outputName,
0140         etaRange=(args.eta_min, args.eta_max),
0141         s=acts.examples.Sequencer(events=args.nevents, numThreads=1),
0142     ).run()
0143