Scripts/Python/digitization.py

0001 #!/usr/bin/env python3
0002
0003 from pathlib import Path
0004 from typing import Optional
0005 import argparse
0006
0007 import acts
0008 import acts.examples
0009
0010
0011 u = acts.UnitConstants
0012
0013
0014 def runDigitization(
0015     trackingGeometry: acts.TrackingGeometry,
0016     field: acts.MagneticFieldProvider,
0017     outputDir: Path,
0018     digiConfigFile: Path,
0019     particlesInput: Optional[Path] = None,
0020     outputRoot: bool = True,
0021     outputCsv: bool = True,
0022     s: Optional[acts.examples.Sequencer] = None,
0023     doMerge: Optional[bool] = None,
0024 ) -> acts.examples.Sequencer:
0025     from acts.examples.simulation import (
0026         addParticleGun,
0027         EtaConfig,
0028         PhiConfig,
0029         ParticleConfig,
0030         addFatras,
0031         addDigitization,
0032     )
0033
0034     s = s or acts.examples.Sequencer(
0035         events=100, numThreads=-1, logLevel=acts.logging.INFO
0036     )
0037     rnd = acts.examples.RandomNumbers(seed=42)
0038
0039     if particlesInput is None:
0040         addParticleGun(
0041             s,
0042             EtaConfig(-4.0, 4.0),
0043             ParticleConfig(4, acts.PdgParticle.eMuon, True),
0044             PhiConfig(0.0, 360.0 * u.degree),
0045             multiplicity=2,
0046             rnd=rnd,
0047         )
0048     else:
0049         # Read input from input collection (e.g. Pythia8 output)
0050         evGen = acts.examples.RootParticleReader(
0051             level=s.config.logLevel,
0052             filePath=str(particlesInput),
0053             outputParticles="particles_generated",
0054         )
0055         s.addReader(evGen)
0056
0057         s.addWhiteboardAlias(
0058             "particles_generated_selected", evGen.config.outputParticles
0059         )
0060
0061     outputDir = Path(outputDir)
0062     addFatras(
0063         s,
0064         trackingGeometry,
0065         field,
0066         rnd=rnd,
0067     )
0068
0069     addDigitization(
0070         s,
0071         trackingGeometry,
0072         field,
0073         digiConfigFile=digiConfigFile,
0074         outputDirCsv=outputDir / "csv" if outputCsv else None,
0075         outputDirRoot=outputDir if outputRoot else None,
0076         rnd=rnd,
0077         doMerge=doMerge,
0078     )
0079
0080     return s
0081
0082
0083 if "__main__" == __name__:
0084
0085     # Parse the command line arguments
0086     p = argparse.ArgumentParser(description="Digitization")
0087     p.add_argument(
0088         "--events",
0089         "-n",
0090         type=int,
0091         help="Number of events",
0092         default=1000,
0093     )
0094     p.add_argument(
0095         "--type",
0096         "-t",
0097         type=str,
0098         help="Type of digitization",
0099         default="smearing",
0100         choices=["smearing", "geometric"],
0101     )
0102     p.add_argument(
0103         "--detector",
0104         "-d",
0105         type=str,
0106         help="Output file",
0107         default="generic",
0108         choices=["generic", "odd"],
0109     )
0110     args = p.parse_args()
0111
0112     if args.detector == "generic":
0113         detector = acts.examples.GenericDetector()
0114     else:
0115         from acts.examples.odd import getOpenDataDetector
0116
0117         detector = getOpenDataDetector()
0118
0119     trackingGeometry = detector.trackingGeometry()
0120
0121     digiConfigFile = (
0122         Path(__file__).resolve().parent.parent.parent.parent
0123         / "Examples/Configs"
0124         / f"{args.detector}-digi-{args.type}-config.json"
0125     )
0126     assert digiConfigFile.exists()
0127
0128     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0129     s = acts.examples.Sequencer(
0130         events=args.events, numThreads=-1, logLevel=acts.logging.INFO
0131     )
0132
0133     runDigitization(
0134         trackingGeometry,
0135         field,
0136         outputDir=Path.cwd(),
0137         digiConfigFile=digiConfigFile,
0138         s=s,
0139     ).run()