Add flat matrix

This commit is contained in:
Anton Reinhard 2024-06-20 15:58:03 +02:00
parent 2dd7627bcf
commit 2b7c02c223
3 changed files with 37 additions and 8 deletions

View File

@ -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")

View File

@ -43,9 +43,8 @@ end
# Feynman Diagram, tree-level, QED
#
struct FeynmanDiagram{N,E,U,T,M} <: AbstractTreeLevelFeynmanDiagram where {N,E,U,T,M}
# TODO: flatten into one list
diagram_structure::NTuple{N,Vector{Int}}
struct FeynmanDiagram{N,E,U,T,M,FM} <: AbstractTreeLevelFeynmanDiagram where {N,E,U,T,M,FM<:FlatMatrix}
diagram_structure::FM
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

23
src/flat_matrix.jl Normal file
View File

@ -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