Scripts/Python/material_mapping.py

0001 #!/usr/bin/env python3
0002
0003 import os
0004 import argparse
0005
0006 import acts
0007 from acts import (
0008     SurfaceMaterialMapper,
0009     VolumeMaterialMapper,
0010     Navigator,
0011     Propagator,
0012     StraightLineStepper,
0013     MaterialMapJsonConverter,
0014 )
0015 from acts.examples import (
0016     Sequencer,
0017     WhiteBoard,
0018     AlgorithmContext,
0019     ProcessCode,
0020     RootMaterialTrackReader,
0021     RootMaterialTrackWriter,
0022     MaterialMapping,
0023     JsonMaterialWriter,
0024     JsonFormat,
0025 )
0026 from acts.examples.odd import getOpenDataDetector
0027
0028
0029 def runMaterialMapping(
0030     trackingGeometry,
0031     decorators,
0032     outputDir,
0033     inputDir,
0034     mapName="material-map",
0035     mapFormat=JsonFormat.Json,
0036     mapSurface=True,
0037     mapVolume=True,
0038     readCachedSurfaceInformation=False,
0039     mappingStep=1,
0040     s=None,
0041 ):
0042     s = s or Sequencer(numThreads=1)
0043
0044     for decorator in decorators:
0045         s.addContextDecorator(decorator)
0046
0047     wb = WhiteBoard(acts.logging.INFO)
0048
0049     context = AlgorithmContext(0, 0, wb)
0050
0051     for decorator in decorators:
0052         assert decorator.decorate(context) == ProcessCode.SUCCESS
0053
0054     # Read material step information from a ROOT TTRee
0055     s.addReader(
0056         RootMaterialTrackReader(
0057             level=acts.logging.INFO,
0058             outputMaterialTracks="material-tracks",
0059             fileList=[
0060                 os.path.join(
0061                     inputDir,
0062                     (
0063                         mapName + "_tracks.root"
0064                         if readCachedSurfaceInformation
0065                         else "geant4_material_tracks.root"
0066                     ),
0067                 )
0068             ],
0069             readCachedSurfaceInformation=readCachedSurfaceInformation,
0070         )
0071     )
0072
0073     stepper = StraightLineStepper()
0074
0075     mmAlgCfg = MaterialMapping.Config(context.geoContext, context.magFieldContext)
0076     mmAlgCfg.trackingGeometry = trackingGeometry
0077     mmAlgCfg.inputMaterialTracks = "material-tracks"
0078
0079     if mapSurface:
0080         navigator = Navigator(
0081             trackingGeometry=trackingGeometry,
0082             resolveSensitive=True,
0083             resolveMaterial=True,
0084             resolvePassive=True,
0085         )
0086         propagator = Propagator(stepper, navigator)
0087         mapper = SurfaceMaterialMapper(level=acts.logging.INFO, propagator=propagator)
0088         mmAlgCfg.materialSurfaceMapper = mapper
0089
0090     if mapVolume:
0091         navigator = Navigator(
0092             trackingGeometry=trackingGeometry,
0093         )
0094         propagator = Propagator(stepper, navigator)
0095         mapper = VolumeMaterialMapper(
0096             level=acts.logging.INFO, propagator=propagator, mappingStep=mappingStep
0097         )
0098         mmAlgCfg.materialVolumeMapper = mapper
0099
0100     jmConverterCfg = MaterialMapJsonConverter.Config(
0101         processSensitives=True,
0102         processApproaches=True,
0103         processRepresenting=True,
0104         processBoundaries=True,
0105         processVolumes=True,
0106         context=context.geoContext,
0107     )
0108
0109     jmw = JsonMaterialWriter(
0110         level=acts.logging.VERBOSE,
0111         converterCfg=jmConverterCfg,
0112         fileName=os.path.join(outputDir, mapName),
0113         writeFormat=mapFormat,
0114     )
0115
0116     mmAlgCfg.materialWriters = [jmw]
0117
0118     s.addAlgorithm(MaterialMapping(level=acts.logging.INFO, config=mmAlgCfg))
0119
0120     s.addWriter(
0121         RootMaterialTrackWriter(
0122             level=acts.logging.INFO,
0123             inputMaterialTracks=mmAlgCfg.mappingMaterialCollection,
0124             filePath=os.path.join(
0125                 outputDir,
0126                 mapName + "_tracks.root",
0127             ),
0128             storeSurface=True,
0129             storeVolume=True,
0130         )
0131     )
0132
0133     return s
0134
0135
0136 if "__main__" == __name__:
0137     p = argparse.ArgumentParser(description="Script to generate ACTS material map")
0138     p.add_argument(
0139         "-o",
0140         "--outFile",
0141         type=str,
0142         default="material-map.json",
0143         help="Output filename for the generated material map. Supported formats: JSON, CBOR.",
0144     )
0145     args = p.parse_args()
0146     if ".json" in args.outFile:
0147         mapFormat = JsonFormat.Json
0148     elif ".cbor" in args.outFile:
0149         mapFormat = JsonFormat.Cbor
0150     else:
0151         print(
0152             "ERROR(material_mapping.py): please provide an output name ending with .json or .cbor"
0153         )
0154         exit()
0155
0156     mapName = args.outFile.split(".")[0]
0157
0158     detector = getOpenDataDetector(None)
0159     trackingGeometry = detector.trackingGeometry()
0160     decorators = detector.contextDecorators()
0161
0162     runMaterialMapping(
0163         trackingGeometry,
0164         decorators,
0165         outputDir=os.getcwd(),
0166         inputDir=os.getcwd(),
0167         readCachedSurfaceInformation=False,
0168         mapName=mapName,
0169         mapFormat=mapFormat,
0170     ).run()