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