physmon/workflows/physmon_ckf_tracking.py

0001 #!/usr/bin/env python3
0002
0003 import tempfile
0004 from pathlib import Path
0005 import shutil
0006
0007 import acts
0008 from acts.examples.simulation import (
0009     addParticleGun,
0010     MomentumConfig,
0011     EtaConfig,
0012     PhiConfig,
0013     ParticleConfig,
0014     addFatras,
0015     addDigitization,
0016 )
0017
0018 from acts.examples.reconstruction import (
0019     addSeeding,
0020     TruthSeedRanges,
0021     ParticleSmearingSigmas,
0022     SeedFinderConfigArg,
0023     SeedFinderOptionsArg,
0024     SeedingAlgorithm,
0025     TruthEstimatedSeedingAlgorithmConfigArg,
0026     CkfConfig,
0027     addCKFTracks,
0028     addAmbiguityResolution,
0029     AmbiguityResolutionConfig,
0030     addVertexFitting,
0031     VertexFinder,
0032     TrackSelectorConfig,
0033 )
0034
0035 from physmon_common import makeSetup
0036
0037 u = acts.UnitConstants
0038
0039 setup = makeSetup()
0040
0041
0042 def run_ckf_tracking(truthSmearedSeeded, truthEstimatedSeeded, label):
0043     with tempfile.TemporaryDirectory() as temp:
0044         s = acts.examples.Sequencer(
0045             events=500,
0046             numThreads=-1,
0047             logLevel=acts.logging.INFO,
0048         )
0049
0050         tp = Path(temp)
0051
0052         for d in setup.decorators:
0053             s.addContextDecorator(d)
0054
0055         rnd = acts.examples.RandomNumbers(seed=42)
0056
0057         addParticleGun(
0058             s,
0059             MomentumConfig(1.0 * u.GeV, 10.0 * u.GeV, transverse=True),
0060             EtaConfig(-3.0, 3.0),
0061             PhiConfig(0.0, 360.0 * u.degree),
0062             ParticleConfig(4, acts.PdgParticle.eMuon, randomizeCharge=True),
0063             vtxGen=acts.examples.GaussianVertexGenerator(
0064                 mean=acts.Vector4(0, 0, 0, 0),
0065                 stddev=acts.Vector4(
0066                     0.0125 * u.mm, 0.0125 * u.mm, 55.5 * u.mm, 1.0 * u.ns
0067                 ),
0068             ),
0069             multiplicity=50,
0070             rnd=rnd,
0071         )
0072
0073         addFatras(
0074             s,
0075             setup.trackingGeometry,
0076             setup.field,
0077             enableInteractions=True,
0078             rnd=rnd,
0079         )
0080
0081         addDigitization(
0082             s,
0083             setup.trackingGeometry,
0084             setup.field,
0085             digiConfigFile=setup.digiConfig,
0086             rnd=rnd,
0087         )
0088
0089         addSeeding(
0090             s,
0091             setup.trackingGeometry,
0092             setup.field,
0093             TruthSeedRanges(pt=(500 * u.MeV, None), nHits=(9, None)),
0094             ParticleSmearingSigmas(
0095                 pRel=0.01
0096             ),  # only used by SeedingAlgorithm.TruthSmeared
0097             SeedFinderConfigArg(
0098                 r=(33 * u.mm, 200 * u.mm),
0099                 deltaR=(1 * u.mm, 60 * u.mm),
0100                 collisionRegion=(-250 * u.mm, 250 * u.mm),
0101                 z=(-2000 * u.mm, 2000 * u.mm),
0102                 maxSeedsPerSpM=1,
0103                 sigmaScattering=5,
0104                 radLengthPerSeed=0.1,
0105                 minPt=500 * u.MeV,
0106                 impactMax=3 * u.mm,
0107             ),
0108             SeedFinderOptionsArg(bFieldInZ=2 * u.T),
0109             TruthEstimatedSeedingAlgorithmConfigArg(deltaR=(10.0 * u.mm, None)),
0110             seedingAlgorithm=SeedingAlgorithm.TruthSmeared
0111             if truthSmearedSeeded
0112             else SeedingAlgorithm.TruthEstimated
0113             if truthEstimatedSeeded
0114             else SeedingAlgorithm.Default
0115             if label == "seeded"
0116             else SeedingAlgorithm.Orthogonal,
0117             initialSigmas=[
0118                 1 * u.mm,
0119                 1 * u.mm,
0120                 1 * u.degree,
0121                 1 * u.degree,
0122                 0.1 / u.GeV,
0123                 1 * u.ns,
0124             ],
0125             initialVarInflation=[1.0] * 6,
0126             geoSelectionConfigFile=setup.geoSel,
0127             rnd=rnd,  # only used by SeedingAlgorithm.TruthSmeared
0128             outputDirRoot=tp,
0129         )
0130
0131         addCKFTracks(
0132             s,
0133             setup.trackingGeometry,
0134             setup.field,
0135             TrackSelectorConfig(
0136                 pt=(500 * u.MeV, None),
0137                 loc0=(-4.0 * u.mm, 4.0 * u.mm),
0138                 nMeasurementsMin=6,
0139                 maxHoles=2,
0140                 maxOutliers=2,
0141             ),
0142             CkfConfig(
0143                 seedDeduplication=False if truthSmearedSeeded else True,
0144                 stayOnSeed=False if truthSmearedSeeded else True,
0145             ),
0146             outputDirRoot=tp,
0147         )
0148
0149         if label in ["seeded", "orthogonal"]:
0150             addAmbiguityResolution(
0151                 s,
0152                 AmbiguityResolutionConfig(
0153                     maximumSharedHits=3,
0154                     maximumIterations=10000,
0155                     nMeasurementsMin=6,
0156                 ),
0157                 outputDirRoot=tp,
0158             )
0159
0160         s.addAlgorithm(
0161             acts.examples.TracksToParameters(
0162                 level=acts.logging.INFO,
0163                 inputTracks="tracks",
0164                 outputTrackParameters="trackParameters",
0165             )
0166         )
0167
0168         # Choosing a seeder only has an effect on VertexFinder.AMVF. For
0169         # VertexFinder.IVF we always use acts.VertexSeedFinder.GaussianSeeder
0170         # (Python binding is not implemented).
0171         # Setting useTime also only has an effect on VertexFinder.AMVF due to
0172         # the same reason.
0173         addVertexFitting(
0174             s,
0175             setup.field,
0176             trackParameters="trackParameters",
0177             outputProtoVertices="ivf_protovertices",
0178             outputVertices="ivf_fittedVertices",
0179             vertexFinder=VertexFinder.Iterative,
0180             outputDirRoot=tp / "ivf",
0181         )
0182
0183         addVertexFitting(
0184             s,
0185             setup.field,
0186             trackParameters="trackParameters",
0187             outputProtoVertices="amvf_protovertices",
0188             outputVertices="amvf_fittedVertices",
0189             seeder=acts.VertexSeedFinder.GaussianSeeder,
0190             useTime=False,  # Time seeding not implemented for the Gaussian seeder
0191             vertexFinder=VertexFinder.AMVF,
0192             outputDirRoot=tp / "amvf",
0193         )
0194
0195         # Use the adaptive grid vertex seeder in combination with the AMVF
0196         # To avoid having too many physmon cases, we only do this for the label
0197         # "seeded"
0198         if label == "seeded":
0199             addVertexFitting(
0200                 s,
0201                 setup.field,
0202                 trackParameters="trackParameters",
0203                 outputProtoVertices="amvf_gridseeder_protovertices",
0204                 outputVertices="amvf_gridseeder_fittedVertices",
0205                 seeder=acts.VertexSeedFinder.AdaptiveGridSeeder,
0206                 useTime=True,
0207                 vertexFinder=VertexFinder.AMVF,
0208                 outputDirRoot=tp / "amvf_gridseeder",
0209             )
0210
0211         s.run()
0212         del s
0213
0214         for vertexing in ["ivf", "amvf"]:
0215             shutil.move(
0216                 tp / f"{vertexing}/performance_vertexing.root",
0217                 tp / f"performance_{vertexing}.root",
0218             )
0219
0220         if label == "seeded":
0221             vertexing = "amvf_gridseeder"
0222             shutil.move(
0223                 tp / f"{vertexing}/performance_vertexing.root",
0224                 tp / f"performance_{vertexing}.root",
0225             )
0226
0227         for stem in (
0228             [
0229                 "performance_ckf",
0230                 "tracksummary_ckf",
0231                 "performance_ivf",
0232                 "performance_amvf",
0233             ]
0234             + (["performance_amvf_gridseeder"] if label == "seeded" else [])
0235             + (
0236                 ["performance_seeding", "performance_ambi"]
0237                 if label in ["seeded", "orthogonal"]
0238                 else ["performance_seeding"]
0239                 if label == "truth_estimated"
0240                 else []
0241             )
0242         ):
0243             perf_file = tp / f"{stem}.root"
0244             assert perf_file.exists(), "Performance file not found"
0245             shutil.copy(perf_file, setup.outdir / f"{stem}_{label}.root")
0246
0247
0248 for truthSmearedSeeded, truthEstimatedSeeded, label in [
0249     (True, False, "truth_smeared"),  # if first is true, second is ignored
0250     (False, True, "truth_estimated"),
0251     (False, False, "seeded"),
0252     (False, False, "orthogonal"),
0253 ]:
0254     run_ckf_tracking(truthSmearedSeeded, truthEstimatedSeeded, label)