From f6084065e5d45ddcebdf47ba60be304f7bf01318 Mon Sep 17 00:00:00 2001
From: AntonReinhard <anton.reinhard@proton.me>
Date: Tue, 9 Jul 2024 18:33:37 +0200
Subject: [PATCH] Small changes

---
 src/QEDprocesses_patch.jl  | 28 ++++++++++++++++++----------
 src/diagrams/diagrams.jl   |  7 ++++++-
 src/generic_process_def.jl |  4 ++--
 3 files changed, 26 insertions(+), 13 deletions(-)

diff --git a/src/QEDprocesses_patch.jl b/src/QEDprocesses_patch.jl
index f6e1378..6f7bc03 100644
--- a/src/QEDprocesses_patch.jl
+++ b/src/QEDprocesses_patch.jl
@@ -1,21 +1,29 @@
 # patch QEDprocesses
 # see issue https://github.com/QEDjl-project/QEDprocesses.jl/issues/77
-@inline function QEDprocesses.number_particles(
-    proc_def::QEDbase.AbstractProcessDefinition, dir::DIR, ::PT
-) where {DIR<:QEDbase.ParticleDirection,PT<:QEDbase.AbstractParticleType}
-    return count(x -> x isa PT, particles(proc_def, dir))
+
+
+function _number_particles_helper(
+    particles::Tuple{}, species::SPECIES, n::Val{N}
+) where {SPECIES,N}
+    return N
 end
 
+function _number_particles_helper(
+    particles::Tuple{SPECIES,Vararg}, species::SPECIES2, n::Val{N}
+) where {SPECIES,SPECIES2,N}
+    return _number_particles_helper(particles[2:end], species, Val(N))
+end
 
-"""
-    number_particles(proc_def::AbstractProcessDefinition, dir::ParticleDirection, species::AbstractParticleType)
+function _number_particles_helper(
+    particles::Tuple{SPECIES,Vararg}, species::SPECIES, n::Val{N}
+) where {SPECIES,N}
+    return _number_particles_helper(particles[2:end], species, Val(N + 1))
+end
 
-Return the number of particles of the given direction and species in the given process definition.
-"""
-@inline function QEDbase.number_particles(
+@inline function QEDprocesses.number_particles(
     proc_def::AbstractProcessDefinition, dir::DIR, species::PT
 ) where {DIR<:ParticleDirection,PT<:AbstractParticleType}
-    return count(x -> x isa PT, particles(proc_def, dir))
+    return _number_particles_helper(particles(proc_def, dir), species, Val(0))
 end
 
 """
diff --git a/src/diagrams/diagrams.jl b/src/diagrams/diagrams.jl
index 8b124fa..5e0e746 100644
--- a/src/diagrams/diagrams.jl
+++ b/src/diagrams/diagrams.jl
@@ -105,6 +105,9 @@ function normalize(virtual_particle::VirtualParticle{P,S,IN_T,OUT_T}) where {P,S
 end
 
 function _momentum_contribution(proc::AbstractProcessDefinition, dir::ParticleDirection, species::AbstractParticleType, index::Int)
+    I = number_incoming_particles(proc)
+    O = number_outgoing_particles(proc)
+
     # get index of n-th "dir species" particle in proc
     particles_seen = 0
     c = 0
@@ -114,9 +117,11 @@ function _momentum_contribution(proc::AbstractProcessDefinition, dir::ParticleDi
             particles_seen += 1
         end
         if particles_seen == index
-            return (ntuple(x -> is_incoming(dir) && x == c, number_incoming_particles(proc)), ntuple(x -> is_outgoing(dir) && x == c, number_outgoing_particles(proc)))
+            return (((is_incoming(dir) && x == c for x in 1:I)...,), ((is_outgoing(dir) && x == c for x in 1:O)...,))
         end
     end
+
+    @assert false "tried to get momentum contribution of $dir $species $index but it does not exist in $proc"
 end
 
 function _momentum_contribution(proc::AbstractProcessDefinition, diagram::FeynmanDiagram{N,E,U,T,M,FM}, n::Int) where {N,E,U,T,M,FM}
diff --git a/src/generic_process_def.jl b/src/generic_process_def.jl
index 71628da..7690888 100644
--- a/src/generic_process_def.jl
+++ b/src/generic_process_def.jl
@@ -25,8 +25,8 @@ struct GenericQEDProcess{INT,OUTT} <: QEDbase.AbstractProcessDefinition where {I
     end
 end
 
-QEDprocesses.incoming_particles(proc::GenericQEDProcess) = proc.incoming_particles
-QEDprocesses.outgoing_particles(proc::GenericQEDProcess) = proc.outgoing_particles
+QEDprocesses.incoming_particles(proc::GenericQEDProcess{INT,OUTT}) where {INT,OUTT} = proc.incoming_particles
+QEDprocesses.outgoing_particles(proc::GenericQEDProcess{INT,OUTT}) where {INT,OUTT} = proc.outgoing_particles
 
 function isphysical(proc::GenericQEDProcess)
     return (number_particles(proc, Incoming(), Electron()) + number_particles(proc, Outgoing(), Positron()) ==