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     RootMaterialTrackReader,
0021     RootMaterialTrackWriter,
0022     MaterialMapping,
0023     JsonMaterialWriter,
0024     JsonFormat,
0025 )
0026
0027
0028 def runMaterialMapping(
0029     trackingGeometry,
0030     decorators,
0031     outputDir,
0032     inputDir,
0033     materialStepsFile,
0034     mapName="material-map",
0035     mapSurface=True,
0036     readCachedSurfaceInformation=False,
0037     dumpMaterialTracks=False,
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                     materialStepsFile,
0061                 )
0062             ],
0063             readCachedSurfaceInformation=readCachedSurfaceInformation,
0064         )
0065     )
0066
0067     stepper = StraightLineStepper()
0068
0069     mmAlgCfg = MaterialMapping.Config(context.geoContext, context.magFieldContext)
0070     mmAlgCfg.trackingGeometry = trackingGeometry
0071     mmAlgCfg.inputMaterialTracks = "material-tracks"
0072
0073     if mapSurface:
0074         navigator = Navigator(
0075             trackingGeometry=trackingGeometry,
0076             resolveSensitive=False,
0077             resolveMaterial=True,
0078             resolvePassive=True,
0079         )
0080         propagator = Propagator(stepper, navigator)
0081         mapper = SurfaceMaterialMapper(level=acts.logging.INFO, propagator=propagator)
0082         mmAlgCfg.materialSurfaceMapper = mapper
0083
0084     jmConverterCfg = MaterialMapJsonConverter.Config(
0085         processSensitives=False,
0086         processNonMaterial=False,
0087         processApproaches=True,
0088         processRepresenting=True,
0089         processBoundaries=True,
0090         processVolumes=True,
0091         context=context.geoContext,
0092     )
0093
0094     jmw = JsonMaterialWriter(
0095         level=acts.logging.VERBOSE,
0096         converterCfg=jmConverterCfg,
0097         fileName=os.path.join(outputDir, mapName),
0098         writeFormat=JsonFormat.Json,
0099     )
0100
0101     if dumpMaterialTracks:
0102         s.addWriter(
0103             RootMaterialTrackWriter(
0104                 level=acts.logging.INFO,
0105                 inputMaterialTracks=mmAlgCfg.collection,
0106                 filePath=os.path.join(
0107                     outputDir,
0108                     mapName + "_tracks.root",
0109                 ),
0110                 storeSurface=True,
0111                 storeVolume=True,
0112             )
0113         )
0114
0115     mmAlgCfg.materialWriters = [jmw]
0116
0117     s.addAlgorithm(MaterialMapping(level=acts.logging.INFO, config=mmAlgCfg))
0118
0119     return s
0120
0121
0122 if "__main__" == __name__:
0123     p = argparse.ArgumentParser(
0124         description="Script to run material mapping on ITk geometry"
0125     )
0126     p.add_argument(
0127         "geo_dir",
0128         help="Input directory containing the ITk standalone geometry. Get in touch if you don't have this.",
0129     )
0130     p.add_argument(
0131         "--inputFile",
0132         type=str,
0133         default="",
0134         help="Input file containing material steps.",
0135     )
0136     p.add_argument(
0137         "--material",
0138         type=str,
0139         default="",
0140         help="Geometry file to define layers used in material mapping",
0141     )
0142
0143     args = p.parse_args()
0144
0145     geo_example_dir = Path(args.geo_dir)
0146     assert geo_example_dir.exists(), "Detector example input directory missing"
0147     assert os.path.exists(
0148         args.inputFile
0149     ), "Invalid file in --inputFile. Please check your input!"
0150     assert os.path.exists(
0151         args.material
0152     ), "Invalid file path/name in --material. Please check your input!"
0153
0154     from acts.examples.itk import buildITkGeometry
0155
0156     detector, trackingGeometry, decorators = buildITkGeometry(
0157         geo_example_dir, customMaterialFile=args.material
0158     )
0159
0160     runMaterialMapping(
0161         trackingGeometry,
0162         decorators,
0163         materialStepsFile=args.inputFile,
0164         outputDir=os.getcwd(),
0165         inputDir=os.getcwd(),
0166         readCachedSurfaceInformation=False,
0167     ).run()