Scripts/Python/itk_seeding.py

0001 #!/usr/bin/env python3
0002 import os
0003 import argparse
0004 import tempfile
0005
0006 import pathlib, acts
0007
0008 from acts.examples import (
0009     CsvSpacePointReader,
0010     TrackParamsEstimationAlgorithm,
0011     SeedingPerformanceWriter,
0012 )
0013 from acts.examples.reconstruction import (
0014     addStandardSeeding,
0015 )
0016
0017 from acts.examples.itk import itkSeedingAlgConfig, InputSpacePointsType
0018
0019 u = acts.UnitConstants
0020 rnd = acts.examples.RandomNumbers(seed=42)
0021
0022
0023 def runITkSeedingFromCsv(detector, trackingGeometry, field, outputDir):
0024     # create temporary file with pixel SPs and run the seeding
0025     with tempfile.TemporaryDirectory() as tmpdirname:
0026         temp = open(tmpdirname + "/event000000000-spacepoints_pixel.csv", "w+t")
0027         print(
0028             "created temporary file: "
0029             + tmpdirname
0030             + "/event000000000-spacepoints_pixel.csv"
0031         )
0032         temp.write(
0033             "measurement_id,sp_type,module_idhash,sp_x,sp_y,sp_z,sp_radius,sp_covr,sp_covz,sp_topHalfStripLength,sp_bottomHalfStripLength,sp_topStripDirection[0],sp_topStripDirection[1],sp_topStripDirection[2],sp_bottomStripDirection[0],sp_bottomStripDirection[1],sp_bottomStripDirection[2],sp_stripCenterDistance[0],sp_stripCenterDistance[1],sp_stripCenterDistance[2],sp_topStripCenterPosition[0],sp_topStripCenterPosition[1],sp_topStripCenterPosition[2]\n 1,0,3139,32.67557144165039,-5.311902523040771,-47.65000152587891,33.10452270507812,0.05999999865889549,0.02999880164861679,0,0,0,0,0,0,0,0,0,0,0,0,0,0\n 2,0,3422,95.14442443847656,-15.46361255645752,-52.125,96.39286804199219,0.05999999865889549,0.01687432639300823,0,0,0,0,0,0,0,0,0,0,0,0,0,0\n 3,0,3650,102.8257064819336,-16.71612739562988,-52.67499923706055,104.1755981445312,0.05999999865889549,0.001875000074505806,0,0,0,0,0,0,0,0,0,0,0,0,0,0\n 4,0,4223,159.4266204833984,-25.91166687011719,-56.75,161.5186157226562,0.05999999865889549,0.02999880164861679,0,0,0,0,0,0,0,0,0,0,0,0,0,0\n 5,0,5015,224.07958984375,-36.37123107910156,-61.40000152587891,227.0121765136719,0.05999999865889549,0.007499700412154198,0,0,0,0,0,0,0,0,0,0,0,0,0,0\n 6,0,6023,284.1485595703125,-46.0638542175293,-65.72499847412109,287.8580932617188,0.05999999865889549,0.001875000074505806,0,0,0,0,0,0,0,0,0,0,0,0,0,0"
0034         )
0035         temp.read()
0036
0037         s = acts.examples.Sequencer(events=1, numThreads=-1, logLevel=acts.logging.INFO)
0038
0039         # Read input space points from input csv files
0040         evReader = CsvSpacePointReader(
0041             level=acts.logging.INFO,
0042             inputStem="spacepoints",
0043             inputCollection="pixel",
0044             inputDir=os.path.dirname(temp.name),
0045             outputSpacePoints="PixelSpacePoint",
0046             extendCollection=False,
0047         )
0048
0049         # add csv reader
0050         s.addReader(evReader)
0051         spacePoints = evReader.config.outputSpacePoints
0052
0053         # run seeding
0054         inputSeeds = addStandardSeeding(
0055             s,
0056             spacePoints,
0057             *acts.examples.itk.itkSeedingAlgConfig(
0058                 InputSpacePointsType.PixelSpacePoints
0059             ),
0060         )
0061
0062         prototracks = "seed-prototracks"
0063         s.addAlgorithm(
0064             acts.examples.SeedsToPrototracks(
0065                 level=acts.logging.INFO,
0066                 inputSeeds=inputSeeds,
0067                 outputProtoTracks=prototracks,
0068             )
0069         )
0070
0071         # estimate seeding performance
0072         parEstimateAlg = TrackParamsEstimationAlgorithm(
0073             level=acts.logging.INFO,
0074             inputSeeds=inputSeeds,
0075             outputTrackParameters="estimatedparameters",
0076             trackingGeometry=trackingGeometry,
0077             magneticField=field,
0078         )
0079         s.addAlgorithm(parEstimateAlg)
0080
0081         s.run()
0082
0083     # create temporary file with strips SPs and run the seeding
0084     with tempfile.TemporaryDirectory() as tmpdirname:
0085         temp = open(tmpdirname + "/event000000000-spacepoints_strip.csv", "w+t")
0086         print(
0087             "created temporary file: "
0088             + tmpdirname
0089             + "/event000000000-spacepoints_strip.csv"
0090         )
0091         temp.write(
0092             "measurement_id,sp_type,module_idhash,sp_x,sp_y,sp_z,sp_radius,sp_covr,sp_covz,sp_topHalfStripLength,sp_bottomHalfStripLength,sp_topStripDirection[0],sp_topStripDirection[1],sp_topStripDirection[2],sp_bottomStripDirection[0],sp_bottomStripDirection[1],sp_bottomStripDirection[2],sp_stripCenterDistance[0],sp_stripCenterDistance[1],sp_stripCenterDistance[2],sp_topStripCenterPosition[0],sp_topStripCenterPosition[1],sp_topStripCenterPosition[2]\n 0,1,0,386.77178955078125,-62.579288482666015625,-72.66841888427734375,391.801727294921875,0.100000001490116119384765625,5.11999988555908203125,12.08999919891357421875,12.08999919891357421875,-0.00864744372665882110595703125,-0.02451671846210956573486328125,0.999662101268768310546875,0.00864744372665882110595703125,0.02451671846210956573486328125,0.999662101268768310546875,-6.43960094451904296875,1.04346692562103271484375,23.157070159912109375,386.6771240234375,-62.847682952880859375,-61.724697113037109375\n 1,1,0,543.9947509765625,-87.7279205322265625,-82.09113311767578125,551.02313232421875,0.100000001490116119384765625,5.11999988555908203125,12.08999919891357421875,12.08999919891357421875,-0.0073835677467286586761474609375,-0.024926505982875823974609375,0.999662101268768310546875,0.0073835677467286586761474609375,0.024926505982875823974609375,0.999662101268768310546875,-6.34883975982666015625,1.17108881473541259765625,-2.3926274776458740234375,544.0279541015625,-87.6157989501953125,-86.58742523193359375\n 2,1,0,544.00189208984375,-87.70365142822265625,-83.064239501953125,551.02630615234375,0.100000001490116119384765625,5.11999988555908203125,12.08999919891357421875,12.08999919891357421875,-0.0073835677467286586761474609375,-0.024926505982875823974609375,0.999662101268768310546875,0.011192028410732746124267578125,0.02346457540988922119140625,0.999662101268768310546875,-24.83704376220703125,4.08867168426513671875,-0.0274789035320281982421875,544.0279541015625,-87.6157989501953125,-86.58742523193359375\n 3,1,0,562.2547607421875,-90.650543212890625,-83.99970245361328125,569.5155029296875,0.100000001490116119384765625,5.11999988555908203125,12.08999919891357421875,12.08999919891357421875,-0.011192028410732746124267578125,-0.02346457540988922119140625,0.999662101268768310546875,0.0073835677467286586761474609375,0.024926505982875823974609375,0.999662101268768310546875,11.8808460235595703125,-1.85768926143646240234375,-0.0588833652436733245849609375,562.25762939453125,-90.6445770263671875,-84.2536773681640625\n 4,1,0,562.26605224609375,-90.62686920166015625,-85.00818634033203125,569.52288818359375,0.100000001490116119384765625,5.11999988555908203125,12.08999919891357421875,12.08999919891357421875,-0.011192028410732746124267578125,-0.02346457540988922119140625,0.999662101268768310546875,0.011192028410732746124267578125,0.02346457540988922119140625,0.999662101268768310546875,-6.60735797882080078125,1.05989348888397216796875,2.3062651157379150390625,562.25762939453125,-90.6445770263671875,-84.2536773681640625\n 5,1,0,750.5875244140625,-120.5648040771484375,-98.9385986328125,760.2088623046875,0.100000001490116119384765625,5.11999988555908203125,24.1799983978271484375,24.1799983978271484375,-0.009588238783180713653564453125,-0.02416430227458477020263671875,0.999662101268768310546875,0.009588238783180713653564453125,0.02416430227458477020263671875,0.999662101268768310546875,-6.041371822357177734375,2.1813886165618896484375,0.3670396506786346435546875,750.81573486328125,-119.98968505859375,-122.7309112548828125\n 6,1,0,979.34979248046875,-156.5580291748046875,-114.63397979736328125,991.78448486328125,0.100000001490116119384765625,5.11999988555908203125,24.1799983978271484375,24.1799983978271484375,-0.00824898667633533477783203125,-0.02465364150702953338623046875,0.999662101268768310546875,0.00824898667633533477783203125,0.02465364150702953338623046875,0.999662101268768310546875,-6.2955722808837890625,1.417438507080078125,-1.45441913604736328125,979.4241943359375,-156.335723876953125,-123.64752960205078125"
0093         )
0094         temp.read()
0095
0096         s = acts.examples.Sequencer(events=1, numThreads=-1, logLevel=acts.logging.INFO)
0097
0098         # Read input space points from input csv files
0099         evReader = CsvSpacePointReader(
0100             level=acts.logging.INFO,
0101             inputStem="spacepoints",
0102             inputCollection="strip",
0103             inputDir=os.path.dirname(temp.name),
0104             outputSpacePoints="StripSpacePoint",
0105             extendCollection=True,
0106         )
0107
0108         # add csv reader
0109         s.addReader(evReader)
0110         spacePoints = evReader.config.outputSpacePoints
0111
0112         # run seeding
0113         inputSeeds = addStandardSeeding(
0114             s,
0115             spacePoints,
0116             *acts.examples.itk.itkSeedingAlgConfig(
0117                 InputSpacePointsType.StripSpacePoints
0118             ),
0119         )
0120
0121         prototracks = "seed-prototracks"
0122         s.addAlgorithm(
0123             acts.examples.SeedsToPrototracks(
0124                 level=acts.logging.INFO,
0125                 inputSeeds=inputSeeds,
0126                 outputProtoTracks=prototracks,
0127             )
0128         )
0129
0130         # estimate seeding performance
0131         parEstimateAlg = TrackParamsEstimationAlgorithm(
0132             level=acts.logging.INFO,
0133             inputSeeds=inputSeeds,
0134             outputTrackParameters="estimatedparameters",
0135             trackingGeometry=trackingGeometry,
0136             magneticField=field,
0137         )
0138         s.addAlgorithm(parEstimateAlg)
0139
0140         s.run()
0141
0142
0143 if "__main__" == __name__:
0144     p = argparse.ArgumentParser(
0145         description="Example script to run ITk seed finding based on CSV spacepoints",
0146     )
0147
0148     geo_dir = pathlib.Path("acts-itk")
0149     outputDir = pathlib.Path.cwd() / "itk_output"
0150
0151     detector, trackingGeometry, decorators = acts.examples.itk.buildITkGeometry(geo_dir)
0152     field = acts.examples.MagneticFieldMapXyz(
0153         str(geo_dir / "bfield/ATLAS-BField-xyz.root")
0154     )
0155
0156     runITkSeedingFromCsv(detector, trackingGeometry, field, outputDir)