Scripts/Python/digitization.py

0001 #!/usr/bin/env python3
0002
0003 from pathlib import Path
0004 from typing import Optional
0005
0006 import acts
0007 import acts.examples
0008
0009
0010 u = acts.UnitConstants
0011
0012
0013 def runDigitization(
0014     trackingGeometry: acts.TrackingGeometry,
0015     field: acts.MagneticFieldProvider,
0016     outputDir: Path,
0017     digiConfigFile: Path,
0018     particlesInput: Optional[Path] = None,
0019     outputRoot: bool = True,
0020     outputCsv: bool = True,
0021     s: Optional[acts.examples.Sequencer] = None,
0022     doMerge: Optional[bool] = None,
0023 ) -> acts.examples.Sequencer:
0024     from acts.examples.simulation import (
0025         addParticleGun,
0026         EtaConfig,
0027         PhiConfig,
0028         ParticleConfig,
0029         addFatras,
0030         addDigitization,
0031     )
0032
0033     s = s or acts.examples.Sequencer(
0034         events=100, numThreads=-1, logLevel=acts.logging.INFO
0035     )
0036     rnd = acts.examples.RandomNumbers(seed=42)
0037
0038     if particlesInput is None:
0039         addParticleGun(
0040             s,
0041             EtaConfig(-2.0, 2.0),
0042             ParticleConfig(4, acts.PdgParticle.eMuon, True),
0043             PhiConfig(0.0, 360.0 * u.degree),
0044             multiplicity=2,
0045             rnd=rnd,
0046         )
0047     else:
0048         # Read input from input collection (e.g. Pythia8 output)
0049         evGen = acts.examples.RootParticleReader(
0050             level=s.config.logLevel,
0051             filePath=str(particlesInput),
0052             outputParticles="particles_generated",
0053         )
0054         s.addReader(evGen)
0055
0056         s.addWhiteboardAlias(
0057             "particles_generated_selected", evGen.config.outputParticles
0058         )
0059
0060     outputDir = Path(outputDir)
0061     addFatras(
0062         s,
0063         trackingGeometry,
0064         field,
0065         rnd=rnd,
0066     )
0067
0068     addDigitization(
0069         s,
0070         trackingGeometry,
0071         field,
0072         digiConfigFile=digiConfigFile,
0073         outputDirCsv=outputDir / "csv" if outputCsv else None,
0074         outputDirRoot=outputDir if outputRoot else None,
0075         rnd=rnd,
0076         doMerge=doMerge,
0077     )
0078
0079     return s
0080
0081
0082 if "__main__" == __name__:
0083     detector = acts.examples.GenericDetector()
0084     trackingGeometry = detector.trackingGeometry()
0085
0086     digiConfigFile = (
0087         Path(__file__).resolve().parent.parent.parent.parent
0088         / "Examples/Configs/generic-digi-smearing-config.json"
0089     )
0090     assert digiConfigFile.exists()
0091
0092     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0093
0094     runDigitization(trackingGeometry, field, outputDir=Path.cwd()).run()