Improve triples generation

src/FeynmanDiagramGenerator.jl

@@ -15,7 +15,7 @@ export FlatMatrix
 export GenericQEDProcess, isphysical
 
 export AbstractTreeLevelFeynmanDiagram, FeynmanVertex, FeynmanDiagram
-export external_particles, virtual_particles, process, vertices
+export external_particles, virtual_particles, process
 
 export VirtualParticle, number_contributions
 export contributions, in_contributions, out_contributions

src/diagrams/diagrams.jl

@@ -158,7 +158,13 @@ function _momentum_contribution(proc::AbstractProcessDefinition, dir::ParticleDi
     @assert false "tried to get momentum contribution of $dir $species $index but it does not exist in $proc"
 end
 
-function _fermion_type(proc::AbstractProcessDefinition, ::FeynmanDiagram{N,E,U,T,M,FM}, n::Int) where {N,E,U,T,M,FM}
+function _fermion_type(proc::AbstractProcessDefinition, n::Int)
+    E = number_particles(proc, Incoming(), Electron()) + number_particles(proc, Outgoing(), Positron())
+    U = 0 # TODO add muons
+    T = 0 # TODO add tauons
+    M = number_particles(proc, Incoming(), Photon()) + number_particles(proc, Outgoing(), Photon())
+    N = E + U + T
+
     # from the fermion index, get (Direction, Species, n) tuple, where n means it's the nth particle of that dir and species
     if (n > 0 && n <= E)
         electron_n = n
@@ -207,16 +213,16 @@ function _fermion_type(proc::AbstractProcessDefinition, ::FeynmanDiagram{N,E,U,T
     end
 end
 
-@inline function _momentum_contribution(proc::AbstractProcessDefinition, diagram::FeynmanDiagram, n::Int)
+@inline function _momentum_contribution(proc::AbstractProcessDefinition, n::Int)
-    return _momentum_contribution(proc, _fermion_type(proc, diagram, n)...)
+    return _momentum_contribution(proc, _fermion_type(proc, n)...)
 end
 
-function _external_particle(proc::AbstractProcessDefinition, diagram::FeynmanDiagram, n::Int)
+function _external_particle(proc::AbstractProcessDefinition, n::Int)
-    (dir, species, _) = _fermion_type(proc, diagram, n)
+    (dir, species, _) = _fermion_type(proc, n)
     if dir == Outgoing()
         species = invert(species)
     end
-    return VirtualParticle(proc, species, _momentum_contribution(proc, diagram, n)...)
+    return VirtualParticle(proc, species, _momentum_contribution(proc, n)...)
 end
 
 function number_contributions(vp::VirtualParticle)
@@ -233,12 +239,27 @@ function Base.isless(a::VirtualParticle, b::VirtualParticle)
     return number_contributions(a) < number_contributions(b)
 end
 
+"""
+    disjunct(a::VirtualParticle, b::VirtualParticle)
+
+Return true if the momenta contributions of `a` and `b` are disjunct.
+"""
+function disjunct(a::VirtualParticle, b::VirtualParticle)
+    for (a_contrib, b_contrib) in Iterators.zip(Iterators.flatten.(contributions.((a, b)))...)
+        if b_contrib && a_contrib
+            return false
+        end
+    end
+
+    return true
+end
+
 """
     contains(a::VirtualParticle, b::VirtualParticle)
 
 Returns true if the set of particles contributing to `a` are contains the set of particles contributing to `b`.
 """
-function contains(a::VirtualParticle{P,S1,IN_T,OUT_T}, b::VirtualParticle{P,S2,IN_T,OUT_T}) where {P,S1,S2,IN_T,OUT_T}
+function contains(a::VirtualParticle, b::VirtualParticle)
     for (a_contrib, b_contrib) in Iterators.zip(Iterators.flatten.(contributions.((a, b)))...)
         if b_contrib && !a_contrib
             return false
@@ -308,24 +329,34 @@ function total_particle_triples(particles::Vector)
     # particle pairs making up the whole graph
     result_triples = Vector{Tuple{VirtualParticle,VirtualParticle,VirtualParticle}}()
 
-    proto_particle = first(particles)
+    proc = QEDbase.process(first(particles))
-    proc = QEDbase.process(proto_particle)
 
-    working_set = vcat(particles, _pseudo_virtual_particles(proc, first(feynman_diagrams(proc))))
+    working_set = vcat(particles, _pseudo_virtual_particles(proc))
 
-    photons = filter(p -> particle_species(p) == Photon(), working_set)
+    photons = filter(p -> is_boson(particle_species(p)), working_set)
-    electrons = filter(p -> particle_species(p) == Electron(), working_set)
+
-    positrons = filter(p -> particle_species(p) == Positron(), working_set)
+    # make electrons a set for fast deletion
+    electrons = Set(filter(p -> is_fermion(particle_species(p)) && is_particle(particle_species(p)), working_set))
+
+    # make positrons a set for fast lookup
+    positrons = Set(filter(p -> is_fermion(particle_species(p)) && is_anti_particle(particle_species(p)), working_set))
 
-    println("photons:   $photons")
-    println("electrons: $electrons")
-    println("positron:  $positrons")
-    println("i photons: $(invert.(photons))")
 
     # no participant can have more than half the external particles, so every possible particle is contained here
-    for (ph, e, p) in Iterators.product(photons, electrons, positrons)
+    # every photon has exactly one electron and positron partner
-        if are_total(ph, e, p)
+    for ph in photons
-            push!(result_triples, (ph, e, p))
+        for e in electrons
+            if !disjunct(ph, e)
+                continue
+            end
+
+            # create the only partner the ph and e could have together, then look for it in the actual positrons
+            expected_p = invert(VirtualParticle(proc, particle_species(e), (contributions(ph) + contributions(e))...))
+
+            if expected_p in positrons
+                @assert are_total(ph, e, expected_p)
+                push!(result_triples, (ph, e, expected_p))
+            end
         end
     end
 
@@ -337,8 +368,8 @@ end
 
 Return a vector of `VirtualParticle` for each external particle. These are not actually virtual particles, but can be helpful as entry points.
 """
-function _pseudo_virtual_particles(proc::AbstractProcessDefinition, diagram::FeynmanDiagram{N,E,U,T,M,FM}) where {N,E,U,T,M,FM}
+function _pseudo_virtual_particles(proc::AbstractProcessDefinition)
-    return sort(_external_particle.(proc, Ref(diagram), [1:number_incoming_particles(proc)+number_outgoing_particles(proc);]))
+    return sort(_external_particle.(proc, [1:number_incoming_particles(proc)+number_outgoing_particles(proc);]))
 end
 
 function virtual_particles(proc::AbstractProcessDefinition, diagram::FeynmanDiagram{N,E,U,T,M,FM}) where {N,E,U,T,M,FM}
@@ -371,12 +402,12 @@ function virtual_particles(proc::AbstractProcessDefinition, diagram::FeynmanDiag
 
         local unknown_photon_momentum = nothing
         # walk line from the *left* (either incoming electron or outgoing positron)
-        (dir, species, _) = _fermion_type(proc, diagram, current_line)
+        (dir, species, _) = _fermion_type(proc, current_line)
         if dir == Outgoing()
             species = invert(species)
         end
 
-        cumulative_mom = _momentum_contribution(proc, diagram, current_line)
+        cumulative_mom = _momentum_contribution(proc, current_line)
 
         for i in 1:length(diagram.diagram_structure, current_line)
             binding_particle = diagram.diagram_structure[current_line, i]
@@ -389,7 +420,7 @@ function virtual_particles(proc::AbstractProcessDefinition, diagram::FeynmanDiag
                     break
                 end
             else # binding_particle is an external photon
-                cumulative_mom += _momentum_contribution(proc, diagram, binding_particle)
+                cumulative_mom += _momentum_contribution(proc, binding_particle)
             end
             push!(result, VirtualParticle(proc, species, cumulative_mom...))
         end
@@ -405,7 +436,7 @@ function virtual_particles(proc::AbstractProcessDefinition, diagram::FeynmanDiag
         right_line = diagram.electron_permutation[current_line]
         species = invert(species)
 
-        cumulative_mom = _momentum_contribution(proc, diagram, right_line)
+        cumulative_mom = _momentum_contribution(proc, right_line)
         for i in length(diagram.diagram_structure, current_line):-1:1   # iterate from the right
             binding_particle = diagram.diagram_structure[current_line, i]
             if (binding_particle <= N) # binding_particle is an internal photon
@@ -422,7 +453,7 @@ function virtual_particles(proc::AbstractProcessDefinition, diagram::FeynmanDiag
                     break
                 end
             else # binding_particle is an external photon
-                cumulative_mom += _momentum_contribution(proc, diagram, binding_particle)
+                cumulative_mom += _momentum_contribution(proc, binding_particle)
             end
             push!(result, VirtualParticle(proc, species, cumulative_mom...))
         end
@@ -431,11 +462,6 @@ function virtual_particles(proc::AbstractProcessDefinition, diagram::FeynmanDiag
     return normalize.(result)[1:end-1]
 end
 
-function vertices(diagram::FeynmanDiagram{N,E,U,T,M,FM}) where {N,E,U,T,M,FM}
-
-    return NTuple{N,VertexType}()
-end
-
 
 #
 # Generate Feynman Diagrams

src/diagrams/interface.jl

@@ -53,13 +53,3 @@ Return a tuple of the incoming and outgoing particles (`QEDbase.AbstractParticle
 function external_particles(diagram::AbstractTreeLevelFeynmanDiagram)
     return (incoming_particles(process(diagram)), outgoing_particles(process(diagram)))
 end
-
-"""
-    vertices(::AbstractTreeLevelFeynmanDiagram)::NTuple{N, VertexType}
-
-Interface function that must be implemented for an instance of [`AbstractTreeLevelFeynmanDiagram`](@ref).
-
-Return an `NTuple` with N elements, where N is the number of vertices in this diagram. For tree-level Feynman diagrams, \$N = k - 2\$, where \$k\$ is the number of external particles.
-The elements of the `NTuple` are instances of [`AbstractFeynmanVertex`](@ref). If necessary, each of these can indicate its position in the diagram.
-"""
-function vertices end