Scripts/Python/geometry.py

0001 #!/usr/bin/env python3
0002
0003 import os
0004 import json
0005 from pathlib import Path
0006
0007 import acts
0008 import acts.examples
0009 from acts.json import MaterialMapJsonConverter, TrackingGeometryJsonConverter
0010 from acts.examples.odd import getOpenDataDetector
0011 from acts.examples import (
0012     WhiteBoard,
0013     AlgorithmContext,
0014     ProcessCode,
0015     CsvTrackingGeometryWriter,
0016     ObjTrackingGeometryWriter,
0017 )
0018
0019 from acts.examples.json import (
0020     JsonSurfacesWriter,
0021     JsonMaterialWriter,
0022     JsonFormat,
0023 )
0024
0025
0026 def runGeometry(
0027     trackingGeometry,
0028     decorators,
0029     outputDir: Path,
0030     events=1,
0031     pyVisProjection="xy",
0032     outputObj=True,
0033     outputPy=False,
0034     outputCsv=True,
0035     outputMaterialMap=True,
0036     outputSurfacesJson=True,
0037     serializeGeometryJson=False,
0038 ):
0039     for ievt in range(events):
0040         eventStore = WhiteBoard(name=f"EventStore#{ievt}", level=acts.logging.INFO)
0041         ialg = 0
0042         ithread = 0
0043
0044         context = AlgorithmContext(ialg, ievt, eventStore, ithread)
0045
0046         for cdr in decorators:
0047             r = cdr.decorate(context)
0048             if r != ProcessCode.SUCCESS:
0049                 raise RuntimeError("Failed to decorate event context")
0050
0051         if outputCsv:
0052             if not os.path.isdir(outputDir / "csv"):
0053                 os.makedirs(outputDir / "csv")
0054             writer = CsvTrackingGeometryWriter(
0055                 level=acts.logging.INFO,
0056                 trackingGeometry=trackingGeometry,
0057                 outputDir=str(outputDir / "csv"),
0058                 writePerEvent=True,
0059             )
0060             writer.write(context)
0061
0062         # The obj and material map outputs go to a single, event-independent
0063         # file each, so writing them once (on the first event) is enough --
0064         # every further event would only overwrite the same file.
0065         if outputObj and ievt == 0:
0066             vis = acts.ObjVisualization3D()
0067             trackingGeometry.visualize(vis, context.recoGeoContext)
0068             vis.write(outputDir / "obj" / "geometry.obj")
0069         if outputPy:
0070             from acts.examples.visualization import PyVisualization2D
0071
0072             vis = PyVisualization2D()
0073             trackingGeometry.visualize(vis, context.recoGeoContext)
0074
0075             vis.plot(projection=pyVisProjection, filename="geometry.pdf")
0076
0077         if outputSurfacesJson:
0078             # if not os.path.isdir(outputDir / "json"):
0079             #    os.makedirs(outputDir / "json")
0080             writer = JsonSurfacesWriter(
0081                 level=acts.logging.INFO,
0082                 trackingGeometry=trackingGeometry,
0083                 outputDir=str(outputDir / "json"),
0084                 writePerEvent=True,
0085                 writeSensitive=True,
0086             )
0087             writer.write(context)
0088
0089             if outputMaterialMap and ievt == 0:
0090                 jmConverterCfg = MaterialMapJsonConverter.Config(
0091                     processSensitives=True,
0092                     processApproaches=True,
0093                     processRepresenting=True,
0094                     processBoundaries=True,
0095                     processVolumes=True,
0096                     processNonMaterial=True,
0097                     context=context.recoGeoContext,
0098                 )
0099
0100                 jmw = JsonMaterialWriter(
0101                     level=acts.logging.VERBOSE,
0102                     converterCfg=jmConverterCfg,
0103                     fileName=str(outputDir / "geometry-map"),
0104                     writeFormat=JsonFormat.Json,
0105                 )
0106
0107                 jmw.write(trackingGeometry)
0108
0109         if serializeGeometryJson:
0110             converter = TrackingGeometryJsonConverter(level=acts.logging.INFO)
0111             jsonStr = converter.toJson(context.recoGeoContext, trackingGeometry)
0112             outPath = outputDir / "json" / "tracking-geometry.json"
0113             outPath.write_text(jsonStr)
0114
0115
0116 if "__main__" == __name__:
0117     # detector = acts.examples.GenericDetector()
0118     detector = getOpenDataDetector()
0119     trackingGeometry = detector.trackingGeometry()
0120     decorators = detector.contextDecorators()
0121
0122     runGeometry(trackingGeometry, decorators, pyVisProjection="rz")
0123     # Uncomment if you want to create the geometry id mapping for DD4hep
0124     # dd4hepIdGeoIdMap = acts.examples.dd4hep.createDD4hepIdGeoIdMap(trackingGeometry)
0125     # dd4hepIdGeoIdValueMap = {}
0126     # for key, value in dd4hepIdGeoIdMap.items():
0127     #     dd4hepIdGeoIdValueMap[key] = value.value
0128
0129     # with open('odd-dd4hep-geoid-mapping.json', 'w') as outfile:
0130     #    json.dump(dd4hepIdGeoIdValueMap, outfile)