Implement convenience functions on the way to full DAG generation, test with new notebook

Project.toml

@@ -6,6 +6,7 @@ version = "0.1.0"
 [deps]
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
+MetagraphOptimization = "3e869610-d48d-4942-ba70-c1b702a33ca4"
 QEDbase = "10e22c08-3ccb-4172-bfcf-7d7aa3d04d93"
 QEDcore = "35dc0263-cb5f-4c33-a114-1d7f54ab753e"
 QEDprocesses = "46de9c38-1bb3-4547-a1ec-da24d767fdad"

src/FeynmanDiagramGenerator.jl

@@ -8,6 +8,7 @@ using Reexport
 include("QEDprocesses_patch.jl")
 
 import Base.==
+import QEDbase.process
 
 export FlatMatrix
 
@@ -15,7 +16,10 @@ export GenericQEDProcess, isphysical
 
 export AbstractTreeLevelFeynmanDiagram, FeynmanVertex, FeynmanDiagram
 export external_particles, virtual_particles, process, vertices
-export VirtualParticle
+
+export VirtualParticle, number_contributions
+export contributions, in_contributions, out_contributions
+export is_virtual, particle_pairs
 
 export Forest
 

src/diagrams/diagrams.jl

@@ -80,9 +80,23 @@ end
 
 function Base.show(io::IO, vp::VirtualParticle)
     pr = x -> x ? "1" : "0"
-    println(io, "$(vp.species): \t$(*(pr.(vp.in_particle_contributions)...)) | $(*(pr.(vp.out_particle_contributions)...))")
+    print(io, "$(vp.species): \t$(*(pr.(vp.in_particle_contributions)...)) | $(*(pr.(vp.out_particle_contributions)...))")
 end
 
+function QEDbase.process(vp::VirtualParticle)
+    return vp.proc
+end
+
+function QEDbase.particle_species(vp::VirtualParticle)
+    return vp.species
+end
+
+in_contributions(vp::VirtualParticle) = vp.in_particle_contributions
+out_contributions(vp::VirtualParticle) = vp.out_particle_contributions
+contributions(vp::VirtualParticle) = ((in_contributions(vp), out_contributions(vp)))
+
+is_virtual(vp::VirtualParticle) = number_contributions(vp) > 1
+
 import Base: +
 
 # "addition" of the bool tuples
@@ -111,9 +125,9 @@ end
 
 # normalize the representation
 function normalize(virtual_particle::VirtualParticle)
-    I = length(virtual_particle.in_particle_contributions)
-    O = length(virtual_particle.out_particle_contributions)
-    data = (virtual_particle.in_particle_contributions, virtual_particle.out_particle_contributions)
+    I = length(in_contributions(virtual_particle))
+    O = length(out_contributions(virtual_particle))
+    data = contributions(virtual_particle)
     s = sum(data[1]) + sum(data[2])
     if s > (I + O) / 2
         return invert(virtual_particle)
