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_input",
0053         )
0054         s.addReader(evGen)
0055
0056     outputDir = Path(outputDir)
0057     addFatras(
0058         s,
0059         trackingGeometry,
0060         field,
0061         rnd=rnd,
0062     )
0063
0064     addDigitization(
0065         s,
0066         trackingGeometry,
0067         field,
0068         digiConfigFile=digiConfigFile,
0069         outputDirCsv=outputDir / "csv" if outputCsv else None,
0070         outputDirRoot=outputDir if outputRoot else None,
0071         rnd=rnd,
0072         doMerge=doMerge,
0073     )
0074
0075     return s
0076
0077
0078 if "__main__" == __name__:
0079     detector = acts.examples.GenericDetector()
0080     trackingGeometry = detector.trackingGeometry()
0081
0082     digiConfigFile = (
0083         Path(__file__).resolve().parent.parent.parent.parent
0084         / "Examples/Algorithms/Digitization/share/default-smearing-config-generic.json"
0085     )
0086     assert digiConfigFile.exists()
0087
0088     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0089
0090     runDigitization(trackingGeometry, field, outputDir=Path.cwd()).run()