From fb8dee171fa4545a1603fc9fe6dbf381407e8196 Mon Sep 17 00:00:00 2001 From: AntonReinhard Date: Mon, 27 May 2024 22:35:46 +0200 Subject: [PATCH] Add GPU tests and fix Compton implementation for the testss --- Project.toml | 13 +- src/QEDprocesses.jl | 1 + src/momentum_generation.jl | 12 +- src/patch_QEDbase.jl | 10 +- src/phase_spaces/access.jl | 12 +- .../perturbative/cross_section.jl | 33 ++-- .../perturbative/total_probability.jl | 3 +- test/gpu/process_interface.jl | 163 ++++++++++++++++++ test/runtests.jl | 4 + .../test_implementation/TestImplementation.jl | 1 + .../test_implementation/random_coordinates.jl | 8 + test/test_implementation/utils.jl | 8 + 12 files changed, 226 insertions(+), 42 deletions(-) create mode 100644 test/gpu/process_interface.jl create mode 100644 test/test_implementation/random_coordinates.jl diff --git a/Project.toml b/Project.toml index 2b082b5..c41da04 100644 --- a/Project.toml +++ b/Project.toml @@ -1,15 +1,10 @@ name = "QEDprocesses" uuid = "46de9c38-1bb3-4547-a1ec-da24d767fdad" -authors = [ - "Uwe Hernandez Acosta ", - "Simeon Ehrig", - "Klaus Steiniger", - "Tom Jungnickel", - "Anton Reinhard", -] +authors = ["Uwe Hernandez Acosta ", "Simeon Ehrig", "Klaus Steiniger", "Tom Jungnickel", "Anton Reinhard"] version = "0.1.0" [deps] +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" QEDbase = "10e22c08-3ccb-4172-bfcf-7d7aa3d04d93" QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" @@ -21,10 +16,12 @@ StaticArrays = "1" julia = "1.6" [extras] +AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Random", "SafeTestsets", "Suppressor", "Test"] +test = ["Random", "SafeTestsets", "Suppressor", "Test", "AMDGPU", "CUDA"] diff --git a/src/QEDprocesses.jl b/src/QEDprocesses.jl index 608a96e..ba4849e 100644 --- a/src/QEDprocesses.jl +++ b/src/QEDprocesses.jl @@ -11,6 +11,7 @@ export AbstractModelDefinition, fundamental_interaction_type export AbstractProcessDefinition, incoming_particles, outgoing_particles export number_incoming_particles, number_outgoing_particles export particles, number_particles +export in_phase_space_dimension, out_phase_space_dimension # probabilities export differential_probability, unsafe_differential_probability diff --git a/src/momentum_generation.jl b/src/momentum_generation.jl index 7e7d278..b301389 100644 --- a/src/momentum_generation.jl +++ b/src/momentum_generation.jl @@ -31,12 +31,12 @@ function _generate_outgoing_momenta end """ _generate_momenta( - proc::AbstractProcessDefinition, - model::AbstractModelDefinition, - phase_space_def::AbstractPhasespaceDefinition, - in_phase_space::NTuple{N,T}, - out_phase_space::NTuple{M,T}, -) where {N,M,T<:Real} + proc::AbstractProcessDefinition, + model::AbstractModelDefinition, + phase_space_def::AbstractPhasespaceDefinition, + in_phase_space::NTuple{N,T}, + out_phase_space::NTuple{M,T}, + ) where {N,M,T<:Real} Return four-momenta for incoming and outgoing particles for given coordinate based phase-space points. """ diff --git a/src/patch_QEDbase.jl b/src/patch_QEDbase.jl index 2cc492f..091b16a 100644 --- a/src/patch_QEDbase.jl +++ b/src/patch_QEDbase.jl @@ -4,7 +4,7 @@ # ############# -# fix: https://github.com/QEDjl-project/QEDbase.jl/pull/61 +# fix: https://github.com/QEDjl-project/QEDbase.jl/pull/65 Base.show(io::IO, ::Electron) = print(io, "electron") Base.show(io::IO, ::Positron) = print(io, "positron") Base.show(io::IO, ::Photon) = print(io, "photon") @@ -24,7 +24,7 @@ Broadcast.broadcastable(part::AbstractParticleType) = Ref(part) Broadcast.broadcastable(spin_or_pol::AbstractSpinOrPolarization) = Ref(spin_or_pol) # fix: https://github.com/QEDjl-project/QEDbase.jl/pull/63 -number_of_spin_pol(::AbstractDefinitePolarization) = 1 -number_of_spin_pol(::AbstractDefiniteSpin) = 1 -number_of_spin_pol(::AbstractIndefinitePolarization) = 2 -number_of_spin_pol(::AbstractIndefiniteSpin) = 2 +multiplicity(::AbstractDefinitePolarization) = 1 +multiplicity(::AbstractDefiniteSpin) = 1 +multiplicity(::AbstractIndefinitePolarization) = 2 +multiplicity(::AbstractIndefiniteSpin) = 2 diff --git a/src/phase_spaces/access.jl b/src/phase_spaces/access.jl index 6436be0..8a881cc 100644 --- a/src/phase_spaces/access.jl +++ b/src/phase_spaces/access.jl @@ -28,8 +28,16 @@ momentum(part::ParticleStateful) = part.mom Return a `Tuple` of all the particles' momenta for the given `ParticleDirection`. """ -momenta(psp::PhaseSpacePoint, ::Incoming) = momentum.(psp.in_particles) -momenta(psp::PhaseSpacePoint, ::Outgoing) = momentum.(psp.out_particles) +function momenta( + psp::PhaseSpacePoint{P,M,D,I,O,E}, ::Incoming +) where {N,P,M,D,I<:Tuple{Vararg{ParticleStateful,N}},O,E} + return NTuple{N,E}(momentum(p) for p in psp.in_particles) +end +function momenta( + psp::PhaseSpacePoint{P,M,D,I,O,E}, ::Outgoing +) where {N,P,M,D,I,O<:Tuple{Vararg{ParticleStateful,N}},E} + return NTuple{N,E}(momentum(p) for p in psp.out_particles) +end """ Base.getindex(psp::PhaseSpacePoint, dir::Incoming, n::Int) diff --git a/src/processes/one_photon_compton/perturbative/cross_section.jl b/src/processes/one_photon_compton/perturbative/cross_section.jl index a451da9..9f49a48 100644 --- a/src/processes/one_photon_compton/perturbative/cross_section.jl +++ b/src/processes/one_photon_compton/perturbative/cross_section.jl @@ -21,7 +21,7 @@ end We average over the initial spins and pols, and sum over final. """ function _averaging_norm(proc::Compton) - normalizations = number_of_spin_pol.(_in_spin_and_pol(proc)) + normalizations = multiplicity.(_in_spin_and_pol(proc)) return inv(prod(normalizations)) end @@ -74,8 +74,8 @@ end in_photon_state = base_state(Photon(), Incoming(), in_photon_mom, proc.in_pol) out_electron_state = base_state(Electron(), Outgoing(), out_electron_mom, proc.out_spin) - out_photon_state = base_state(Photon(), Outgoing(), out_photon_mom, proc.out_pol) + return _pert_compton_matrix_element( in_electron_mom, in_electron_state, @@ -105,23 +105,18 @@ function _pert_compton_matrix_element( QEDbase._as_svec(out_photon_state), ) - matrix_elements = Vector{ComplexF64}() - sizehint!(matrix_elements, length(base_states_comb)) - for (in_el, in_ph, out_el, out_ph) in base_states_comb - push!( - matrix_elements, - _pert_compton_matrix_element_single( - in_electron_mom, - in_el, - in_photon_mom, - in_ph, - out_electron_mom, - out_el, - out_photon_mom, - out_ph, - ), - ) - end + matrix_elements = NTuple{length(base_states_comb),ComplexF64}( + _pert_compton_matrix_element_single( + in_electron_mom, + in_el, + in_photon_mom, + in_ph, + out_electron_mom, + out_el, + out_photon_mom, + out_ph, + ) for (in_el, in_ph, out_el, out_ph) in base_states_comb + ) return matrix_elements end diff --git a/src/processes/one_photon_compton/perturbative/total_probability.jl b/src/processes/one_photon_compton/perturbative/total_probability.jl index 6699b17..c43122d 100644 --- a/src/processes/one_photon_compton/perturbative/total_probability.jl +++ b/src/processes/one_photon_compton/perturbative/total_probability.jl @@ -1,6 +1,6 @@ function _total_probability(in_psp::InPhaseSpacePoint{<:Compton,<:PerturbativeQED}) - omega = getE(momentum(in_psp[Incoming(), 2])) + omega = getE(momentum(in_psp, Incoming(), 2)) function func(x) return unsafe_differential_probability( @@ -9,7 +9,6 @@ function _total_probability(in_psp::InPhaseSpacePoint{<:Compton,<:PerturbativeQE end tot_prob, _ = quadgk(func, -1, 1; rtol=sqrt(eps(omega))) - tot_prob *= 2 * pi # phi integration is trivial return tot_prob end diff --git a/test/gpu/process_interface.jl b/test/gpu/process_interface.jl new file mode 100644 index 0000000..8f965ca --- /dev/null +++ b/test/gpu/process_interface.jl @@ -0,0 +1,163 @@ +using QEDprocesses +using QEDbase +using AMDGPU +using CUDA +using Random +using SafeTestsets + +include("../test_implementation/TestImplementation.jl") + +GPU_VECTOR_TYPES = Vector{Type}() +if AMDGPU.functional() + println("Testing with AMDGPU.jl") + push!(GPU_VECTOR_TYPES, ROCVector) +end +if CUDA.functional() + println("Testing with CUDA.jl") + push!(GPU_VECTOR_TYPES, CuVector) +end + +if isempty(GPU_VECTOR_TYPES) + println("No functional GPUs found!") + return nothing +end + +PROC_DEF_TUPLES = [ + ( + Compton(), + PerturbativeQED(), + PhasespaceDefinition(SphericalCoordinateSystem(), ElectronRestFrame()), + ), + ( + Compton(SpinUp(), PolX(), SpinDown(), PolY()), + PerturbativeQED(), + PhasespaceDefinition(SphericalCoordinateSystem(), ElectronRestFrame()), + ), +] + +RNG = Random.MersenneTwister(573) + +@testset "Testing with $VECTOR_TYPE" for VECTOR_TYPE in GPU_VECTOR_TYPES + @testset "$proc $model $ps_def" for (proc, model, ps_def) in PROC_DEF_TUPLES + N = 10_000 + + psps = [ + PhaseSpacePoint( + proc, + model, + ps_def, + TestImplementation._rand_coordinates(RNG, proc, model, ps_def)..., + ) for _ in 1:N + ] + procs = [proc for _ in 1:N] + + gpupsps = VECTOR_TYPE(psps) + gpuprocs = VECTOR_TYPE(procs) + + @testset "PSP interface" begin + in_moms_gpu = Vector(momenta.(gpupsps, Incoming())) + out_moms_gpu = Vector(momenta.(gpupsps, Outgoing())) + in_moms = momenta.(psps, Incoming()) + out_moms = momenta.(psps, Outgoing()) + + @test getindex.(in_moms_gpu, Ref(1)) == getindex.(in_moms, Ref(1)) + @test getindex.(in_moms_gpu, Ref(2)) == getindex.(in_moms, Ref(2)) + @test getindex.(out_moms_gpu, Ref(1)) == getindex.(out_moms, Ref(1)) + @test getindex.(out_moms_gpu, Ref(2)) == getindex.(out_moms, Ref(2)) + end + + @testset "Private Process Functions" begin + @test all( + isapprox.( + Vector(QEDprocesses._averaging_norm.(gpuprocs)), + QEDprocesses._averaging_norm.(procs), + ), + ) + end + + @testset "Public Process Functions" begin + @test Vector(incoming_particles.(gpuprocs)) == incoming_particles.(procs) + @test Vector(outgoing_particles.(gpuprocs)) == outgoing_particles.(procs) + + @test Vector(particles.(gpuprocs, Incoming())) == particles.(procs, Incoming()) + @test Vector(particles.(gpuprocs, Outgoing())) == particles.(procs, Outgoing()) + + @test Vector(number_incoming_particles.(gpuprocs)) == + number_incoming_particles.(procs) + @test Vector(number_outgoing_particles.(gpuprocs)) == + number_outgoing_particles.(procs) + + @test Vector(number_particles.(gpuprocs, Incoming())) == + number_particles.(procs, Incoming()) + @test Vector(number_particles.(gpuprocs, Outgoing())) == + number_particles.(procs, Outgoing()) + + @test Vector(in_phase_space_dimension.(gpuprocs, model)) == + in_phase_space_dimension.(procs, model) + @test Vector(out_phase_space_dimension.(gpuprocs, model)) == + out_phase_space_dimension.(procs, model) + end + + @testset "Private PhaseSpacePoint Interface" begin + @test all( + isapprox.( + Vector(QEDprocesses._incident_flux.(gpupsps)), + QEDprocesses._incident_flux.(psps), + ), + ) + + @test all( + TestImplementation.tuple_isapprox.( + Vector(QEDprocesses._matrix_element.(gpupsps)), + QEDprocesses._matrix_element.(psps), + ), + ) + + @test Vector(QEDprocesses._is_in_phasespace.(gpupsps)) == + QEDprocesses._is_in_phasespace.(psps) + + @test all( + isapprox.( + Vector(QEDprocesses._phase_space_factor.(gpupsps)), + QEDprocesses._phase_space_factor.(psps), + ), + ) + + #= TODO: this is currently broken + @test all( + isapprox.( + Vector(QEDprocesses._total_probability.(gpupsps)), + QEDprocesses._total_probability.(psps), + ), + )=# + end + + @testset "Public PhaseSpacePoint Interface" begin + @test all( + isapprox.( + Vector(differential_probability.(gpupsps)), + differential_probability.(psps), + ), + ) + + @test all( + isapprox.( + Vector(QEDprocesses._is_in_phasespace.(gpupsps)), + QEDprocesses._is_in_phasespace.(psps), + ), + ) + + @test all( + isapprox.( + Vector(differential_cross_section.(gpupsps)), + differential_cross_section.(psps), + ), + ) + + #= TODO: this is currently broken + @test all( + isapprox.(Vector(total_cross_section.(gpupsps)), total_cross_section.(psps)) + )=# + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 9ecc6ee..aa58584 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -28,4 +28,8 @@ begin # scattering processes include("processes/run_process_test.jl") + + @time @safetestset "GPU process interface" begin + include("gpu/process_interface.jl") + end end diff --git a/test/test_implementation/TestImplementation.jl b/test/test_implementation/TestImplementation.jl index b9d4faa..dd22083 100644 --- a/test/test_implementation/TestImplementation.jl +++ b/test/test_implementation/TestImplementation.jl @@ -34,6 +34,7 @@ include("groundtruths.jl") include("test_model.jl") include("test_process.jl") include("random_momenta.jl") +include("random_coordinates.jl") include("utils.jl") end diff --git a/test/test_implementation/random_coordinates.jl b/test/test_implementation/random_coordinates.jl new file mode 100644 index 0000000..8e2623e --- /dev/null +++ b/test/test_implementation/random_coordinates.jl @@ -0,0 +1,8 @@ +""" +Return a tuple of tuples of incoming and outgoing coordinates for a given process, model and ps_def that make up a physical phase space point. +""" +function _rand_coordinates( + rng::AbstractRNG, ::PROCESS, ::MODEL, ::PS_DEF +) where {PROCESS<:Compton,MODEL<:PerturbativeQED,PS_DEF<:PhasespaceDefinition} + return ((rand(rng, Float64),), (rand(rng, Float64), rand(rng, Float64))) +end diff --git a/test/test_implementation/utils.jl b/test/test_implementation/utils.jl index 3a141d9..11e3f68 100644 --- a/test/test_implementation/utils.jl +++ b/test/test_implementation/utils.jl @@ -42,3 +42,11 @@ end function _furl_moms(moms::NTuple{N,Float64}) where {N} return Tuple(_furl_moms(Vector{Float64}([moms...]))) end + +tuple_isapprox(::Tuple{}, ::Tuple{}; atol=0.0, rtol=eps()) = true +function tuple_isapprox( + a::Tuple{<:Number,Vararg}, b::Tuple{<:Number,Vararg}; atol=0.0, rtol=eps() +) + return isapprox(a[1], b[1]; atol=atol, rtol=rtol) && + tuple_isapprox(a[2:end], b[2:end]; atol=atol, rtol=rtol) +end