@@ -193,11 +207,141 @@ function _fermion_type(proc::AbstractProcessDefinition, ::FeynmanDiagram{N,E,U,T
     end
 end
 
-function _momentum_contribution(proc::AbstractProcessDefinition, diagram::FeynmanDiagram, n::Int)
+@inline function _momentum_contribution(proc::AbstractProcessDefinition, diagram::FeynmanDiagram, n::Int)
     return _momentum_contribution(proc, _fermion_type(proc, diagram, n)...)
 end
 
-function virtual_particles(proc::QEDbase.AbstractProcessDefinition, diagram::FeynmanDiagram{N,E,U,T,M,FM}) where {N,E,U,T,M,FM}
+function _external_particle(proc::AbstractProcessDefinition, diagram::FeynmanDiagram, n::Int)
+    (dir, species, _) = _fermion_type(proc, diagram, n)
+    if dir == Outgoing()
+        species = invert(species)
+    end
+    return VirtualParticle(proc, species, _momentum_contribution(proc, diagram, n)...)
+end
+
+function number_contributions(vp::VirtualParticle)
+    return sum(vp.in_particle_contributions) + sum(vp.out_particle_contributions)
+end
+
+function Base.isless(a::VirtualParticle, b::VirtualParticle)
+    if number_contributions(a) == number_contributions(b)
+        if a.in_particle_contributions == b.in_particle_contributions
+            return a.out_particle_contributions < b.out_particle_contributions
+        end
+        return a.in_particle_contributions < b.in_particle_contributions
+    end
+    return number_contributions(a) < number_contributions(b)
+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}
+    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
+
+"""
+    make_up(a::VirtualParticle, b::VirtualParticle, c::VirtualParticle)
+    
+For virtual particles `a`, `b`, and `c`, return true if `a` and `b`'s joint momentum contributions add up to `c`'s momentum contributions.
+"""
+function make_up(a::VirtualParticle, b::VirtualParticle, c::VirtualParticle)
+    if particle_species(a) == Photon() && particle_species(b) == Photon()
+        return false
+    end
+    # it should be unnecessary to check here that a and b can actually react. if a + b = c they must be able to if a, b and c all exist in the diagram.
+    for (a_contrib, b_contrib, c_contrib) in Iterators.zip(Iterators.flatten.(contributions.((a, b, c)))...)
+        if c_contrib != a_contrib + b_contrib
+            return false
+        end
+    end
+
+    return true
+end
+
+"""
+    are_total(a::VirtualParticle, b::VirtualParticle, c::VirtualParticle)
+
+Return true if a, b and c combined contain all external particles exactly once.
+"""
+function are_total(a::VirtualParticle, b::VirtualParticle, c::VirtualParticle)
+    for (a_contrib, b_contrib, c_contrib) in Iterators.zip(Iterators.flatten.(contributions.((a, b, c)))...)
+        if a_contrib + b_contrib + c_contrib != 1
+            return false
+        end
+    end
+
+    return true
+end
+
+"""
+    pairs
+
+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}}()
+
+    # 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))
+            end
+        end
+    end
+
+    return result_pairs
+end
+
+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(proto_particle)
+
+    working_set = vcat(particles, _pseudo_virtual_particles(proc, first(feynman_diagrams(proc))))
+
+    photons = filter(p -> particle_species(p) == Photon(), working_set)
+    electrons = filter(p -> particle_species(p) == Electron(), working_set)
+    positrons = filter(p -> particle_species(p) == Positron(), 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)
+        if are_total(ph, e, p)
+            push!(result_triples, (ph, e, p))
+        end
+    end
+
+    return result_triples
+end
+
+"""
+    _pseudo_virtual_particles
+
+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}
+    return sort(_external_particle.(proc, Ref(diagram), [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}
     fermion_lines = PriorityQueue{Int64,Int64}()
 
     # count number of internal photons in each fermion line and make a priority queue for fermion line => number of internal photons
@@ -226,8 +370,11 @@ function virtual_particles(proc::QEDbase.AbstractProcessDefinition, diagram::Fey
         current_line = dequeue!(fermion_lines)
 
         local unknown_photon_momentum = nothing
-        # walk line from the *left*
-        species = _fermion_type(proc, diagram, current_line)[2]
+        # walk line from the *left* (either incoming electron or outgoing positron)
+        (dir, species, _) = _fermion_type(proc, diagram, current_line)
+        if dir == Outgoing()
+            species = invert(species)
+        end
 
         cumulative_mom = _momentum_contribution(proc, diagram, current_line)
 
@@ -281,7 +428,7 @@ function virtual_particles(proc::QEDbase.AbstractProcessDefinition, diagram::Fey
         end
     end
 
-    return ntuple(x -> normalize(result[x]), length(result) - 1)
+    return normalize.(result)[1:end-1]
 end
 
 function vertices(diagram::FeynmanDiagram{N,E,U,T,M,FM}) where {N,E,U,T,M,FM}