Scripts/Python/material_mapping_core.py

0001 #!/usr/bin/env python3
0002
0003 import argparse
0004
0005 import acts
0006 from acts import (
0007     MaterialMapper,
0008     IntersectionMaterialAssigner,
0009     BinnedSurfaceMaterialAccumulater,
0010     MaterialMapJsonConverter,
0011 )
0012
0013 from acts.examples import (
0014     Sequencer,
0015     WhiteBoard,
0016     AlgorithmContext,
0017     RootMaterialTrackReader,
0018     RootMaterialTrackWriter,
0019     CoreMaterialMapping,
0020     JsonMaterialWriter,
0021     RootMaterialWriter,
0022     JsonFormat,
0023 )
0024
0025 from acts.examples.dd4hep import (
0026     DD4hepDetector,
0027     DD4hepDetectorOptions,
0028     DD4hepGeometryService,
0029 )
0030
0031 from acts.examples.odd import getOpenDataDetector, getOpenDataDetectorDirectory
0032
0033
0034 def runMaterialMapping(surfaces, inputFile, outputFile, outputMap, loglevel):
0035     # Create a sequencer
0036     print("Creating the sequencer with 1 thread (inter event information needed)")
0037
0038     s = Sequencer(numThreads=1)
0039
0040     # IO for material tracks reading
0041     wb = WhiteBoard(acts.logging.INFO)
0042
0043     # Read material step information from a ROOT TTRee
0044     s.addReader(
0045         RootMaterialTrackReader(
0046             level=acts.logging.INFO,
0047             outputMaterialTracks="material-tracks",
0048             fileList=[inputFile],
0049             readCachedSurfaceInformation=False,
0050         )
0051     )
0052
0053     # Assignment setup : Intersection assigner
0054     materialAssingerConfig = IntersectionMaterialAssigner.Config()
0055     materialAssingerConfig.surfaces = surfaces
0056     materialAssinger = IntersectionMaterialAssigner(materialAssingerConfig, loglevel)
0057
0058     # Accumulation setup : Binned surface material accumulater
0059     materialAccumulaterConfig = BinnedSurfaceMaterialAccumulater.Config()
0060     materialAccumulaterConfig.materialSurfaces = surfaces
0061     materialAccumulater = BinnedSurfaceMaterialAccumulater(
0062         materialAccumulaterConfig, loglevel
0063     )
0064
0065     # Mapper setup
0066     materialMapperConfig = MaterialMapper.Config()
0067     materialMapperConfig.assignmentFinder = materialAssinger
0068     materialMapperConfig.surfaceMaterialAccumulater = materialAccumulater
0069     materialMapper = MaterialMapper(materialMapperConfig, loglevel)
0070
0071     # Add the map writer(s)
0072     mapWriters = []
0073     # json map writer
0074     context = AlgorithmContext(0, 0, wb)
0075     jmConverterCfg = MaterialMapJsonConverter.Config(
0076         processSensitives=True,
0077         processApproaches=True,
0078         processRepresenting=True,
0079         processBoundaries=True,
0080         processVolumes=True,
0081         context=context.geoContext,
0082     )
0083     mapWriters.append(
0084         JsonMaterialWriter(
0085             level=loglevel,
0086             converterCfg=jmConverterCfg,
0087             fileName=outputMap + "",
0088             writeFormat=JsonFormat.Json,
0089         )
0090     )
0091     mapWriters.append(RootMaterialWriter(level=loglevel, filePath=outputMap + ".root"))
0092
0093     # Mapping Algorithm
0094     coreMaterialMappingConfig = CoreMaterialMapping.Config()
0095     coreMaterialMappingConfig.materialMapper = materialMapper
0096     coreMaterialMappingConfig.inputMaterialTracks = "material-tracks"
0097     coreMaterialMappingConfig.mappedMaterialTracks = "mapped-material-tracks"
0098     coreMaterialMappingConfig.unmappedMaterialTracks = "unmapped-material-tracks"
0099     coreMaterialMappingConfig.materiaMaplWriters = mapWriters
0100     coreMaterialMapping = CoreMaterialMapping(coreMaterialMappingConfig, loglevel)
0101     s.addAlgorithm(coreMaterialMapping)
0102
0103     # Add the mapped material tracks writer
0104     s.addWriter(
0105         RootMaterialTrackWriter(
0106             level=acts.logging.INFO,
0107             inputMaterialTracks="mapped-material-tracks",
0108             filePath=outputFile + "_mapped.root",
0109             storeSurface=True,
0110             storeVolume=True,
0111         )
0112     )
0113
0114     # Add the unmapped material tracks writer
0115     s.addWriter(
0116         RootMaterialTrackWriter(
0117             level=acts.logging.INFO,
0118             inputMaterialTracks="unmapped-material-tracks",
0119             filePath=outputFile + "_unmapped.root",
0120             storeSurface=True,
0121             storeVolume=True,
0122         )
0123     )
0124
0125     return s
0126
0127
0128 if "__main__" == __name__:
0129     p = argparse.ArgumentParser()
0130
0131     p.add_argument(
0132         "-n", "--events", type=int, default=1000, help="Number of events to process"
0133     )
0134     p.add_argument(
0135         "-i", "--input", type=str, default="", help="Input file with material tracks"
0136     )
0137     p.add_argument(
0138         "-o", "--output", type=str, default="", help="Output file (core) name"
0139     )
0140     p.add_argument(
0141         "-m", "--map", type=str, default="", help="Output file for the material map"
0142     )
0143
0144     p.add_argument(
0145         "--matconfig", type=str, default="", help="Material configuration file"
0146     )
0147
0148     p.add_argument(
0149         "--experimental",
0150         action=argparse.BooleanOptionalAction,
0151         help="Construct experimental geometry",
0152     )
0153
0154     args = p.parse_args()
0155
0156     if args.experimental:
0157         odd_xml = getOpenDataDetectorDirectory() / "xml" / "OpenDataDetector.xml"
0158
0159         # Create the dd4hep geometry service and detector
0160         dd4hepConfig = DD4hepGeometryService.Config()
0161         dd4hepConfig.logLevel = acts.logging.INFO
0162         dd4hepConfig.xmlFileNames = [str(odd_xml)]
0163         dd4hepGeometryService = DD4hepGeometryService(dd4hepConfig)
0164         dd4hepDetector = DD4hepDetector(dd4hepGeometryService)
0165
0166         cOptions = DD4hepDetectorOptions(logLevel=acts.logging.INFO, emulateToGraph="")
0167
0168         # Context and options
0169         geoContext = acts.GeometryContext()
0170         [detector, contextors, store] = dd4hepDetector.finalize(geoContext, cOptions)
0171
0172         materialSurfaces = detector.extractMaterialSurfaces()
0173         print("Extracted number of material surfaces: ", len(materialSurfaces))
0174
0175     else:
0176         matDeco = None
0177         if args.matconfig != "":
0178             matDeco = acts.IMaterialDecorator.fromFile(args.matconfig)
0179
0180         [detector, trackingGeometry, decorators] = getOpenDataDetector(matDeco)
0181
0182         materialSurfaces = trackingGeometry.extractMaterialSurfaces()
0183         print("Extracted number of material surfaces: ", len(materialSurfaces))
0184
0185     runMaterialMapping(
0186         materialSurfaces, args.input, args.output, args.map, acts.logging.INFO
0187     ).run()