docs/examples/test_custom_algorithm.py

0001 def test_custom_algorithm():
0002     from acts.examples.simulation import (
0003         addParticleGun,
0004         MomentumConfig,
0005         EtaConfig,
0006         ParticleConfig,
0007     )
0008
0009     #! [Custom algorithm filtering particles by pT]
0010     import acts
0011     from acts import UnitConstants as u
0012     import acts.examples
0013
0014     class HighPtParticleFilter(acts.examples.IAlgorithm):
0015         def __init__(self, ptMin, level):
0016             acts.examples.IAlgorithm.__init__(self, "HighPtParticleFilter", level)
0017             self.ptMin = ptMin
0018             self.inputParticles = acts.examples.ReadDataHandle(
0019                 self, acts.examples.SimParticleContainer, "InputParticles"
0020             )
0021             self.inputParticles.initialize("particles_generated")
0022             self.outputParticles = acts.examples.WriteDataHandle(
0023                 self, acts.examples.SimParticleContainer, "OutputParticles"
0024             )
0025             self.outputParticles.initialize("particles_high_pt")
0026
0027         def execute(self, context):
0028             particles = self.inputParticles(context.eventStore)
0029
0030             kept = acts.examples.SimParticleContainer()
0031             for particle in particles:
0032                 if particle.transverseMomentum >= self.ptMin:
0033                     kept.insert(particle)
0034
0035             self.outputParticles(context, kept)
0036             return acts.examples.ProcessCode.SUCCESS
0037
0038     #! [Custom algorithm filtering particles by pT]
0039
0040     s = acts.examples.Sequencer(events=5, numThreads=1)
0041     rnd = acts.examples.RandomNumbers(seed=42)
0042
0043     addParticleGun(
0044         s,
0045         MomentumConfig(0.1 * u.GeV, 10 * u.GeV, transverse=True),
0046         EtaConfig(-2.0, 2.0),
0047         ParticleConfig(20, acts.PdgParticle.eMuon, randomizeCharge=True),
0048         rnd=rnd,
0049     )
0050
0051     s.addAlgorithm(HighPtParticleFilter(ptMin=1 * u.GeV, level=acts.logging.INFO))
0052
0053     s.run()