Examples/tests/test_data_augmentation.py

0001 from pathlib import Path
0002 from typing import Optional
0003
0004 from pathlib import Path
0005 from typing import Optional
0006
0007 import acts
0008 import acts.examples
0009 import pytest
0010 import uproot as ur
0011 import awkward as ak
0012 import numpy as np
0013
0014 u = acts.UnitConstants
0015
0016
0017 def runTruthTracking(
0018     trackingGeometry: acts.TrackingGeometry,
0019     field: acts.MagneticFieldProvider,
0020     digiConfigFile: Path,
0021     outputDir: Path,
0022     inputParticlePath: Optional[Path] = None,
0023     inputSimHitsPath: Optional[Path] = None,
0024     decorators=[],
0025     s: acts.examples.Sequencer = None,
0026     n_events: int = 10,
0027 ):
0028     from acts.examples.simulation import (
0029         addParticleGun,
0030         ParticleConfig,
0031         EtaConfig,
0032         PhiConfig,
0033         MomentumConfig,
0034         addFatras,
0035         addDigitization,
0036         ParticleSelectorConfig,
0037         addDigiParticleSelection,
0038     )
0039     from acts.examples.reconstruction import (
0040         addSeeding,
0041         SeedingAlgorithm,
0042         TrackSmearingSigmas,
0043         addTruthTrackingGsf,
0044     )
0045     from acts.examples.root import (
0046         RootParticleReader,
0047         RootSimHitReader,
0048         RootTrackSummaryWriter,
0049     )
0050
0051     s = s or acts.examples.Sequencer(
0052         events=n_events, numThreads=-1, logLevel=acts.logging.INFO
0053     )
0054
0055     for d in decorators:
0056         s.addContextDecorator(d)
0057
0058     rnd = acts.examples.RandomNumbers(seed=42)
0059     outputDir = Path(outputDir)
0060     logger = acts.getDefaultLogger("GSF Example", acts.logging.INFO)
0061
0062     if inputParticlePath is None:
0063         addParticleGun(
0064             s,
0065             ParticleConfig(num=1, pdg=acts.PdgParticle.eElectron, randomizeCharge=True),
0066             EtaConfig(-3.0, 3.0, uniform=True),
0067             MomentumConfig(1.0 * u.GeV, 100.0 * u.GeV, transverse=True),
0068             PhiConfig(0.0, 360.0 * u.degree),
0069             vtxGen=acts.examples.GaussianVertexGenerator(
0070                 mean=acts.Vector4(0, 0, 0, 0),
0071                 stddev=acts.Vector4(0.015, 0.015, 55.0, 0),
0072             ),
0073             multiplicity=1,
0074             rnd=rnd,
0075             outputDirRoot=outputDir,
0076         )
0077     else:
0078         logger.info("Reading particles from {}", inputParticlePath.resolve())
0079         assert inputParticlePath.exists()
0080         s.addReader(
0081             RootParticleReader(
0082                 level=acts.logging.INFO,
0083                 filePath=str(inputParticlePath.resolve()),
0084                 outputParticles="particles_generated",
0085             )
0086         )
0087         s.addWhiteboardAlias("particles", "particles_generated")
0088
0089     if inputSimHitsPath is None:
0090         addFatras(
0091             s,
0092             trackingGeometry,
0093             field,
0094             rnd=rnd,
0095             enableInteractions=True,
0096         )
0097     else:
0098         logger.info("Reading hits from {}", inputSimHitsPath.resolve())
0099         s.addReader(
0100             RootSimHitReader(
0101                 level=acts.logging.INFO,
0102                 filePath=str(inputSimHitsPath.resolve()),
0103                 outputSimHits="simhits",
0104             )
0105         )
0106         s.addWhiteboardAlias("particles_simulated_selected", "particles_generated")
0107
0108     addDigitization(
0109         s,
0110         trackingGeometry,
0111         field,
0112         digiConfigFile=digiConfigFile,
0113         rnd=rnd,
0114     )
0115
0116     addDigiParticleSelection(
0117         s,
0118         ParticleSelectorConfig(
0119             pt=(0.9 * u.GeV, None),
0120             measurements=(7, None),
0121             removeNeutral=True,
0122             removeSecondaries=True,
0123         ),
0124     )
0125
0126     addSeeding(
0127         s,
0128         trackingGeometry,
0129         field,
0130         rnd=rnd,
0131         inputParticles="particles_generated",
0132         seedingAlgorithm=SeedingAlgorithm.TruthSmeared,
0133         trackSmearingSigmas=TrackSmearingSigmas(
0134             # zero everything so the GSF has a chance to find the measurements
0135             loc0=0,
0136             loc0PtA=0,
0137             loc0PtB=0,
0138             loc1=0,
0139             loc1PtA=0,
0140             loc1PtB=0,
0141             time=0,
0142             phi=0,
0143             theta=0,
0144             ptRel=0,
0145         ),
0146         particleHypothesis=acts.ParticleHypothesis.electron,
0147         initialSigmas=[
0148             1 * u.mm,
0149             1 * u.mm,
0150             1 * u.degree,
0151             1 * u.degree,
0152             0 / u.GeV,
0153             1 * u.ns,
0154         ],
0155         initialSigmaQoverPt=0.1 / u.GeV,
0156         initialSigmaPtRel=0.1,
0157         initialVarInflation=[1e0, 1e0, 1e0, 1e0, 1e0, 1e0],
0158     )
0159
0160     addTruthTrackingGsf(
0161         s,
0162         trackingGeometry,
0163         field,
0164     )
0165
0166     s.addAlgorithm(
0167         acts.examples.TrackSelectorAlgorithm(
0168             level=acts.logging.INFO,
0169             inputTracks="tracks",
0170             outputTracks="selected-tracks",
0171             selectorConfig=acts.TrackSelector.Config(
0172                 minMeasurements=7,
0173             ),
0174         )
0175     )
0176     s.addWhiteboardAlias("tracks", "selected-tracks")
0177
0178     s.addWriter(
0179         RootTrackSummaryWriter(
0180             level=acts.logging.INFO,
0181             inputTracks="tracks",
0182             inputParticles="particles_selected",
0183             inputTrackParticleMatching="track_particle_matching",
0184             filePath=str(outputDir / "tracksummary.root"),
0185             writeGsfSpecific=True,
0186         )
0187     )
0188
0189     return s
0190
0191
0192 def readExampleRootData(outputDir: Path):
0193     outputDir = Path(outputDir)
0194     summary_tree = ur.open(outputDir / "tracksummary.root")["tracksummary"]
0195     particles_tree = ur.open(outputDir / "particles.root")["particles"]
0196     summary_fields = ["t_d0", "t_z0", "t_phi", "t_theta"]
0197     summary = summary_tree.arrays(summary_fields, library="ak")
0198     masks = []
0199     for field in summary_fields:
0200         nonempty = ak.to_numpy(ak.num(summary[field]) > 0)
0201
0202         firsts = ak.firsts(summary[field])
0203         arr = ak.to_numpy(firsts)
0204         arr = np.asarray(arr).squeeze()
0205         # finite: numeric finite values (will be False for NaN/Inf)
0206         finite = np.isfinite(arr)
0207
0208         combined_filter = np.logical_and(nonempty, finite)
0209         masks.append(combined_filter)
0210
0211     particle_fields = ["vx", "vy", "vz", "px", "py", "pz", "q"]
0212     particles = particles_tree.arrays(particle_fields, library="ak")
0213     for field in particle_fields:
0214         nonempty = ak.to_numpy(ak.num(particles[field]) > 0)
0215
0216         firsts = ak.firsts(particles[field])
0217         arr = ak.to_numpy(firsts)
0218         arr = np.asarray(arr).squeeze()
0219         # finite: numeric finite values (will be False for NaN/Inf)
0220         finite = np.isfinite(arr)
0221         combined = np.logical_and(nonempty, finite)
0222         masks.append(combined)
0223
0224     combined_mask = np.logical_and.reduce(masks)
0225     simulation_data = {
0226         "d0": ak.to_numpy(summary["t_d0"][combined_mask][:, 0]),
0227         "z0": ak.to_numpy(summary["t_z0"][combined_mask][:, 0]),
0228         "phi": ak.to_numpy(summary["t_phi"][combined_mask][:, 0]),
0229         "theta": ak.to_numpy(summary["t_theta"][combined_mask][:, 0]),
0230         "vx": ak.to_numpy(particles["vx"][combined_mask][:, 0]),
0231         "vy": ak.to_numpy(particles["vy"][combined_mask][:, 0]),
0232         "vz": ak.to_numpy(particles["vz"][combined_mask][:, 0]),
0233         "px": ak.to_numpy(particles["px"][combined_mask][:, 0]),
0234         "py": ak.to_numpy(particles["py"][combined_mask][:, 0]),
0235         "pz": ak.to_numpy(particles["pz"][combined_mask][:, 0]),
0236         "q": ak.to_numpy(particles["q"][combined_mask][:, 0]),
0237     }
0238     return simulation_data
0239
0240
0241 def sampleUniformPointsIn3D(x_offset: float, y_offset: float, z_offset: float, n: int):
0242     if x_offset == 0 and y_offset == 0 and z_offset == 0:
0243         return np.zeros((n, 3))
0244     semi_axes = np.array([x_offset, y_offset, z_offset])
0245     points = np.empty((0, 3))
0246     # Expected acceptance rate of a candidate is pi/6 (volume of ellipsoid over
0247     # bounding box), so oversample accordingly to reduce the number of
0248     # rejection loop iterations.
0249     acceptance_rate = np.pi / 6
0250     while len(points) < n:
0251         remaining = n - len(points)
0252         n_candidates = int(remaining / acceptance_rate) + 10
0253         candidates = np.random.uniform(
0254             low=-semi_axes, high=semi_axes, size=(n_candidates, 3)
0255         )
0256         mask = np.sum((candidates / semi_axes) ** 2, axis=-1) <= 1
0257         points = np.vstack([points, candidates[mask]])
0258     return points[:n]
0259
0260
0261 def vacuumTrackParameterAndBeamspotPropagation(
0262     beamspots,
0263     truth_params,
0264 ):
0265     geo_context = acts.GeometryContext.dangerouslyDefaultConstruct()
0266     mag_field_context = acts.MagneticFieldContext()
0267     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0268     propagator = acts.EigenVoidPropagator(
0269         acts.EigenStepper(field), acts.VoidNavigator()
0270     )
0271     propagator_options = acts.PropagatorPlainOptions(geo_context, mag_field_context)
0272     # The particle hypothesis only affects the mass, which does not enter the
0273     # vacuum propagation used here (no material interactions), so a fixed pion
0274     # hypothesis is used regardless of the true particle species.
0275     particle_hypothesis = acts.ParticleHypothesis(211, 0.13957 * u.GeV, 1.0)
0276
0277     n = truth_params.shape[0]
0278     beamspot_pocas = np.empty((n, 5))
0279     for i, truth_params_set in enumerate(truth_params):
0280         beamspot = beamspots[i]
0281         vtx = truth_params_set[:3]
0282         mom = truth_params_set[3:6]
0283         q = truth_params_set[6]
0284
0285         pos4 = acts.Vector4(vtx[0], vtx[1], vtx[2], 0.0)
0286         qOverP = q / np.linalg.norm(mom)
0287         start = acts.BoundTrackParameters.createCurvilinear(
0288             pos4, acts.Vector3(*mom), qOverP, None, particle_hypothesis
0289         )
0290         target = acts.Surface.createPerigee(acts.Vector3(beamspot[0], beamspot[1], 0.0))
0291         result = propagator.propagateToSurface(start, target, propagator_options)
0292         beamspot_pocas[i] = np.array(result.parameters)[:5]
0293     return beamspot_pocas
0294
0295
0296 @pytest.mark.odd
0297 def test_track_propagation_returns_values(odd_detector_config):
0298     field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
0299
0300     outputDir = Path.cwd()
0301     n = 100
0302     with odd_detector_config.detector:
0303         runTruthTracking(
0304             trackingGeometry=odd_detector_config.trackingGeometry,
0305             field=field,
0306             digiConfigFile=odd_detector_config.digiConfigFile,
0307             outputDir=outputDir,
0308             decorators=odd_detector_config.decorators,
0309             n_events=n,
0310         ).run()
0311
0312     simulation_data = readExampleRootData(outputDir)
0313     truth_params = np.column_stack(
0314         [
0315             simulation_data["vx"],
0316             simulation_data["vy"],
0317             simulation_data["vz"],
0318             simulation_data["px"],
0319             simulation_data["py"],
0320             simulation_data["pz"],
0321             simulation_data["q"],
0322         ]
0323     )
0324
0325     x_offset = 0.1
0326     y_offset = 1.1
0327     z_offset = 10
0328     beamspots = sampleUniformPointsIn3D(x_offset, y_offset, z_offset, len(truth_params))
0329
0330     augmented_perigees = vacuumTrackParameterAndBeamspotPropagation(
0331         beamspots,
0332         truth_params,
0333     )
0334     assert augmented_perigees.shape == (len(truth_params), 5)
0335     assert np.all(np.isfinite(augmented_perigees))