Fix virtual particle species

src/diagrams/diagrams.jl

@@ -78,6 +78,11 @@ struct VirtualParticle{PROC<:QEDbase.AbstractProcessDefinition,PT<:QEDbase.Abstr
     out_particle_contributions::NTuple{O,Bool}
 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)...))")
+end
+
 import Base: +
 
 # "addition" of the bool tuples
@@ -89,16 +94,31 @@ function +(a::Tuple{NTuple{I,Bool},NTuple{O,Bool}}, b::Tuple{NTuple{I,Bool},NTup
     return (ntuple(i -> a[1][i] || b[1][i], I), ntuple(i -> a[2][i] || b[2][i], O))
 end
 
+invert(::Electron) = Positron()
+invert(::Positron) = Electron()
+invert(::Photon) = Photon()
+invert(::AbstractParticleType) = throw(InvalidInputError("unimplemented for this particle type"))
+
+function invert(virtual_particle::VirtualParticle)
+    I = length(virtual_particle.in_particle_contributions)
+    O = length(virtual_particle.out_particle_contributions)
+    return VirtualParticle(
+        virtual_particle.proc,
+        invert(virtual_particle.species),
+        ntuple(x -> !virtual_particle.in_particle_contributions[x], I),
+        ntuple(x -> !virtual_particle.out_particle_contributions[x], O))
+end
+
 # normalize the representation
-function normalize(virtual_particle::VirtualParticle{P,S,IN_T,OUT_T}) where {P,S,IN_T,OUT_T}
+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)
     s = sum(data[1]) + sum(data[2])
     if s > (I + O) / 2
-        return VirtualParticle(virtual_particle.proc, virtual_particle.species, ntuple(x -> !data[1][x], I), ntuple(x -> !data[2][x], O))
+        return invert(virtual_particle)
     elseif s == (I + O) / 2 && data[1][1] == false
-        return VirtualParticle(virtual_particle.proc, virtual_particle.species, ntuple(x -> !data[1][x], I), ntuple(x -> !data[2][x], O))
+        return invert(virtual_particle)
     else
         return virtual_particle
     end
@@ -124,16 +144,14 @@ 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 _momentum_contribution(proc::AbstractProcessDefinition, diagram::FeynmanDiagram{N,E,U,T,M,FM}, n::Int) where {N,E,U,T,M,FM}
+function _fermion_type(proc::AbstractProcessDefinition, ::FeynmanDiagram{N,E,U,T,M,FM}, n::Int) where {N,E,U,T,M,FM}
+    # 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)
-        # left electron n
         electron_n = n
         if electron_n > number_particles(proc, Incoming(), Electron())
-            # outgoing positron
+            return (Outgoing(), Positron(), electron_n - number_particles(proc, Incoming(), Electron()))
-            return _momentum_contribution(proc, Outgoing(), Positron(), electron_n - number_particles(proc, Incoming(), Electron()))
         else
-            # incoming electron
+            return (Incoming(), Electron(), electron_n)
-            return _momentum_contribution(proc, Incoming(), Electron(), electron_n)
         end
     elseif (n > E && n <= E + U)
         # left muon n - E
@@ -147,21 +165,19 @@ function _momentum_contribution(proc::AbstractProcessDefinition, diagram::Feynma
         # photon
         photon_n = n - N
         if photon_n > number_particles(proc, Incoming(), Photon())
-            # outgoing photon
+            return (Outgoing(), Photon(), photon_n - number_particles(proc, Incoming(), Photon()))
-            return _momentum_contribution(proc, Outgoing(), Photon(), photon_n - number_particles(proc, Incoming(), Photon()))
         else
-            # incoming photon
+            return (Incoming(), Photon(), photon_n)
-            return _momentum_contribution(proc, Incoming(), Photon(), photon_n)
         end
     elseif (n > N + M && n <= N + M + E)
         # right electron
         electron_n = n - N - M
         if electron_n > number_particles(proc, Outgoing(), Electron())
             # incoming positron
-            return _momentum_contribution(proc, Incoming(), Positron(), electron_n - number_particles(proc, Outgoing(), Electron()))
+            return (Incoming(), Positron(), electron_n - number_particles(proc, Outgoing(), Electron()))
         else
             # outgoing electron
-            return _momentum_contribution(proc, Outgoing(), Electron(), electron_n)
+            return (Outgoing(), Electron(), electron_n)
         end
     elseif (n > N + M + E && n <= N + M + E + U)
         # right muon
@@ -173,10 +189,14 @@ function _momentum_contribution(proc::AbstractProcessDefinition, diagram::Feynma
         throw(InvalidInputError("unimplemented for tauons"))
     else
         # error
-        throw(InvalidInputError("invalid index given for _momentum_contribution()"))
+        throw(InvalidInputError("invalid index given"))
     end
 end
 
+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}
     fermion_lines = PriorityQueue{Int64,Int64}()
 
@@ -206,8 +226,9 @@ function virtual_particles(proc::QEDbase.AbstractProcessDefinition, diagram::Fey
         current_line = dequeue!(fermion_lines)
 
         local unknown_photon_momentum = nothing
-        # walk line from the *left* (everything looks like an electron/muon/tauon)
+        # walk line from the *left*
-        species = current_line <= E ? Electron() : throw(InvalidInputError("muon/taun not implemented yet"))
+        species = _fermion_type(proc, diagram, current_line)[2]
+
         cumulative_mom = _momentum_contribution(proc, diagram, current_line)
 
         for i in 1:length(diagram.diagram_structure, current_line)
@@ -233,10 +254,9 @@ function virtual_particles(proc::QEDbase.AbstractProcessDefinition, diagram::Fey
             continue
         end
 
-        # walk line from the *right* (everything looks like a positron/antimuon/antitauon)
-        species = current_line <= E ? Positron() : throw(InvalidInputError("muon/taun not implemented yet"))
         # find right side of the line
         right_line = diagram.electron_permutation[current_line]
+        species = invert(species)
 
         cumulative_mom = _momentum_contribution(proc, diagram, right_line)
         for i in length(diagram.diagram_structure, current_line):-1:1   # iterate from the right