Examples/tests/test_histogram_fit_backends.py

0001 """Equivalence test for two Gaussian resolution-fit backends: ROOT's
0002 `TH1::Fit` (via `ActsPlugins::RootHistogramFit`) and `acts.examples.scipy`'s
0003 `curve_fit`-based one. A synthetic `IAlgorithm` writes tracks/particles/hits
0004 straight to the whiteboard; the real `TrackTruthMatcher` then derives the
0005 matching, and two `PythonTrackParameterPerformanceWriter`s differing only in
0006 `fitFunction` see bit-identical histograms so only the fit itself can differ.
0007 """
0008
0009 import numpy as np
0010 import pytest
0011
0012 import acts
0013 import acts.examples
0014 import acts.examples.scipy as acts_scipy
0015 from acts.examples import PythonTrackParameterPerformanceWriter
0016
0017 u = acts.UnitConstants
0018
0019 pytestmark = pytest.mark.root
0020
0021
0022 class _SyntheticTrackAlgorithm(acts.examples.IAlgorithm):
0023     """Writes `nTracks` synthetic tracks/particles/measurement-particle-map
0024     entries to the whiteboard every event: truth d0 = 0, fitted d0 = a
0025     residual drawn from `sampler` (a callable `rng -> float`).
0026
0027     Deliberately does NOT write a TrackParticleMatching itself; instead it
0028     writes one source link per track plus a MeasurementParticlesMap entry, so
0029     the real TrackTruthMatcher (see `_run_backends`) derives the matching
0030     itself, as it would from real digitized hits.
0031     """
0032
0033     def __init__(self, sampler, nTracks, seed):
0034         super().__init__(name="SyntheticTrackAlgorithm", level=acts.logging.WARNING)
0035         self._sampler = sampler
0036         self._nTracks = nTracks
0037         self._rng = np.random.default_rng(seed)
0038
0039         self.outputTracks = acts.examples.WriteDataHandle(
0040             self, acts.examples.ConstTrackContainer, "OutputTracks"
0041         )
0042         self.outputTracks.initialize("tracks")
0043         self.outputParticles = acts.examples.WriteDataHandle(
0044             self, acts.examples.SimParticleContainer, "OutputParticles"
0045         )
0046         self.outputParticles.initialize("particles_selected")
0047         self.outputMeasurementParticlesMap = acts.examples.WriteDataHandle(
0048             self, acts.examples.MeasurementParticlesMap, "OutputMeasurementParticlesMap"
0049         )
0050         self.outputMeasurementParticlesMap.initialize("measurement_particles_map")
0051
0052     def execute(self, context):
0053         tc = acts.examples.TrackContainer()
0054         particles = acts.examples.SimParticleContainer()
0055         measurementParticlesMap = acts.examples.MeasurementParticlesMap()
0056
0057         surface = acts.Surface.createPerigee(acts.Vector3(0, 0, 0))
0058         # Identity covariance makes pull == residual (enough for this test).
0059         cov = acts.BoundMatrix.Identity()
0060         geoId = acts.GeometryIdentifier()
0061
0062         for i in range(self._nTracks):
0063             barcode = acts.examples.SimBarcode()
0064             barcode.particle = i
0065             particle = acts.examples.SimParticle(barcode, acts.PdgParticle.eMuon)
0066             # Transverse direction, matching the track parameters below --
0067             # avoids a degenerate intersection with the perigee surface.
0068             particle.direction = acts.Vector3(1, 0, 0)
0069             particle.absoluteMomentum = 1.0 * u.GeV
0070             particles.insert(particle)
0071
0072             # One measurement per track, own truth particle -> always Matched.
0073             hitIndex = i
0074             measurementParticlesMap.insert(hitIndex, barcode)
0075
0076             residual = self._sampler(self._rng)
0077             track = tc.makeTrack()
0078             track.referenceSurface = surface
0079             track.parameters = acts.BoundVector(residual, 0.0, 0.0, np.pi / 2, 1.0, 0.0)
0080             track.covariance = cov
0081             track.particleHypothesis = acts.ParticleHypothesis.muon
0082             track.nMeasurements = 1
0083
0084             state = track.appendTrackState()
0085             state.typeFlags.isMeasurement = True
0086             state.uncalibratedSourceLink = acts.examples.IndexSourceLink(
0087                 geoId, hitIndex
0088             ).toSourceLink()
0089
0090         self.outputTracks(context, tc.makeConst())
0091         self.outputParticles(context, particles)
0092         self.outputMeasurementParticlesMap(context, measurementParticlesMap)
0093
0094         return acts.examples.ProcessCode.SUCCESS
0095
0096
0097 def _small_res_plot_config():
0098     """Shrink Eta/Phi/Pt from their 40-bin defaults to 2 bins each -- every
0099     synthetic track lands in the same (eta, phi, pT) bin, so this just avoids
0100     fitting ~1600 empty slices per parameter for nothing.
0101     """
0102     cfg = acts.examples.ResPlotToolConfig()
0103     cfg.varBinning["Eta"] = acts.Axis.regular(2, -4.0, 4.0, "#eta")
0104     cfg.varBinning["Phi"] = acts.Axis.regular(2, -np.pi, np.pi, "#phi")
0105     cfg.varBinning["Pt"] = acts.Axis.regular(2, 0.0, 100.0, "pT [GeV/c]")
0106     return cfg
0107
0108
0109 def _run_backends(sampler, nTracks, seed):
0110     """Run the synthetic algorithm + the real TrackTruthMatcher once, score
0111     the result with both fit backends, and return `{backend: histogram_dict}`.
0112     """
0113     import acts.examples.root as acts_root
0114
0115     s = acts.examples.Sequencer(events=1, numThreads=1, logLevel=acts.logging.WARNING)
0116     s.addAlgorithm(_SyntheticTrackAlgorithm(sampler, nTracks, seed))
0117     s.addAlgorithm(
0118         acts.examples.TrackTruthMatcher(
0119             level=acts.logging.WARNING,
0120             config=acts.examples.TrackTruthMatcher.Config(
0121                 inputTracks="tracks",
0122                 inputParticles="particles_selected",
0123                 inputMeasurementParticlesMap="measurement_particles_map",
0124                 outputTrackParticleMatching="track_particle_matching",
0125                 outputParticleTrackMatching="particle_track_matching",
0126             ),
0127         )
0128     )
0129
0130     writers = {}
0131     for backend, fitFn in [
0132         ("root", acts_root.makeRootHistogramFitFunction()),
0133         ("scipy", acts_scipy.makeScipyHistogramFitFunction()),
0134     ]:
0135         cfg = acts.examples.PythonTrackParameterPerformanceWriter.Config(
0136             inputTracks="tracks",
0137             inputParticles="particles_selected",
0138             inputTrackParticleMatching="track_particle_matching",
0139             fitFunction=fitFn,
0140             resPlotToolConfig=_small_res_plot_config(),
0141         )
0142         writers[backend] = acts.examples.PythonTrackParameterPerformanceWriter(
0143             config=cfg, level=acts.logging.WARNING
0144         )
0145         s.addWriter(writers[backend])
0146
0147     s.run()
0148     # histograms() re-runs fitProfiles() on every call; cache it once.
0149     return {backend: w.histograms() for backend, w in writers.items()}
0150
0151
0152 def _fitted_bins(histograms, key, backend):
0153     """`(rootVals, otherVals, both)` for `key`, restricted to bins where both
0154     ROOT and `backend` succeeded (an unfitted bin has error == 0).
0155     """
0156     root = histograms["root"].get(key)
0157     other = histograms[backend].get(key)
0158     assert root is not None, f"ROOT produced no {key} (fit failed everywhere)"
0159     assert other is not None, f"{backend} produced no {key} (fit failed everywhere)"
0160
0161     rootVals = np.asarray(root.histogram.values())
0162     rootErrs = np.asarray(root.histogram.errors())
0163     otherVals = np.asarray(other.histogram.values())
0164     otherErrs = np.asarray(other.histogram.errors())
0165
0166     both = (rootErrs > 0) & (otherErrs > 0)
0167     assert (
0168         np.count_nonzero(both) >= 1
0169     ), f"no bin where both root and {backend} succeeded fitting {key}"
0170     return rootVals[both], otherVals[both], both
0171
0172
0173 def _assert_backend_agrees(histograms, key, backend, rtol, atol):
0174     rootVals, otherVals, _ = _fitted_bins(histograms, key, backend)
0175     np.testing.assert_allclose(
0176         otherVals,
0177         rootVals,
0178         rtol=rtol,
0179         atol=atol,
0180         err_msg=f"{backend} vs root disagree on {key}",
0181     )
0182
0183
0184 # Observed agreement is ~1e-5 to 1e-6 relative; 1e-3 leaves ample margin.
0185 # resmean also gets an absolute floor since it's close to zero.
0186 _RTOL = 1e-3
0187 _MEAN_ATOL = 1e-3
0188
0189 _SCENARIOS = {
0190     "gaussian": lambda rng: rng.normal(0.0, 0.02),
0191     "gaussian_with_outliers": lambda rng: (
0192         rng.uniform(-0.5, 0.5) if rng.uniform() < 0.02 else rng.normal(0.0, 0.02)
0193     ),
0194 }
0195
0196 # Fixed, not hash(scenario): Python randomizes string hashing per-process
0197 # (PYTHONHASHSEED), so seeding off hash() would make this test's sample
0198 # non-reproducible from run to run.
0199 _SEEDS = {"gaussian": 2, "gaussian_with_outliers": 3}
0200
0201
0202 @pytest.mark.parametrize("scenario", list(_SCENARIOS.keys()))
0203 def test_fit_backends_agree(scenario):
0204     histograms = _run_backends(
0205         _SCENARIOS[scenario], nTracks=5000, seed=_SEEDS[scenario]
0206     )
0207
0208     _assert_backend_agrees(
0209         histograms, "reswidth_d0_vs_eta", "scipy", rtol=_RTOL, atol=0.0
0210     )
0211     _assert_backend_agrees(
0212         histograms, "resmean_d0_vs_eta", "scipy", rtol=_RTOL, atol=_MEAN_ATOL
0213     )