Add convenient accessor functions for phase space point

QEDjl-project · May 6, 2024 · aa50ca9 · aa50ca9
1 parent b9b4369
commit aa50ca9
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 16 deletions.
diff --git a/src/QEDprocesses.jl b/src/QEDprocesses.jl
@@ -31,7 +31,7 @@ export AbstractCoordinateSystem, SphericalCoordinateSystem
 export AbstractFrameOfReference, CenterOfMomentumFrame, ElectronRestFrame
 export AbstractPhasespaceDefinition, PhasespaceDefinition
 export ParticleStateful, PhaseSpacePoint
-export spin, pol
+export spin, pol, nth_momentum, getindex
 
 using QEDbase
 

diff --git a/src/phase_spaces.jl b/src/phase_spaces.jl
@@ -142,6 +142,23 @@ struct ParticleStateful{ElType<:AbstractFourMomentum} <: AbstractParticle
     end
 end
 
+@inline is_incoming(particle::ParticleStateful) = is_incoming(particle.dir)
+@inline is_outgoing(particle::ParticleStateful) = is_outgoing(particle.dir)
+@inline is_fermion(particle::ParticleStateful) = is_fermion(particle.species)
+@inline is_boson(particle::ParticleStateful) = is_boson(particle.species)
+@inline is_particle(particle::ParticleStateful) = is_particle(particle.species)
+@inline is_anti_particle(particle::ParticleStateful) = is_anti_particle(particle.species)
+@inline mass(particle::ParticleStateful) = mass(particle.species)
+@inline charge(particle::ParticleStateful) = charge(particle.species)
+
+@inline _spin(::Species, particle::ParticleStateful) where {Species<:FermionLike} =
+    particle.spin_or_pol
+@inline spin(particle::ParticleStateful) = _spin(particle.species, particle)
+
+@inline _pol(::Species, particle::ParticleStateful) where {Species<:BosonLike} =
+    particle.spin_or_pol
+@inline pol(particle::ParticleStateful) = _pol(particle.species, particle)
+
 """
     PhaseSpacePoint
 
@@ -216,19 +233,29 @@ struct PhaseSpacePoint{
     end
 end
 
-@inline is_incoming(particle::ParticleStateful) = is_incoming(particle.dir)
-@inline is_outgoing(particle::ParticleStateful) = is_outgoing(particle.dir)
-@inline is_fermion(particle::ParticleStateful) = is_fermion(particle.species)
-@inline is_boson(particle::ParticleStateful) = is_boson(particle.species)
-@inline is_particle(particle::ParticleStateful) = is_particle(particle.species)
-@inline is_anti_particle(particle::ParticleStateful) = is_anti_particle(particle.species)
-@inline mass(particle::ParticleStateful) = mass(particle.species)
-@inline charge(particle::ParticleStateful) = charge(particle.species)
+"""
+    Base.getindex(psp::PhaseSpacePoint, dir::Incoming, n::Int)
 
-@inline _spin(::Species, particle::ParticleStateful) where {Species<:FermionLike} =
-    particle.spin_or_pol
-@inline spin(particle::ParticleStateful) = _spin(particle.species, particle)
+Overload for the array indexing operator `[]`. Returns the nth incoming particle in this phase space point.
+"""
+function Base.getindex(psp::PhaseSpacePoint, ::Incoming, n::Int)
+    return psp.in_particles[n]
+end
 
-@inline _pol(::Species, particle::ParticleStateful) where {Species<:BosonLike} =
-    particle.spin_or_pol
-@inline pol(particle::ParticleStateful) = _pol(particle.species, particle)
+"""
+    Base.getindex(psp::PhaseSpacePoint, dir::Outgoing, n::Int)
+
+Overload for the array indexing operator `[]`. Returns the nth outgoing particle in this phase space point.
+"""
+function Base.getindex(psp::PhaseSpacePoint, ::Outgoing, n::Int)
+    return psp.out_particles[n]
+end
+
+"""
+    nth_momentum(psp::PhaseSpacePoint, dir::ParticleDirection, n::Int)
+
+Returns the momentum of the `n`th particle in the given [`PhaseSpacePoint`](@ref) which has direction `dir`. If `n` is outside the valid range for this phase space point, an [`BoundsError`](@ref) is thrown.
+"""
+function nth_momentum(psp::PhaseSpacePoint, dir::ParticleDirection, n::Int)
+    return psp[dir, n].mom
+end
diff --git a/test/phase_spaces.jl b/test/phase_spaces.jl
@@ -72,7 +72,19 @@ end
     )
     phasespace_def = TESTPSDEF
 
-    PhaseSpacePoint(process, model, phasespace_def, in_particles_valid, out_particles_valid)
+    psp = PhaseSpacePoint(
+        process, model, phasespace_def, in_particles_valid, out_particles_valid
+    )
+
+    @test nth_momentum(psp, Incoming(), 1) == in_el.mom
+    @test nth_momentum(psp, Incoming(), 2) == in_ph.mom
+    @test nth_momentum(psp, Outgoing(), 1) == out_el.mom
+    @test nth_momentum(psp, Outgoing(), 2) == out_ph.mom
+
+    @test psp[Incoming(), 1] == in_el
+    @test psp[Incoming(), 2] == in_ph
+    @test psp[Outgoing(), 1] == out_el
+    @test psp[Outgoing(), 2] == out_ph
 
     if (VERSION >= v"1.8")
         # julia versions before 1.8 did not have support for regex matching in @test_throws
@@ -93,6 +105,16 @@ end
         )
     end
 
+    @test_throws BoundsError nth_momentum(psp, Incoming(), -1)
+    @test_throws BoundsError nth_momentum(psp, Outgoing(), -1)
+    @test_throws BoundsError nth_momentum(psp, Incoming(), 4)
+    @test_throws BoundsError nth_momentum(psp, Outgoing(), 4)
+
+    @test_throws BoundsError psp[Incoming(), -1]
+    @test_throws BoundsError psp[Outgoing(), -1]
+    @test_throws BoundsError psp[Incoming(), 4]
+    @test_throws BoundsError psp[Outgoing(), 4]
+
     @test_throws InvalidInputError PhaseSpacePoint(
         process, model, phasespace_def, in_particles_invalid, out_particles_valid
     )