Scripts/Python/geomodel_G4.py

0001 import os
0002 import acts
0003 import argparse
0004 from acts import (
0005     logging,
0006     GeometryContext,
0007 )
0008
0009 from acts.examples import (
0010     AlgorithmContext,
0011     WhiteBoard,
0012     ObjTrackingGeometryWriter,
0013 )
0014 from acts import geomodel as gm
0015 from acts.examples import geomodel as gmexample
0016 from acts import examples
0017 import acts.examples.geomodel as gm_ex
0018
0019 from pathlib import Path
0020 from propagation import runPropagation
0021
0022
0023 def runGeant4(
0024     detector,
0025     trackingGeometry,
0026     field,
0027     outputDir,
0028     materialMappings=[],
0029     volumeMappings=[],
0030     s: acts.examples.Sequencer = None,
0031     events=100,
0032     seed=1602,
0033     nMuonPerEvt=1,
0034 ):
0035     from acts.examples.simulation import (
0036         addParticleGun,
0037         addGeant4,
0038         EtaConfig,
0039         MomentumConfig,
0040         PhiConfig,
0041         ParticleConfig,
0042     )
0043     from pathlib import Path
0044
0045     s = s or acts.examples.Sequencer(events=events, numThreads=1)
0046     s.config.logLevel = acts.logging.INFO
0047     rnd = acts.examples.RandomNumbers(acts.examples.RandomNumbers.Config(seed=seed))
0048     u = acts.UnitConstants
0049     outputDir = Path(outputDir)
0050     addParticleGun(
0051         s,
0052         outputDirRoot=outputDir / "PG",
0053         momentumConfig=MomentumConfig(50.0 * u.GeV, 60.0 * u.GeV, transverse=True),
0054         etaConfig=EtaConfig(-1.0, 1.0, uniform=True),
0055         phiConfig=PhiConfig(0.0, 360.0 * u.degree),
0056         particleConfig=ParticleConfig(1, acts.PdgParticle.eMuon, randomizeCharge=True),
0057         vtxGen=acts.examples.GaussianVertexGenerator(
0058             mean=acts.Vector4(0, 0, 0, 0),
0059             stddev=acts.Vector4(0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 1.0 * u.ns),
0060         ),
0061         multiplicity=nMuonPerEvt,
0062         rnd=rnd,
0063         printParticles=False,
0064     )
0065
0066     addGeant4(
0067         s,
0068         detector,
0069         trackingGeometry,
0070         field,
0071         outputDirCsv=outputDir / "csv",
0072         outputDirRoot=outputDir,
0073         outputDirObj=outputDir / "obj",
0074         rnd=rnd,
0075         materialMappings=materialMappings,
0076         volumeMappings=volumeMappings,
0077     )
0078     return s
0079
0080
0081 def main():
0082     from argparse import ArgumentParser
0083
0084     u = acts.UnitConstants
0085
0086     parser = ArgumentParser()
0087     parser.add_argument(
0088         "-i",
0089         "--input",
0090         type=str,
0091         default="",  # "/eos/user/c/cimuonsw/GeometryFiles/MockUp.db",
0092         help="Input SQL file",
0093     )
0094     parser.add_argument(
0095         "--mockupDetector",
0096         type=str,
0097         choices=["Muon"],
0098         help="Predefined mockup detector which is built transiently",
0099         default="Muon",
0100     )
0101     parser.add_argument("--outDir", default="./", help="Output")
0102     parser.add_argument("--nEvents", default=100, type=int, help="Number of events")
0103     parser.add_argument(
0104         "--randomSeed", default=1602, type=int, help="Random seed for event generation"
0105     )
0106     parser.add_argument(
0107         "--geoSvgDump",
0108         default=False,
0109         action="store_true",
0110         help="Dump the tracking geometry in an obj format",
0111     )
0112
0113     args = parser.parse_args()
0114
0115     gContext = acts.GeometryContext.dangerouslyDefaultConstruct()
0116     logLevel = logging.INFO
0117
0118     # Create the tracking geometry builder for the muon system
0119     gmBuilderConfig = gm_ex.GeoModelMuonMockupBuilder.Config()
0120
0121     # Read the geometry model from the database
0122     gmTree = None
0123     ### Use an external geo model file
0124     if len(args.input):
0125         gmTree = gm.readFromDb(args.input)
0126         gmBuilderConfig.stationNames = ["BIL", "BML", "BOL"]
0127
0128     elif args.mockupDetector == "Muon":
0129         mockUpCfg = gm_ex.GeoMuonMockupExperiment.Config()
0130         mockUpCfg.dumpTree = True
0131         mockUpCfg.dbName = "ActsGeoMS.db"
0132         mockUpCfg.nSectors = 12
0133         mockUpCfg.nEtaStations = 8
0134         mockUpCfg.buildEndcaps = True
0135         mockUpBuilder = gm_ex.GeoMuonMockupExperiment(
0136             mockUpCfg, "GeoMockUpMS", logLevel
0137         )
0138         gmBuilderConfig.stationNames = [
0139             "BI",
0140             "BM",
0141             "BO",
0142             "EAI",
0143             "EAM",
0144             "EAO",
0145             "ECI",
0146             "ECM",
0147             "ECO",
0148         ]
0149
0150         gmTree = mockUpBuilder.constructMS()
0151     else:
0152         raise RuntimeError(f"{args.mockupDetector} not implemented yet")
0153
0154     gmFactoryConfig = gm.GeoModelDetectorObjectFactory.Config()
0155     gmFactoryConfig.nameList = [
0156         "RpcGasGap",
0157         "MDTDriftGas",
0158         "TgcGasGap",
0159         "SmallWheelGasGap",
0160     ]
0161     gmFactoryConfig.convertSubVolumes = True
0162     gmFactoryConfig.convertBox = ["MDT", "RPC", "SmallWheel", "TGC"]
0163
0164     gmFactory = gm.GeoModelDetectorObjectFactory(gmFactoryConfig, logLevel)
0165     # The options
0166     gmFactoryOptions = gm.GeoModelDetectorObjectFactory.Options()
0167     gmFactoryOptions.queries = ["Muon"]
0168
0169     # The Cache & construct call
0170     gmFactoryCache = gm.GeoModelDetectorObjectFactory.Cache()
0171     gmFactory.construct(gmFactoryCache, gContext, gmTree, gmFactoryOptions)
0172
0173     gmBuilderConfig.volumeBoxFPVs = gmFactoryCache.boundingBoxes
0174
0175     gmDetectorCfg = gm_ex.GeoModelDetector.Config()
0176     gmDetectorCfg.geoModelTree = gmTree
0177     detector = gm_ex.GeoModelDetector(gmDetectorCfg)
0178
0179     field = acts.ConstantBField(acts.Vector3(0, 0, 0 * u.T))
0180
0181     trackingGeometryBuilder = gm_ex.GeoModelMuonMockupBuilder(
0182         gmBuilderConfig, "GeoModelMuonMockupBuilder", logLevel
0183     )
0184
0185     trackingGeometry = detector.buildTrackingGeometry(gContext, trackingGeometryBuilder)
0186
0187     algSequence = runGeant4(
0188         detector=detector,
0189         trackingGeometry=trackingGeometry,
0190         field=field,
0191         outputDir=args.outDir,
0192         volumeMappings=gmFactoryConfig.nameList,
0193         events=args.nEvents,
0194         seed=args.randomSeed,
0195     )
0196
0197     from acts.examples import MuonSpacePointDigitizer
0198
0199     digiAlg = MuonSpacePointDigitizer(
0200         randomNumbers=acts.examples.RandomNumbers(
0201             acts.examples.RandomNumbers.Config(seed=2 * args.randomSeed)
0202         ),
0203         trackingGeometry=trackingGeometry,
0204         dumpVisualization=False,
0205         digitizeTime=True,
0206         outputSpacePoints="MuonSpacePoints",
0207         level=logLevel,
0208     )
0209     algSequence.addAlgorithm(digiAlg)
0210
0211     from acts.examples.root import RootMuonSpacePointWriter
0212
0213     algSequence.addWriter(
0214         RootMuonSpacePointWriter(
0215             level=logLevel,
0216             inputSpacePoints="MuonSpacePoints",
0217             filePath=f"{args.outDir}/MS_SpacePoints.root",
0218         )
0219     )
0220
0221     if args.geoSvgDump:
0222         wb = WhiteBoard(acts.logging.INFO)
0223         context = AlgorithmContext(0, 0, wb, 10)
0224         obj_dir = Path(args.outDir) / "obj"
0225         obj_dir.mkdir(exist_ok=True)
0226         writer = ObjTrackingGeometryWriter(
0227             level=acts.logging.INFO, outputDir=str(obj_dir)
0228         )
0229
0230         writer.write(context, trackingGeometry)
0231
0232     algSequence.run()
0233
0234
0235 if __name__ == "__main__":
0236     main()