Add number_particles overloads (#90)

QEDjl-project · Jul 9, 2024 · 6c227ce · 6c227ce
1 parent c150e90
commit 6c227ce
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 3 deletions.
diff --git a/src/QEDbase.jl b/src/QEDbase.jl
@@ -106,8 +106,8 @@ include("particles/direction.jl")
 include("particles/spin_pol.jl")
 
 include("interfaces/phase_space.jl")
-include("interfaces/process.jl")
 include("interfaces/particle_stateful.jl")
+include("interfaces/process.jl")
 include("interfaces/phase_space_point.jl")
 
 include("cross_section/diff_probability.jl")

diff --git a/src/interfaces/process.jl b/src/interfaces/process.jl
@@ -191,15 +191,37 @@ Convenience function dispatching to [`incoming_particles`](@ref) or [`outgoing_p
     outgoing_particles(proc_def)
 
 """
-    number_particles(proc_def::AbstractProcessDefinition, ::ParticleDirection)
+    number_particles(proc_def::AbstractProcessDefinition, dir::ParticleDirection)
 
-Convenience function dispatching to [`number_incoming_particles`](@ref) or [`number_outgoing_particles`](@ref) depending on the given direction.
+Convenience function dispatching to [`number_incoming_particles`](@ref) or [`number_outgoing_particles`](@ref) depending on the given direction, returning the number of incoming or outgoing particles, respectively.
 """
 @inline number_particles(proc_def::AbstractProcessDefinition, ::Incoming) =
     number_incoming_particles(proc_def)
 @inline number_particles(proc_def::AbstractProcessDefinition, ::Outgoing) =
     number_outgoing_particles(proc_def)
 
+"""
+    number_particles(proc_def::AbstractProcessDefinition, dir::ParticleDirection, species::AbstractParticleType)
+
+Return the number of particles of the given direction and species in the given process definition.
+"""
+@inline function number_particles(
+    proc_def::AbstractProcessDefinition, dir::DIR, species::PT
+) where {DIR<:ParticleDirection,PT<:AbstractParticleType}
+    return count(x -> x isa PT, particles(proc_def, dir))
+end
+
+"""
+    number_particles(proc_def::AbstractProcessDefinition, particle::AbstractParticleStateful)
+
+Return the number of particles of the given particle's direction and species in the given process definition.
+"""
+@inline function number_particles(
+    proc_def::AbstractProcessDefinition, ps::AbstractParticleStateful
+)
+    return number_particles(proc_def, particle_direction(ps), particle_species(ps))
+end
+
 #####
 # Generation of four-momenta from coordinates
 #