Improve particle_pair generation

src/diagrams/diagrams.jl

@@ -303,26 +303,33 @@ function are_total(a::VirtualParticle, b::VirtualParticle, c::VirtualParticle)
     return true
 end
 
-"""
-    pairs
+function particle_pairs(particles::Vector)
+    pairs = Dict{VirtualParticle,Vector{Tuple{VirtualParticle,VirtualParticle}}}()
 
-For a particle `vp` and a vector of `particles`, return a Vector of pairs of particles from the vector containing all pairs whose joint contributions are equal to the contributions of `vp`.
-"""
-function particle_pairs(vp::VirtualParticle, particles::Vector)
-    # momentum contributions are unique among particles, so each particle can only have at most one partner
-    result_pairs = Vector{Tuple{VirtualParticle,VirtualParticle}}()
+    proc = QEDbase.process(first(particles))
+    # make sure the "smallest" particles come first, i.e. those with few contributors
+    all_particles = sort([_pseudo_virtual_particles(proc)..., particles...])
 
-    # TODO: make this less awful
-    for i in eachindex(particles)
-        a = particles[i]
-        for b in particles[i+1:end]
-            if make_up(a, b, vp)
-                push!(result_pairs, (a, b))
+    # find pairs for every particle after the external ones (those can't have pairs)
+    for p_i in number_incoming_particles(proc)+number_outgoing_particles(proc)+1:length(all_particles)
+        p = all_particles[p_i]
+        pairs[p] = Vector{Tuple{VirtualParticle,VirtualParticle}}()
+
+        # only need to consider external particles and virtual particles that come before p_i
+        for p_a_i in 1:p_i-2
+            # and only partners between a and p_i
+            for p_b_i in p_a_i+1:p_i-1
+                p_a = all_particles[p_a_i]
+                p_b = all_particles[p_b_i]
+
+                if make_up(p_a, p_b, p)
+                    push!(pairs[p], (p_a, p_b))
+                end
             end
         end
     end
 
-    return result_pairs
+    return pairs
 end
 
 function total_particle_triples(particles::Vector)