Add feynman_diagrams() iterator

src/FeynmanDiagramGenerator.jl

@@ -10,7 +10,7 @@ export external_particles, virtual_particles, process, vertices
 
 export Forest
 
-export plane_trees, labelled_plane_trees
+export plane_trees, labelled_plane_trees, feynman_diagrams
 
 export can_interact
 export QED

src/diagrams/diagrams.jl

@@ -58,3 +58,136 @@ function vertices(::AbstractTreeLevelFeynmanDiagram)
 
     return NTuple{N,VertexType}()
 end
+
+
+#
+# Generate Feynman Diagrams
+#
+
+mutable struct ExternalPhotonIterator
+    N::Int  # number of fermion lines
+    M::Int  # number of external photons
+    fermion_structures::Vector{Vector{Vector{Int}}}
+    fermion_structure_index::Int
+    photon_structures::Vector{Vector{Vector{Int}}}
+    photon_structure_index::Int
+end
+
+function _feynman_structures(n::Int, m::Int)
+    f = labelled_plane_trees(n)
+    return ExternalPhotonIterator(
+        n,
+        m,
+        f,
+        1,
+        _external_photon_multiplicity(f[1], n, m),
+        1
+    )
+end
+
+function Base.length(it::ExternalPhotonIterator)
+    return factorial(it.M + 3 * it.N - 3, 2 * it.N - 1)
+end
+
+function Base.iterate(iter::ExternalPhotonIterator)
+    return (iter.photon_structures[iter.photon_structure_index], nothing)
+end
+
+function Base.iterate(iter::ExternalPhotonIterator, n::Nothing)
+    if iter.photon_structure_index == length(iter.photon_structures)
+        if iter.fermion_structure_index == length(iter.fermion_structures)
+            return nothing
+        end
+
+        iter.fermion_structure_index += 1
+        iter.photon_structure_index = 1
+        iter.photon_structures = _external_photon_multiplicity(iter.fermion_structures[iter.fermion_structure_index], iter.N, iter.M)
+    else
+        iter.photon_structure_index += 1
+    end
+
+    return (iter.photon_structures[iter.photon_structure_index], nothing)
+end
+
+function _external_photon_multiplicity(fermion_structure::Vector{Vector{Int}}, n::Int, m::Int)
+    if m == 0
+        return [fermion_structure]
+    end
+
+    res = Vector{Vector{Vector{Int}}}()
+
+    # go through lines
+    for line_index in eachindex(fermion_structure)
+
+        # go through indices start to end
+        for index in 1:length(fermion_structure[line_index])+1
+            # copy to prevent muting
+            new_photon_structure = deepcopy(fermion_structure)
+            # add new photon
+            insert!(new_photon_structure[line_index], index, n + 1)
+            # recurse
+            append!(res, _external_photon_multiplicity(new_photon_structure, n + 1, m - 1))
+        end
+    end
+
+    return res
+end
+
+mutable struct FeynmanDiagramIterator
+    E::Int  # number of electron lines  (indices 1      - e)
+    e_perms::Vector{Vector{Int}} # list of all the possible permutations of the electrons
+    e_index::Int
+    U::Int  # number of muon lines      (indices e+1    - e+u)
+    u_perms::Vector{Vector{Int}}
+    u_index::Int
+    T::Int  # number of tauon lines     (indices e+u+1  - e+u+t)
+    t_perms::Vector{Vector{Int}}
+    t_index::Int
+    M::Int  # number of external photons
+    photon_structure_iter::ExternalPhotonIterator
+    photon_structure::Vector{Vector{Int}} # current structure that's being permuted
+end
+
+function Base.length(it::FeynmanDiagramIterator)
+    N = it.E + it.U + it.T
+    return factorial(it.M + 3 * N - 3, 2 * N - 1) * factorial(it.E) * factorial(it.U) * factorial(it.T)
+end
+
+function Base.iterate(iter::FeynmanDiagramIterator)
+    return ((iter.photon_structure, iter.e_perms[iter.e_index], iter.u_perms[iter.u_index], iter.t_perms[iter.t_index]), nothing)
+end
+
+function Base.iterate(iter::FeynmanDiagramIterator, ::Nothing)
+    iter.t_index += 1
+
+    if iter.t_index > length(iter.t_perms)
+        iter.t_index = 1
+        iter.u_index += 1
+    end
+
+    if iter.u_index > length(iter.u_perms)
+        iter.u_index = 1
+        iter.e_index += 1
+    end
+
+    if iter.e_index > length(iter.e_perms)
+        iter.e_index = 1
+        photon_iter_result = iterate(iter.photon_structure_iter, nothing)
+        if isnothing(photon_iter_result)
+            return nothing
+        end
+        (iter.photon_structure, _) = photon_iter_result
+    end
+
+    return ((iter.photon_structure, iter.e_perms[iter.e_index], iter.u_perms[iter.u_index], iter.t_perms[iter.t_index]), nothing)
+end
+
+function feynman_diagrams(e::Int, u::Int, t::Int, m::Int)
+    f_iter = _feynman_structures(e + u + t, m)
+    e_perms = collect(permutations(Int[1:e;]))
+    u_perms = collect(permutations(Int[e+1:e+u;]))
+    t_perms = collect(permutations(Int[e+u+1:e+u+t;]))
+    first_photon_structure, _ = iterate(f_iter)
+
+    return FeynmanDiagramIterator(e, e_perms, 1, u, u_perms, 1, t, t_perms, 1, m, f_iter, first_photon_structure)
+end

src/trees/labelled_plane_trees.jl

@@ -122,6 +122,10 @@ function _plane_tree_permutations(t::Vector{Vector{Int64}})
 end
 
 function labelled_plane_trees(N::Int64)
+    if N == 1
+        return [[Int[]]]
+    end
+
     all_trees = Vector{Vector{Vector{Int64}}}()
 
     sizehint!(all_trees, factorial(3N - 3, 2N - 1))