Scripts/Python/digitization.py

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