Scripts/Python/material_mapping_itk.py

0001 #!/usr/bin/env python3
0002
0003 import os
0004 import argparse
0005 from pathlib import Path
0006
0007 import acts
0008 from acts import (
0009     SurfaceMaterialMapper,
0010     Navigator,
0011     Propagator,
0012     StraightLineStepper,
0013     MaterialMapJsonConverter,
0014 )
0015 from acts.examples import (
0016     Sequencer,
0017     WhiteBoard,
0018     AlgorithmContext,
0019     ProcessCode,
0020     MaterialMapping,
0021 )
0022
0023 from acts.examples.root import (
0024     RootMaterialTrackReader,
0025     RootMaterialTrackWriter,
0026 )
0027
0028 from acts.examples.json import (
0029     JsonMaterialWriter,
0030     JsonFormat,
0031 )
0032
0033
0034 def runMaterialMapping(
0035     trackingGeometry,
0036     decorators,
0037     outputDir,
0038     inputDir,
0039     materialStepsFile,
0040     mapName="material-map",
0041     mapSurface=True,
0042     readCachedSurfaceInformation=False,
0043     dumpMaterialTracks=False,
0044     s=None,
0045 ):
0046     s = s or Sequencer(numThreads=1)
0047
0048     for decorator in decorators:
0049         s.addContextDecorator(decorator)
0050
0051     wb = WhiteBoard(acts.logging.INFO)
0052
0053     context = AlgorithmContext(0, 0, wb)
0054
0055     for decorator in decorators:
0056         assert decorator.decorate(context) == ProcessCode.SUCCESS
0057
0058     # Read material step information from a ROOT TTRee
0059     s.addReader(
0060         RootMaterialTrackReader(
0061             level=acts.logging.INFO,
0062             outputMaterialTracks="material-tracks",
0063             fileList=[
0064                 os.path.join(
0065                     inputDir,
0066                     materialStepsFile,
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=False,
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     jmConverterCfg = MaterialMapJsonConverter.Config(
0091         processSensitives=False,
0092         processNonMaterial=False,
0093         processApproaches=True,
0094         processRepresenting=True,
0095         processBoundaries=True,
0096         processVolumes=True,
0097         context=context.geoContext,
0098     )
0099
0100     jmw = JsonMaterialWriter(
0101         level=acts.logging.VERBOSE,
0102         converterCfg=jmConverterCfg,
0103         fileName=os.path.join(outputDir, mapName),
0104         writeFormat=JsonFormat.Json,
0105     )
0106
0107     if dumpMaterialTracks:
0108         s.addWriter(
0109             RootMaterialTrackWriter(
0110                 level=acts.logging.INFO,
0111                 inputMaterialTracks=mmAlgCfg.collection,
0112                 filePath=os.path.join(
0113                     outputDir,
0114                     mapName + "_tracks.root",
0115                 ),
0116                 storeSurface=True,
0117                 storeVolume=True,
0118             )
0119         )
0120
0121     mmAlgCfg.materialWriters = [jmw]
0122
0123     s.addAlgorithm(MaterialMapping(level=acts.logging.INFO, config=mmAlgCfg))
0124
0125     return s
0126
0127
0128 if "__main__" == __name__:
0129     p = argparse.ArgumentParser(
0130         description="Script to run material mapping on ITk geometry"
0131     )
0132     p.add_argument(
0133         "geo_dir",
0134         help="Input directory containing the ITk standalone geometry. Get in touch if you don't have this.",
0135     )
0136     p.add_argument(
0137         "--inputFile",
0138         type=str,
0139         default="",
0140         help="Input file containing material steps.",
0141     )
0142     p.add_argument(
0143         "--material",
0144         type=str,
0145         default="",
0146         help="Geometry file to define layers used in material mapping",
0147     )
0148
0149     args = p.parse_args()
0150
0151     geo_example_dir = Path(args.geo_dir)
0152     assert geo_example_dir.exists(), "Detector example input directory missing"
0153     assert os.path.exists(
0154         args.inputFile
0155     ), "Invalid file in --inputFile. Please check your input!"
0156     assert os.path.exists(
0157         args.material
0158     ), "Invalid file path/name in --material. Please check your input!"
0159
0160     from acts.examples.itk import buildITkGeometry
0161
0162     detector = buildITkGeometry(geo_example_dir, customMaterialFile=args.material)
0163     trackingGeometry = detector.trackingGeometry()
0164     decorators = detector.contextDecorators()
0165
0166     runMaterialMapping(
0167         trackingGeometry,
0168         decorators,
0169         materialStepsFile=args.inputFile,
0170         outputDir=os.getcwd(),
0171         inputDir=os.getcwd(),
0172         readCachedSurfaceInformation=False,
0173     ).run()