Scripts/Python/material_mapping.py

0001 #!/usr/bin/env python3
0002
0003 import os
0004
0005 import acts
0006 from acts import (
0007     SurfaceMaterialMapper,
0008     VolumeMaterialMapper,
0009     Navigator,
0010     Propagator,
0011     StraightLineStepper,
0012     MaterialMapJsonConverter,
0013 )
0014 from acts.examples import (
0015     Sequencer,
0016     WhiteBoard,
0017     AlgorithmContext,
0018     ProcessCode,
0019     RootMaterialTrackReader,
0020     RootMaterialTrackWriter,
0021     MaterialMapping,
0022     JsonMaterialWriter,
0023     JsonFormat,
0024 )
0025 from acts.examples.odd import getOpenDataDetector
0026
0027
0028 def runMaterialMapping(
0029     trackingGeometry,
0030     decorators,
0031     outputDir,
0032     inputDir,
0033     mapName="material-map",
0034     mapSurface=True,
0035     mapVolume=True,
0036     readCachedSurfaceInformation=False,
0037     mappingStep=1,
0038     s=None,
0039 ):
0040     s = s or Sequencer(numThreads=1)
0041
0042     for decorator in decorators:
0043         s.addContextDecorator(decorator)
0044
0045     wb = WhiteBoard(acts.logging.INFO)
0046
0047     context = AlgorithmContext(0, 0, wb)
0048
0049     for decorator in decorators:
0050         assert decorator.decorate(context) == ProcessCode.SUCCESS
0051
0052     # Read material step information from a ROOT TTRee
0053     s.addReader(
0054         RootMaterialTrackReader(
0055             level=acts.logging.INFO,
0056             outputMaterialTracks="material-tracks",
0057             fileList=[
0058                 os.path.join(
0059                     inputDir,
0060                     mapName + "_tracks.root"
0061                     if readCachedSurfaceInformation
0062                     else "geant4_material_tracks.root",
0063                 )
0064             ],
0065             readCachedSurfaceInformation=readCachedSurfaceInformation,
0066         )
0067     )
0068
0069     stepper = StraightLineStepper()
0070
0071     mmAlgCfg = MaterialMapping.Config(context.geoContext, context.magFieldContext)
0072     mmAlgCfg.trackingGeometry = trackingGeometry
0073     mmAlgCfg.inputMaterialTracks = "material-tracks"
0074
0075     if mapSurface:
0076         navigator = Navigator(
0077             trackingGeometry=trackingGeometry,
0078             resolveSensitive=True,
0079             resolveMaterial=True,
0080             resolvePassive=True,
0081         )
0082         propagator = Propagator(stepper, navigator)
0083         mapper = SurfaceMaterialMapper(level=acts.logging.INFO, propagator=propagator)
0084         mmAlgCfg.materialSurfaceMapper = mapper
0085
0086     if mapVolume:
0087         navigator = Navigator(
0088             trackingGeometry=trackingGeometry,
0089         )
0090         propagator = Propagator(stepper, navigator)
0091         mapper = VolumeMaterialMapper(
0092             level=acts.logging.INFO, propagator=propagator, mappingStep=mappingStep
0093         )
0094         mmAlgCfg.materialVolumeMapper = mapper
0095
0096     jmConverterCfg = MaterialMapJsonConverter.Config(
0097         processSensitives=True,
0098         processApproaches=True,
0099         processRepresenting=True,
0100         processBoundaries=True,
0101         processVolumes=True,
0102         context=context.geoContext,
0103     )
0104
0105     jmw = JsonMaterialWriter(
0106         level=acts.logging.VERBOSE,
0107         converterCfg=jmConverterCfg,
0108         fileName=os.path.join(outputDir, mapName),
0109         writeFormat=JsonFormat.Json,
0110     )
0111
0112     mmAlgCfg.materialWriters = [jmw]
0113
0114     s.addAlgorithm(MaterialMapping(level=acts.logging.INFO, config=mmAlgCfg))
0115
0116     s.addWriter(
0117         RootMaterialTrackWriter(
0118             level=acts.logging.INFO,
0119             inputMaterialTracks=mmAlgCfg.mappingMaterialCollection,
0120             filePath=os.path.join(
0121                 outputDir,
0122                 mapName + "_tracks.root",
0123             ),
0124             storeSurface=True,
0125             storeVolume=True,
0126         )
0127     )
0128
0129     return s
0130
0131
0132 if "__main__" == __name__:
0133     matDeco = acts.IMaterialDecorator.fromFile("geometry-map.json")
0134     detector, trackingGeometry, decorators = getOpenDataDetector(matDeco)
0135
0136     runMaterialMapping(
0137         trackingGeometry,
0138         decorators,
0139         outputDir=os.getcwd(),
0140         inputDir=os.getcwd(),
0141         readCachedSurfaceInformation=False,
0142     ).run()