Wrap feynman diagram iterator results in usable object

src/FeynmanDiagramGenerator.jl

@@ -5,7 +5,7 @@ using QEDprocesses
 
 import Base.==
 
-export AbstractTreeLevelFeynmanDiagram, FeynmanVertex
+export AbstractTreeLevelFeynmanDiagram, FeynmanVertex, FeynmanDiagram
 export external_particles, virtual_particles, process, vertices
 
 export Forest

src/diagrams/diagrams.jl

@@ -43,8 +43,34 @@ end
 # Feynman Diagram, tree-level, QED
 #
 
-struct FeynmanDiagram <: AbstractTreeLevelFeynmanDiagram
+struct FeynmanDiagram{PROC,N,E,U,T,M} <: AbstractTreeLevelFeynmanDiagram where {P<:AbstractProcessDefinition,N,E,U,T,M}
-    proc::AbstractProcessDefinition
+    # E, U, T, and M can be inferred from the PROC, but not necessarily in 0 runtime
+    proc::PROC
+
+    diagram_structure::NTuple{N,Vector{Int}}
+
+    electron_permutation::NTuple{E,Int}
+    muon_permutation::NTuple{U,Int}
+    tauon_permutation::NTuple{T,Int}
+
+    function FeynmanDiagram(
+        proc::PROC,
+        structure::Vector{Vector{Int}},
+        elec_perm::Vector{Int},
+        muon_perm::Vector{Int},
+        tauon_perm::Vector{Int},
+        ::Val{E},
+        ::Val{U},
+        ::Val{T},
+        ::Val{M}
+    ) where {PROC<:AbstractProcessDefinition,E,U,T,M}
+        @assert E == length(elec_perm)
+        @assert U == length(muon_perm)
+        @assert T == length(tauon_perm)
+        N = E + U + T
+
+        return new{PROC,N,E,U,T,M}(proc, NTuple{N,Vector{Int}}(structure), NTuple{E,Int}(elec_perm), NTuple{U,Int}(muon_perm), NTuple{T,Int}(tauon_perm))
+    end
 end
 
 process(diagram::FeynmanDiagram) = d.proc
@@ -133,28 +159,32 @@ function _external_photon_multiplicity(fermion_structure::Vector{Vector{Int}}, n
     return res
 end
 
-mutable struct FeynmanDiagramIterator
+mutable struct FeynmanDiagramIterator{PROC<:AbstractProcessDefinition,E,U,T,M}
-    E::Int  # number of electron lines  (indices 1      - e)
+    process::PROC
+    e::Val{E}  # 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::Val{U}  # 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::Val{T}  # 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
+    m::Val{M}  # 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)
+function Base.length(it::FeynmanDiagramIterator{<:AbstractProcessDefinition,E,U,T,M}) where {E,U,T,M}
-    N = it.E + it.U + it.T
+    N = E + U + T
-    return factorial(it.M + 3 * N - 3, 2 * N - 1) * factorial(it.E) * factorial(it.U) * factorial(it.T)
+    return factorial(M + 3 * N - 3, 2 * N - 1) * factorial(E) * factorial(U) * factorial(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)
+    return (
+        FeynmanDiagram(iter.process, iter.photon_structure, iter.e_perms[iter.e_index], iter.u_perms[iter.u_index], iter.t_perms[iter.t_index], iter.e, iter.u, iter.t, iter.m),
+        nothing
+    )
 end
 
 function Base.iterate(iter::FeynmanDiagramIterator, ::Nothing)
@@ -179,15 +209,22 @@ function Base.iterate(iter::FeynmanDiagramIterator, ::Nothing)
         (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)
+    return (
+        FeynmanDiagram(iter.process, iter.photon_structure, iter.e_perms[iter.e_index], iter.u_perms[iter.u_index], iter.t_perms[iter.t_index], iter.e, iter.u, iter.t, iter.m),
+        nothing
+    )
+end
+
+function feynman_diagrams(proc::PROC, e::Int, u::Int, t::Int, m::Int) where {PROC<:AbstractProcessDefinition}
+    if (e + u + t) * 2 + m < 4
+        throw(InvalidInputError("At least 4 particles must participate in a scattering process for a physical result!"))
     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)
+    return FeynmanDiagramIterator(proc, Val(e), e_perms, 1, Val(u), u_perms, 1, Val(t), t_perms, 1, Val(m), f_iter, first_photon_structure)
 end