Add flat matrix

src/FeynmanDiagramGenerator.jl

@@ -7,6 +7,8 @@ include("QEDprocesses_patch.jl")
 
 import Base.==
 
+export FlatMatrix
+
 export GenericQEDProcess, isphysical
 
 export AbstractTreeLevelFeynmanDiagram, FeynmanVertex, FeynmanDiagram
@@ -19,6 +21,8 @@ export plane_trees, labelled_plane_trees, feynman_diagrams
 export can_interact
 export QED
 
+include("flat_matrix.jl")
+
 include("trees/labelled_plane_trees.jl")
 include("trees/trees.jl")
 include("trees/iterator.jl")

src/diagrams/diagrams.jl

@@ -43,9 +43,8 @@ end
 # Feynman Diagram, tree-level, QED
 #
 
-struct FeynmanDiagram{N,E,U,T,M} <: AbstractTreeLevelFeynmanDiagram where {N,E,U,T,M}
+struct FeynmanDiagram{N,E,U,T,M,FM} <: AbstractTreeLevelFeynmanDiagram where {N,E,U,T,M,FM<:FlatMatrix}
-    # TODO: flatten into one list
+    diagram_structure::FM
-    diagram_structure::NTuple{N,Vector{Int}}
 
     electron_permutation::NTuple{E,Int}
     muon_permutation::NTuple{U,Int}
@@ -66,13 +65,14 @@ struct FeynmanDiagram{N,E,U,T,M} <: AbstractTreeLevelFeynmanDiagram where {N,E,U
         @assert T == length(tauon_perm)
         N = E + U + T
 
-        return new{N,E,U,T,M}(NTuple{N,Vector{Int}}(structure), NTuple{E,Int}(elec_perm), NTuple{U,Int}(muon_perm), NTuple{T,Int}(tauon_perm))
+        fm = FlatMatrix(structure)
+        return new{N,E,U,T,M,typeof(fm)}(fm, NTuple{E,Int}(elec_perm), NTuple{U,Int}(muon_perm), NTuple{T,Int}(tauon_perm))
     end
 end
 
-function virtual_particles(diagram::FeynmanDiagram)
+function virtual_particles(diagram::FeynmanDiagram{N,E,U,T,M,FM})
 
-    return NTuple{N,Tuple{QEDbase.AbstractParticleType,BitArray}}()
+    return NTuple{N,Tuple{QEDbase.AbstractParticleType,BitArray,BitArray}}()
 end
 
 function vertices(::AbstractTreeLevelFeynmanDiagram)
@@ -175,8 +175,9 @@ function Base.length(it::FeynmanDiagramIterator{E,U,T,M}) where {E,U,T,M}
 end
 
 function Base.iterate(iter::FeynmanDiagramIterator)
+    f = FeynmanDiagram(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)
     return (
-        FeynmanDiagram(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),
+        f,
         nothing
     )
 end
@@ -203,8 +204,9 @@ function Base.iterate(iter::FeynmanDiagramIterator, ::Nothing)
         (iter.photon_structure, _) = photon_iter_result
     end
 
+    f = FeynmanDiagram(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)
     return (
-        FeynmanDiagram(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),
+        f,
         nothing
     )
 end

src/flat_matrix.jl

@@ -0,0 +1,23 @@
+# array of arrays but with a given number of arrays (M) and given total length (N)
+
+struct FlatMatrix{T,N,M}
+    values::NTuple{N,T}
+    indices::NTuple{M,Int}
+
+    function FlatMatrix(v::Vector{Vector{T}}) where {T}
+        M = length(v)
+        N = sum(length.(v))
+
+        values = NTuple{N,T}(vcat(v...))
+        indices = ntuple(i -> sum(length.(v[1:i-1])), M)
+
+        return new{Int,N,M}(values, indices)
+    end
+end
+
+function Base.getindex(m::FlatMatrix{T,N,M}, x::Int, y::Int) where {T,N,M}
+    x <= M || throw(InvalidInputError("invalid indices ($x, $y) for flat matrix $m"))
+    (x > 0 && y > 0) || throw(InvalidInputError("invalid indices ($x, $y) for flat matrix $m"))
+    x == M ? m.indices[x] + y <= N : m.indices[x] + y <= m.indices[x+1] || throw(InvalidInputError("invalid indices ($x, $y) for flat matrix $m"))
+    return m.values[m.indices[x]+y]
+end