metagraphoptimization.jl/src/QEDprocesses_patch.jl

72 lines
2.1 KiB
Julia

# patch QEDprocesses
# see issue https://github.com/QEDjl-project/QEDprocesses.jl/issues/77
@inline function QEDprocesses.number_particles(
proc_def::QEDbase.AbstractProcessDefinition,
dir::DIR,
::PT,
) where {DIR <: QEDbase.ParticleDirection, PT <: QEDbase.AbstractParticleType}
return count(x -> x isa PT, particles(proc_def, dir))
end
@inline function QEDprocesses.number_particles(
proc_def::QEDbase.AbstractProcessDefinition,
::PS,
) where {
DIR <: QEDbase.ParticleDirection,
PT <: QEDbase.AbstractParticleType,
EL <: AbstractFourMomentum,
PS <: ParticleStateful{DIR, PT, EL},
}
return QEDprocesses.number_particles(proc_def, DIR(), PT())
end
@inline function QEDprocesses.number_particles(
proc_def::QEDbase.AbstractProcessDefinition,
::Type{PS},
) where {
DIR <: QEDbase.ParticleDirection,
PT <: QEDbase.AbstractParticleType,
EL <: AbstractFourMomentum,
PS <: ParticleStateful{DIR, PT, EL},
}
return QEDprocesses.number_particles(proc_def, DIR(), PT())
end
@inline function QEDcore.ParticleStateful{DIR, SPECIES}(
mom::AbstractFourMomentum,
) where {DIR <: ParticleDirection, SPECIES <: AbstractParticleType}
return ParticleStateful(DIR(), SPECIES(), mom)
end
@inline function QEDcore.ParticleStateful{DIR, SPECIES, EL}(
mom::EL,
) where {DIR <: ParticleDirection, SPECIES <: AbstractParticleType, EL <: AbstractFourMomentum}
return ParticleStateful(DIR(), SPECIES(), mom)
end
@inline function QEDbase.momentum(
psp::AbstractPhaseSpacePoint{MODEL, PROC, PS_DEF, INT, OUTT},
dir::ParticleDirection,
species::AbstractParticleType,
n::Int,
) where {MODEL, PROC, PS_DEF, INT, OUTT}
# TODO: can be done through fancy template recursion too with 0 overhead
i = 0
c = n
for p in particles(psp, dir)
i += 1
if particle_species(p) isa typeof(species)
c -= 1
end
if c == 0
break
end
end
if c != 0 || n <= 0
throw(InvalidInputError("could not get $n-th momentum of $dir $species, does not exist"))
end
return momenta(psp, dir)[i]
end