From 55501c15c8c92f73dc90f7f8e98cf4dc19382e9d Mon Sep 17 00:00:00 2001 From: Rubydragon Date: Thu, 27 Jun 2024 13:21:53 +0200 Subject: [PATCH 1/6] WIP --- examples/congruent_in_ph.jl | 123 +++++++++++++++++++ src/models/physics_models/qed/create.jl | 47 ++++---- src/models/physics_models/qed/diagrams.jl | 10 +- test/unit_tests_qedmodel.jl | 140 +++++++++++----------- 4 files changed, 223 insertions(+), 97 deletions(-) create mode 100644 examples/congruent_in_ph.jl diff --git a/examples/congruent_in_ph.jl b/examples/congruent_in_ph.jl new file mode 100644 index 0000000..a0639bf --- /dev/null +++ b/examples/congruent_in_ph.jl @@ -0,0 +1,123 @@ +using MetagraphOptimization +using QEDbase +using QEDcore +using Random +using UUIDs + +RNG = Random.MersenneTwister(123) + +function mock_machine() + return Machine( + [ + MetagraphOptimization.NumaNode( + 0, + 1, + MetagraphOptimization.default_strategy(MetagraphOptimization.NumaNode), + -1.0, + UUIDs.uuid1(), + ), + ], + [-1.0;;], + ) +end + +function congruent_input(processDescription::QEDProcessDescription, omega::Number) + # generate an input sample for given e + nk -> e' + k' process, where the nk are equal + massSum = 0 + inputMasses = Vector{Float64}() + for (particle, n) in processDescription.inParticles + for _ in 1:n + massSum += mass(particle) + push!(inputMasses, mass(particle)) + end + end + outputMasses = Vector{Float64}() + for (particle, n) in processDescription.outParticles + for _ in 1:n + massSum += mass(particle) + push!(outputMasses, mass(particle)) + end + end + + initialMomenta = [ + i == 1 ? SFourMomentum(1, 0, 0, 0) : SFourMomentum(omega, 0, 0, omega) for + i in 1:length(inputMasses) + ] + + # add some extra random mass to allow for some momentum + ss = sqrt(sum(initialMomenta) * sum(initialMomenta)) + + result = Vector{QEDProcessInput}() + sizehint!(result, 16) + + spin_pol_combinations = Iterators.product( + [SpinUp, SpinDown], [SpinUp, SpinDown], [PolX, PolY], [PolX, PolY] + ) + for (in_spin, out_spin, in_pol, out_pol) in spin_pol_combinations + + # get the electron first, then the n photons + particles = Vector{QEDParticle}() + + for (particle, n) in processDescription.inParticles + if particle <: FermionStateful + mom = initialMomenta[1] + push!(particles, particle(mom, in_spin())) + elseif particle <: PhotonStateful + for i in 1:n + mom = initialMomenta[i + 1] + push!(particles, particle(mom, in_pol())) + end + else + @assert false + end + end + + final_momenta = MetagraphOptimization.generate_physical_massive_moms( + RNG, ss, outputMasses + ) + index = 1 + for (particle, n) in processDescription.outParticles + for _ in 1:n + if particle <: FermionStateful + push!(particles, particle(final_momenta[index], out_spin())) + elseif particle <: PhotonStateful + push!(particles, particle(final_momenta[index], out_pol())) + end + index += 1 + end + end + + inFerms = MetagraphOptimization._svector_from_type( + processDescription, FermionStateful{Incoming,in_spin}, particles + ) + outFerms = MetagraphOptimization._svector_from_type( + processDescription, FermionStateful{Outgoing,out_spin}, particles + ) + inAntiferms = MetagraphOptimization._svector_from_type( + processDescription, AntiFermionStateful{Incoming,in_spin}, particles + ) + outAntiferms = MetagraphOptimization._svector_from_type( + processDescription, AntiFermionStateful{Outgoing,out_spin}, particles + ) + inPhotons = MetagraphOptimization._svector_from_type( + processDescription, PhotonStateful{Incoming,in_pol}, particles + ) + outPhotons = MetagraphOptimization._svector_from_type( + processDescription, PhotonStateful{Outgoing,out_pol}, particles + ) + + processInput = QEDProcessInput( + processDescription, + inFerms, + outFerms, + inAntiferms, + outAntiferms, + inPhotons, + outPhotons, + ) + + push!(result, processInput) + end + + return result +end diff --git a/src/models/physics_models/qed/create.jl b/src/models/physics_models/qed/create.jl index b8a8a50..94454d8 100644 --- a/src/models/physics_models/qed/create.jl +++ b/src/models/physics_models/qed/create.jl @@ -8,7 +8,6 @@ function _svector_from_type(processDescription::GenericQEDProcess, type, particl if haskey(out_particles(processDescription), type) return SVector{out_particles(processDescription)[type], type}(filter(x -> typeof(x) <: type, particles)) end - return SVector{0, type}() end """ @@ -64,9 +63,9 @@ function gen_graph(process_description::GenericQEDProcess) # TODO: Not all diagram outputs should always be summed at the end, if they differ by fermion exchange they need to be diffed # Should not matter for n-Photon Compton processes though - sum_node = insert_node!(graph, make_node(ComputeTaskQED_Sum(0)), track = false, invalidate_cache = false) - global_data_out = insert_node!(graph, make_node(DataTask(COMPLEX_SIZE)), track = false, invalidate_cache = false) - insert_edge!(graph, sum_node, global_data_out, track = false, invalidate_cache = false) + sum_node = insert_node!(graph, make_node(ComputeTaskQED_Sum(0)); track = false, invalidate_cache = false) + global_data_out = insert_node!(graph, make_node(DataTask(COMPLEX_SIZE)); track = false, invalidate_cache = false) + insert_edge!(graph, sum_node, global_data_out; track = false, invalidate_cache = false) # remember the data out nodes for connection dataOutNodes = Dict() @@ -75,16 +74,16 @@ function gen_graph(process_description::GenericQEDProcess) # generate data in and U tasks data_in = insert_node!( graph, - make_node(DataTask(PARTICLE_VALUE_SIZE), String(particle)), + make_node(DataTask(PARTICLE_VALUE_SIZE), String(particle)); track = false, invalidate_cache = false, ) # read particle data node - compute_u = insert_node!(graph, make_node(ComputeTaskQED_U()), track = false, invalidate_cache = false) # compute U node + compute_u = insert_node!(graph, make_node(ComputeTaskQED_U()); track = false, invalidate_cache = false) # compute U node data_out = - insert_node!(graph, make_node(DataTask(PARTICLE_VALUE_SIZE)), track = false, invalidate_cache = false) # transfer data out from u (one ABCParticleValue object) + insert_node!(graph, make_node(DataTask(PARTICLE_VALUE_SIZE)); track = false, invalidate_cache = false) # transfer data out from u (one ABCParticleValue object) - insert_edge!(graph, data_in, compute_u, track = false, invalidate_cache = false) - insert_edge!(graph, compute_u, data_out, track = false, invalidate_cache = false) + insert_edge!(graph, data_in, compute_u; track = false, invalidate_cache = false) + insert_edge!(graph, compute_u, data_out; track = false, invalidate_cache = false) # remember the data_out node for future edges dataOutNodes[String(particle)] = data_out @@ -100,19 +99,19 @@ function gen_graph(process_description::GenericQEDProcess) data_in1 = dataOutNodes[String(vertex.in1)] data_in2 = dataOutNodes[String(vertex.in2)] - compute_V = insert_node!(graph, make_node(ComputeTaskQED_V()), track = false, invalidate_cache = false) # compute vertex + compute_V = insert_node!(graph, make_node(ComputeTaskQED_V()); track = false, invalidate_cache = false) # compute vertex - insert_edge!(graph, data_in1, compute_V, track = false, invalidate_cache = false) - insert_edge!(graph, data_in2, compute_V, track = false, invalidate_cache = false) + insert_edge!(graph, data_in1, compute_V; track = false, invalidate_cache = false) + insert_edge!(graph, data_in2, compute_V; track = false, invalidate_cache = false) data_V_out = insert_node!( graph, - make_node(DataTask(PARTICLE_VALUE_SIZE)), + make_node(DataTask(PARTICLE_VALUE_SIZE)); track = false, invalidate_cache = false, ) - insert_edge!(graph, compute_V, data_V_out, track = false, invalidate_cache = false) + insert_edge!(graph, compute_V, data_V_out; track = false, invalidate_cache = false) if (vertex.out == tie.in1 || vertex.out == tie.in2) # out particle is part of the tie -> there will be an S2 task with it later, don't make S1 task @@ -122,18 +121,18 @@ function gen_graph(process_description::GenericQEDProcess) # otherwise, add S1 task compute_S1 = - insert_node!(graph, make_node(ComputeTaskQED_S1()), track = false, invalidate_cache = false) # compute propagator + insert_node!(graph, make_node(ComputeTaskQED_S1()); track = false, invalidate_cache = false) # compute propagator - insert_edge!(graph, data_V_out, compute_S1, track = false, invalidate_cache = false) + insert_edge!(graph, data_V_out, compute_S1; track = false, invalidate_cache = false) data_S1_out = insert_node!( graph, - make_node(DataTask(PARTICLE_VALUE_SIZE)), + make_node(DataTask(PARTICLE_VALUE_SIZE)); track = false, invalidate_cache = false, ) - insert_edge!(graph, compute_S1, data_S1_out, track = false, invalidate_cache = false) + insert_edge!(graph, compute_S1, data_S1_out; track = false, invalidate_cache = false) # overrides potentially different nodes from previous diagrams, which is intentional dataOutNodes[String(vertex.out)] = data_S1_out @@ -144,16 +143,16 @@ function gen_graph(process_description::GenericQEDProcess) data_in1 = dataOutNodes[String(tie.in1)] data_in2 = dataOutNodes[String(tie.in2)] - compute_S2 = insert_node!(graph, make_node(ComputeTaskQED_S2()), track = false, invalidate_cache = false) + compute_S2 = insert_node!(graph, make_node(ComputeTaskQED_S2()); track = false, invalidate_cache = false) - data_S2 = insert_node!(graph, make_node(DataTask(PARTICLE_VALUE_SIZE)), track = false, invalidate_cache = false) + data_S2 = insert_node!(graph, make_node(DataTask(PARTICLE_VALUE_SIZE)); track = false, invalidate_cache = false) - insert_edge!(graph, data_in1, compute_S2, track = false, invalidate_cache = false) - insert_edge!(graph, data_in2, compute_S2, track = false, invalidate_cache = false) + insert_edge!(graph, data_in1, compute_S2; track = false, invalidate_cache = false) + insert_edge!(graph, data_in2, compute_S2; track = false, invalidate_cache = false) - insert_edge!(graph, compute_S2, data_S2, track = false, invalidate_cache = false) + insert_edge!(graph, compute_S2, data_S2; track = false, invalidate_cache = false) - insert_edge!(graph, data_S2, sum_node, track = false, invalidate_cache = false) + insert_edge!(graph, data_S2, sum_node; track = false, invalidate_cache = false) add_child!(task(sum_node)) end diff --git a/src/models/physics_models/qed/diagrams.jl b/src/models/physics_models/qed/diagrams.jl index e3f94a6..27082e5 100644 --- a/src/models/physics_models/qed/diagrams.jl +++ b/src/models/physics_models/qed/diagrams.jl @@ -152,8 +152,9 @@ function ==(d1::FeynmanDiagram, d2::FeynmanDiagram) )=# end -copy(fd::FeynmanDiagram) = - FeynmanDiagram(deepcopy(fd.vertices), copy(fd.tie[]), deepcopy(fd.particles), copy(fd.type_ids)) +function copy(fd::FeynmanDiagram) + return FeynmanDiagram(deepcopy(fd.vertices), copy(fd.tie[]), deepcopy(fd.particles), copy(fd.type_ids)) +end """ id_for_type(d::FeynmanDiagram, t::Type{<:ParticleStateful}) @@ -509,7 +510,9 @@ end Helper function for [`gen_compton_diagrams`](@Ref). Generates a single diagram for the given order and n input and m output photons. """ -function gen_compton_diagram_from_order_one_side(order::Vector{Int}, inFerm, outFerm, n::Int, m::Int) +function gen_compton_diagram_from_order_one_side( + order::Vector{Int}, inFerm, outFerm, n::Int, m::Int +) photons = vcat( [FeynmanParticle(ParticleStateful{Incoming, Photon}, i) for i in 1:n], [FeynmanParticle(ParticleStateful{Outgoing, Photon}, i) for i in 1:m], @@ -587,7 +590,6 @@ function gen_compton_diagrams(n::Int, m::Int) return vcat(diagrams...) end - """ gen_compton_diagrams_one_side(n::Int, m::Int) diff --git a/test/unit_tests_qedmodel.jl b/test/unit_tests_qedmodel.jl index 6a9defe..9a48050 100644 --- a/test/unit_tests_qedmodel.jl +++ b/test/unit_tests_qedmodel.jl @@ -56,8 +56,8 @@ function compton_groundtruth(input::PhaseSpacePoint) virt2_mom = p1.momentum + k1.momentum @test isapprox(p2.momentum + k2.momentum, virt2_mom) - s_p2_k1 = propagator(Electron(), virt1_mom) - s_p1_k1 = propagator(Electron(), virt2_mom) + s_p2_k1 = QEDbase.propagator(Electron(), virt1_mom) + s_p1_k1 = QEDbase.propagator(Electron(), virt2_mom) diagram1 = u_p2 * (eps_1 * QED_vertex()) * s_p2_k1 * (eps_2 * QED_vertex()) * u_p1 diagram2 = u_p2 * (eps_2 * QED_vertex()) * s_p1_k1 * (eps_1 * QED_vertex()) * u_p1 @@ -65,7 +65,6 @@ function compton_groundtruth(input::PhaseSpacePoint) return diagram1 + diagram2 end - @testset "Interaction Result" begin import MetagraphOptimization.QED_conserve_momentum @@ -202,97 +201,97 @@ end graph = DAG() # s to output (exit node) - d_exit = insert_node!(graph, make_node(DataTask(16)), track = false) + d_exit = insert_node!(graph, make_node(DataTask(16)); track=false) - sum_node = insert_node!(graph, make_node(ComputeTaskQED_Sum(2)), track = false) + sum_node = insert_node!(graph, make_node(ComputeTaskQED_Sum(2)); track=false) - d_s0_sum = insert_node!(graph, make_node(DataTask(16)), track = false) - d_s1_sum = insert_node!(graph, make_node(DataTask(16)), track = false) + d_s0_sum = insert_node!(graph, make_node(DataTask(16)); track=false) + d_s1_sum = insert_node!(graph, make_node(DataTask(16)); track=false) # final s compute - s0 = insert_node!(graph, make_node(ComputeTaskQED_S2()), track = false) - s1 = insert_node!(graph, make_node(ComputeTaskQED_S2()), track = false) + s0 = insert_node!(graph, make_node(ComputeTaskQED_S2()); track=false) + s1 = insert_node!(graph, make_node(ComputeTaskQED_S2()); track=false) # data from v0 and v1 to s0 - d_v0_s0 = insert_node!(graph, make_node(DataTask(96)), track = false) - d_v1_s0 = insert_node!(graph, make_node(DataTask(96)), track = false) - d_v2_s1 = insert_node!(graph, make_node(DataTask(96)), track = false) - d_v3_s1 = insert_node!(graph, make_node(DataTask(96)), track = false) + d_v0_s0 = insert_node!(graph, make_node(DataTask(96)); track=false) + d_v1_s0 = insert_node!(graph, make_node(DataTask(96)); track=false) + d_v2_s1 = insert_node!(graph, make_node(DataTask(96)); track=false) + d_v3_s1 = insert_node!(graph, make_node(DataTask(96)); track=false) # v0 and v1 compute - v0 = insert_node!(graph, make_node(ComputeTaskQED_V()), track = false) - v1 = insert_node!(graph, make_node(ComputeTaskQED_V()), track = false) - v2 = insert_node!(graph, make_node(ComputeTaskQED_V()), track = false) - v3 = insert_node!(graph, make_node(ComputeTaskQED_V()), track = false) + v0 = insert_node!(graph, make_node(ComputeTaskQED_V()); track=false) + v1 = insert_node!(graph, make_node(ComputeTaskQED_V()); track=false) + v2 = insert_node!(graph, make_node(ComputeTaskQED_V()); track=false) + v3 = insert_node!(graph, make_node(ComputeTaskQED_V()); track=false) # data from uPhIn, uPhOut, uElIn, uElOut to v0 and v1 - d_uPhIn_v0 = insert_node!(graph, make_node(DataTask(96)), track = false) - d_uElIn_v0 = insert_node!(graph, make_node(DataTask(96)), track = false) - d_uPhOut_v1 = insert_node!(graph, make_node(DataTask(96)), track = false) - d_uElOut_v1 = insert_node!(graph, make_node(DataTask(96)), track = false) + d_uPhIn_v0 = insert_node!(graph, make_node(DataTask(96)); track=false) + d_uElIn_v0 = insert_node!(graph, make_node(DataTask(96)); track=false) + d_uPhOut_v1 = insert_node!(graph, make_node(DataTask(96)); track=false) + d_uElOut_v1 = insert_node!(graph, make_node(DataTask(96)); track=false) # data from uPhIn, uPhOut, uElIn, uElOut to v2 and v3 - d_uPhOut_v2 = insert_node!(graph, make_node(DataTask(96)), track = false) - d_uElIn_v2 = insert_node!(graph, make_node(DataTask(96)), track = false) - d_uPhIn_v3 = insert_node!(graph, make_node(DataTask(96)), track = false) - d_uElOut_v3 = insert_node!(graph, make_node(DataTask(96)), track = false) + d_uPhOut_v2 = insert_node!(graph, make_node(DataTask(96)); track=false) + d_uElIn_v2 = insert_node!(graph, make_node(DataTask(96)); track=false) + d_uPhIn_v3 = insert_node!(graph, make_node(DataTask(96)); track=false) + d_uElOut_v3 = insert_node!(graph, make_node(DataTask(96)); track=false) # uPhIn, uPhOut, uElIn and uElOut computes - uPhIn = insert_node!(graph, make_node(ComputeTaskQED_U()), track = false) - uPhOut = insert_node!(graph, make_node(ComputeTaskQED_U()), track = false) - uElIn = insert_node!(graph, make_node(ComputeTaskQED_U()), track = false) - uElOut = insert_node!(graph, make_node(ComputeTaskQED_U()), track = false) + uPhIn = insert_node!(graph, make_node(ComputeTaskQED_U()); track=false) + uPhOut = insert_node!(graph, make_node(ComputeTaskQED_U()); track=false) + uElIn = insert_node!(graph, make_node(ComputeTaskQED_U()); track=false) + uElOut = insert_node!(graph, make_node(ComputeTaskQED_U()); track=false) # data into U - d_uPhIn = insert_node!(graph, make_node(DataTask(16), "ki1"), track = false) - d_uPhOut = insert_node!(graph, make_node(DataTask(16), "ko1"), track = false) - d_uElIn = insert_node!(graph, make_node(DataTask(16), "ei1"), track = false) - d_uElOut = insert_node!(graph, make_node(DataTask(16), "eo1"), track = false) + d_uPhIn = insert_node!(graph, make_node(DataTask(16), "ki1"); track=false) + d_uPhOut = insert_node!(graph, make_node(DataTask(16), "ko1"); track=false) + d_uElIn = insert_node!(graph, make_node(DataTask(16), "ei1"); track=false) + d_uElOut = insert_node!(graph, make_node(DataTask(16), "eo1"); track=false) # now for all the edges - insert_edge!(graph, d_uPhIn, uPhIn, track = false) - insert_edge!(graph, d_uPhOut, uPhOut, track = false) - insert_edge!(graph, d_uElIn, uElIn, track = false) - insert_edge!(graph, d_uElOut, uElOut, track = false) + insert_edge!(graph, d_uPhIn, uPhIn; track=false) + insert_edge!(graph, d_uPhOut, uPhOut; track=false) + insert_edge!(graph, d_uElIn, uElIn; track=false) + insert_edge!(graph, d_uElOut, uElOut; track=false) - insert_edge!(graph, uPhIn, d_uPhIn_v0, track = false) - insert_edge!(graph, uPhOut, d_uPhOut_v1, track = false) - insert_edge!(graph, uElIn, d_uElIn_v0, track = false) - insert_edge!(graph, uElOut, d_uElOut_v1, track = false) + insert_edge!(graph, uPhIn, d_uPhIn_v0; track=false) + insert_edge!(graph, uPhOut, d_uPhOut_v1; track=false) + insert_edge!(graph, uElIn, d_uElIn_v0; track=false) + insert_edge!(graph, uElOut, d_uElOut_v1; track=false) - insert_edge!(graph, uPhIn, d_uPhIn_v3, track = false) - insert_edge!(graph, uPhOut, d_uPhOut_v2, track = false) - insert_edge!(graph, uElIn, d_uElIn_v2, track = false) - insert_edge!(graph, uElOut, d_uElOut_v3, track = false) + insert_edge!(graph, uPhIn, d_uPhIn_v3; track=false) + insert_edge!(graph, uPhOut, d_uPhOut_v2; track=false) + insert_edge!(graph, uElIn, d_uElIn_v2; track=false) + insert_edge!(graph, uElOut, d_uElOut_v3; track=false) - insert_edge!(graph, d_uPhIn_v0, v0, track = false) - insert_edge!(graph, d_uPhOut_v1, v1, track = false) - insert_edge!(graph, d_uElIn_v0, v0, track = false) - insert_edge!(graph, d_uElOut_v1, v1, track = false) + insert_edge!(graph, d_uPhIn_v0, v0; track=false) + insert_edge!(graph, d_uPhOut_v1, v1; track=false) + insert_edge!(graph, d_uElIn_v0, v0; track=false) + insert_edge!(graph, d_uElOut_v1, v1; track=false) - insert_edge!(graph, d_uPhIn_v3, v3, track = false) - insert_edge!(graph, d_uPhOut_v2, v2, track = false) - insert_edge!(graph, d_uElIn_v2, v2, track = false) - insert_edge!(graph, d_uElOut_v3, v3, track = false) + insert_edge!(graph, d_uPhIn_v3, v3; track=false) + insert_edge!(graph, d_uPhOut_v2, v2; track=false) + insert_edge!(graph, d_uElIn_v2, v2; track=false) + insert_edge!(graph, d_uElOut_v3, v3; track=false) - insert_edge!(graph, v0, d_v0_s0, track = false) - insert_edge!(graph, v1, d_v1_s0, track = false) - insert_edge!(graph, v2, d_v2_s1, track = false) - insert_edge!(graph, v3, d_v3_s1, track = false) + insert_edge!(graph, v0, d_v0_s0; track=false) + insert_edge!(graph, v1, d_v1_s0; track=false) + insert_edge!(graph, v2, d_v2_s1; track=false) + insert_edge!(graph, v3, d_v3_s1; track=false) - insert_edge!(graph, d_v0_s0, s0, track = false) - insert_edge!(graph, d_v1_s0, s0, track = false) + insert_edge!(graph, d_v0_s0, s0; track=false) + insert_edge!(graph, d_v1_s0, s0; track=false) - insert_edge!(graph, d_v2_s1, s1, track = false) - insert_edge!(graph, d_v3_s1, s1, track = false) + insert_edge!(graph, d_v2_s1, s1; track=false) + insert_edge!(graph, d_v3_s1, s1; track=false) - insert_edge!(graph, s0, d_s0_sum, track = false) - insert_edge!(graph, s1, d_s1_sum, track = false) + insert_edge!(graph, s0, d_s0_sum; track=false) + insert_edge!(graph, s1, d_s1_sum; track=false) - insert_edge!(graph, d_s0_sum, sum_node, track = false) - insert_edge!(graph, d_s1_sum, sum_node, track = false) + insert_edge!(graph, d_s0_sum, sum_node; track=false) + insert_edge!(graph, d_s1_sum, sum_node; track=false) - insert_edge!(graph, sum_node, d_exit, track = false) + insert_edge!(graph, sum_node, d_exit; track=false) input = [gen_process_input(process) for _ in 1:1000] @@ -305,9 +304,12 @@ end @test isapprox(compton_function.(input), compton_groundtruth.(input)) end -@testset "Equal results after optimization" for optimizer in - [ReductionOptimizer(), RandomWalkOptimizer(MersenneTwister(0))] - @testset "Process $proc_str" for proc_str in ["ke->ke", "kp->kp", "kk->ep", "ep->kk", "ke->kke", "ke->kkke"] +@testset "Equal results after optimization" for optimizer in [ + ReductionOptimizer(), RandomWalkOptimizer(MersenneTwister(0)) +] + @testset "Process $proc_str" for proc_str in [ + "ke->ke", "kp->kp", "kk->ep", "ep->kk", "ke->kke", "ke->kkke" + ] model = QEDModel() process = parse_process(proc_str, model) machine = Machine( -- 2.49.0 From 92f534f6bf061ca7024acc4bfb703b64aa94d988 Mon Sep 17 00:00:00 2001 From: Anton Reinhard Date: Thu, 4 Jul 2024 17:03:18 +0200 Subject: [PATCH 2/6] Fix congruent ph script --- examples/congruent_in_ph.jl | 156 ++++++++++++++++-------------------- 1 file changed, 67 insertions(+), 89 deletions(-) diff --git a/examples/congruent_in_ph.jl b/examples/congruent_in_ph.jl index a0639bf..1c91757 100644 --- a/examples/congruent_in_ph.jl +++ b/examples/congruent_in_ph.jl @@ -1,8 +1,10 @@ using MetagraphOptimization using QEDbase using QEDcore +using QEDprocesses using Random using UUIDs +using CSV RNG = Random.MersenneTwister(123) @@ -21,103 +23,79 @@ function mock_machine() ) end -function congruent_input(processDescription::QEDProcessDescription, omega::Number) +function congruent_input_momenta(processDescription::GenericQEDProcess, omega::Number) # generate an input sample for given e + nk -> e' + k' process, where the nk are equal massSum = 0 inputMasses = Vector{Float64}() - for (particle, n) in processDescription.inParticles - for _ in 1:n - massSum += mass(particle) - push!(inputMasses, mass(particle)) - end + for particle in incoming_particles(processDescription) + massSum += mass(particle) + push!(inputMasses, mass(particle)) end outputMasses = Vector{Float64}() - for (particle, n) in processDescription.outParticles - for _ in 1:n - massSum += mass(particle) - push!(outputMasses, mass(particle)) - end + for particle in outgoing_particles(processDescription) + massSum += mass(particle) + push!(outputMasses, mass(particle)) end - initialMomenta = [ - i == 1 ? SFourMomentum(1, 0, 0, 0) : SFourMomentum(omega, 0, 0, omega) for - i in 1:length(inputMasses) - ] + initial_momenta = + [i == 1 ? SFourMomentum(1, 0, 0, 0) : SFourMomentum(omega, 0, 0, omega) for i in 1:length(inputMasses)] # add some extra random mass to allow for some momentum - ss = sqrt(sum(initialMomenta) * sum(initialMomenta)) + ss = sqrt(sum(initial_momenta) * sum(initial_momenta)) + final_momenta = MetagraphOptimization.generate_physical_massive_moms(RNG, ss, outputMasses) - result = Vector{QEDProcessInput}() - sizehint!(result, 16) - - spin_pol_combinations = Iterators.product( - [SpinUp, SpinDown], [SpinUp, SpinDown], [PolX, PolY], [PolX, PolY] - ) - for (in_spin, out_spin, in_pol, out_pol) in spin_pol_combinations - - # get the electron first, then the n photons - particles = Vector{QEDParticle}() - - for (particle, n) in processDescription.inParticles - if particle <: FermionStateful - mom = initialMomenta[1] - push!(particles, particle(mom, in_spin())) - elseif particle <: PhotonStateful - for i in 1:n - mom = initialMomenta[i + 1] - push!(particles, particle(mom, in_pol())) - end - else - @assert false - end - end - - final_momenta = MetagraphOptimization.generate_physical_massive_moms( - RNG, ss, outputMasses - ) - index = 1 - for (particle, n) in processDescription.outParticles - for _ in 1:n - if particle <: FermionStateful - push!(particles, particle(final_momenta[index], out_spin())) - elseif particle <: PhotonStateful - push!(particles, particle(final_momenta[index], out_pol())) - end - index += 1 - end - end - - inFerms = MetagraphOptimization._svector_from_type( - processDescription, FermionStateful{Incoming,in_spin}, particles - ) - outFerms = MetagraphOptimization._svector_from_type( - processDescription, FermionStateful{Outgoing,out_spin}, particles - ) - inAntiferms = MetagraphOptimization._svector_from_type( - processDescription, AntiFermionStateful{Incoming,in_spin}, particles - ) - outAntiferms = MetagraphOptimization._svector_from_type( - processDescription, AntiFermionStateful{Outgoing,out_spin}, particles - ) - inPhotons = MetagraphOptimization._svector_from_type( - processDescription, PhotonStateful{Incoming,in_pol}, particles - ) - outPhotons = MetagraphOptimization._svector_from_type( - processDescription, PhotonStateful{Outgoing,out_pol}, particles - ) - - processInput = QEDProcessInput( - processDescription, - inFerms, - outFerms, - inAntiferms, - outAntiferms, - inPhotons, - outPhotons, - ) - - push!(result, processInput) - end - - return result + return (tuple(initial_momenta...), tuple(final_momenta...)) +end + +function build_psp(processDescription::GenericQEDProcess, momenta) + return PhaseSpacePoint( + processDescription, + PerturbativeQED(), + PhasespaceDefinition(SphericalCoordinateSystem(), ElectronRestFrame()), + momenta[1], + momenta[2], + ) +end + +n = 1000000 +photons = 4 +omega = 2e-3 + +# n is the number of incoming photons +# omega is the number + +println("Generating $n inputs for $photons photons, omega=$omega...") + +# temp process to generate momenta +temp_process = parse_process("k"^photons * "e->ke", QEDModel(), PolX(), SpinUp(), PolX(), SpinUp()) +input_momenta = [congruent_input_momenta(temp_process, omega) for _ in 1:n] +results = [0.0im for _ in 1:n] + +for (in_pol, in_spin, out_pol, out_spin) in + Iterators.product([PolX(), PolY()], [SpinUp(), SpinDown()], [PolX(), PolY()], [SpinUp(), SpinDown()]) + + print("Calculating for spin/pol config: $in_pol, $in_spin -> $out_pol, $out_spin... Preparing inputs... ") + process = parse_process("k"^photons * "e->ke", QEDModel(), in_pol, in_spin, out_pol, out_spin) + inputs = build_psp.(Ref(process), input_momenta) + + print("Preparing graph... ") + # prepare function + graph = gen_graph(process) + optimize_to_fixpoint!(ReductionOptimizer(), graph) + print("Preparing function... ") + func = get_compute_function(graph, process, mock_machine()) + + print("Calculating... ") + for i in 1:n + results[i] += func(inputs[i])^2 + end + println("Done.") +end + +println("Writing results") +open("$(photons)_congruent_photons_omega_$(omega).csv", "w") do f + println(f, "out_photon_momentum;out_electron_momentum;result") + for (momentum, result) in Iterators.zip(input_momenta, results) + println(f, "$(momentum[2][1]);$(momentum[2][2]);$(result)") + end end -- 2.49.0 From 813d40cd30b8623e39f13260adbab770f4c49853 Mon Sep 17 00:00:00 2001 From: Anton Reinhard Date: Thu, 4 Jul 2024 20:24:39 +0200 Subject: [PATCH 3/6] Working state --- examples/congruent_in_ph.jl | 176 ++++++++++++++++++----- src/code_gen/function.jl | 4 +- src/code_gen/tape_machine.jl | 14 +- src/code_gen/type.jl | 2 +- src/models/interface.jl | 4 +- src/models/physics_models/qed/compute.jl | 24 +++- 6 files changed, 176 insertions(+), 48 deletions(-) diff --git a/examples/congruent_in_ph.jl b/examples/congruent_in_ph.jl index 1c91757..1f8367d 100644 --- a/examples/congruent_in_ph.jl +++ b/examples/congruent_in_ph.jl @@ -5,6 +5,7 @@ using QEDprocesses using Random using UUIDs using CSV +using JLD2 RNG = Random.MersenneTwister(123) @@ -37,16 +38,66 @@ function congruent_input_momenta(processDescription::GenericQEDProcess, omega::N push!(outputMasses, mass(particle)) end - initial_momenta = - [i == 1 ? SFourMomentum(1, 0, 0, 0) : SFourMomentum(omega, 0, 0, omega) for i in 1:length(inputMasses)] + initial_momenta = [ + i == length(inputMasses) ? SFourMomentum(1, 0, 0, 0) : SFourMomentum(omega, 0, 0, omega) for + i in 1:length(inputMasses) + ] - # add some extra random mass to allow for some momentum ss = sqrt(sum(initial_momenta) * sum(initial_momenta)) final_momenta = MetagraphOptimization.generate_physical_massive_moms(RNG, ss, outputMasses) return (tuple(initial_momenta...), tuple(final_momenta...)) end +# cos_theta ∈ [-1, 1] and phi ∈ [0, 2π] +function congruent_input_momenta_scenario_2( + processDescription::GenericQEDProcess, + omega::Number, + theta::Number, + phi::Number, +) + # ------------- + # same as above + + # generate an input sample for given e + nk -> e' + k' process, where the nk are equal + massSum = 0 + inputMasses = Vector{Float64}() + for particle in incoming_particles(processDescription) + massSum += mass(particle) + push!(inputMasses, mass(particle)) + end + outputMasses = Vector{Float64}() + for particle in outgoing_particles(processDescription) + massSum += mass(particle) + push!(outputMasses, mass(particle)) + end + + initial_momenta = [ + i == length(inputMasses) ? SFourMomentum(1, 0, 0, 0) : SFourMomentum(omega, 0, 0, omega) for + i in 1:length(inputMasses) + ] + + ss = sqrt(sum(initial_momenta) * sum(initial_momenta)) + + # up to here + # ---------- + + # now calculate the final_momenta from omega, cos_theta and phi + n = number_particles(processDescription, ParticleStateful{Incoming, Photon, SFourMomentum}) + + cos_theta = cos(theta) + omega_prime = (n * omega) / (1 + n * omega * (1 - cos_theta)) + + k_prime = + omega_prime * SFourMomentum(1, sqrt(1 - cos_theta^2) * cos(phi), sqrt(1 - cos_theta^2) * sin(phi), cos_theta) + p_prime = sum(initial_momenta) - k_prime + + final_momenta = (k_prime, p_prime) + + return (tuple(initial_momenta...), tuple(final_momenta...)) + +end + function build_psp(processDescription::GenericQEDProcess, momenta) return PhaseSpacePoint( processDescription, @@ -57,45 +108,106 @@ function build_psp(processDescription::GenericQEDProcess, momenta) ) end +return 0 + +# scenario 2 +N = 1000 +M = 1000 + +thetas = collect(LinRange(0, 2π, N)) +phis = collect(LinRange(0, 2π, M)) + +for photons in [6] + # temp process to generate momenta + for omega in [2e-3, 2e-6] + println("Generating $(N*M) inputs for $photons photons (Scenario 2 grid walk)...") + temp_process = parse_process("k"^photons * "e->ke", QEDModel(), PolX(), SpinUp(), PolX(), SpinUp()) + + input_momenta = [ + congruent_input_momenta_scenario_2(temp_process, omega, theta, phi) for + (theta, phi) in Iterators.product(thetas, phis) + ] + results = [0.0 for _ in 1:length(input_momenta)] + + i = 1 + for (in_pol, in_spin, out_pol, out_spin) in + Iterators.product([PolX(), PolY()], [SpinUp(), SpinDown()], [PolX(), PolY()], [SpinUp(), SpinDown()]) + + print( + "[$i/16] Calculating for spin/pol config: $in_pol, $in_spin -> $out_pol, $out_spin... Preparing inputs... ", + ) + process = parse_process("k"^photons * "e->ke", QEDModel(), in_pol, in_spin, out_pol, out_spin) + inputs = build_psp.(Ref(process), input_momenta) + + print("Preparing graph... ") + # prepare function + graph = gen_graph(process) + optimize_to_fixpoint!(ReductionOptimizer(), graph) + print("Preparing function... ") + func = get_compute_function(graph, process, mock_machine()) + + print("Calculating... ") + Threads.@threads for i in 1:length(inputs) + results[i] += abs2(func(inputs[i])) + end + println("Done.") + i += 1 + end + + println("Writing results") + + out_ph_moms = getindex.(getindex.(input_momenta, 2), 1) + out_el_moms = getindex.(getindex.(input_momenta, 2), 2) + + @save "$(photons)_congruent_photons_omega_$(omega)_grid.jld2" out_ph_moms out_el_moms results + end +end + + n = 1000000 -photons = 4 -omega = 2e-3 # n is the number of incoming photons # omega is the number -println("Generating $n inputs for $photons photons, omega=$omega...") +for photons in [6] + # temp process to generate momenta + for omega in [2e-3, 2e-6] + println("Generating $n inputs for $photons photons...") + temp_process = parse_process("k"^photons * "e->ke", QEDModel(), PolX(), SpinUp(), PolX(), SpinUp()) -# temp process to generate momenta -temp_process = parse_process("k"^photons * "e->ke", QEDModel(), PolX(), SpinUp(), PolX(), SpinUp()) -input_momenta = [congruent_input_momenta(temp_process, omega) for _ in 1:n] -results = [0.0im for _ in 1:n] + input_momenta = [congruent_input_momenta(temp_process, omega) for _ in 1:n] + results = [0.0 for _ in 1:length(input_momenta)] -for (in_pol, in_spin, out_pol, out_spin) in - Iterators.product([PolX(), PolY()], [SpinUp(), SpinDown()], [PolX(), PolY()], [SpinUp(), SpinDown()]) + i = 1 + for (in_pol, in_spin, out_pol, out_spin) in + Iterators.product([PolX(), PolY()], [SpinUp(), SpinDown()], [PolX(), PolY()], [SpinUp(), SpinDown()]) - print("Calculating for spin/pol config: $in_pol, $in_spin -> $out_pol, $out_spin... Preparing inputs... ") - process = parse_process("k"^photons * "e->ke", QEDModel(), in_pol, in_spin, out_pol, out_spin) - inputs = build_psp.(Ref(process), input_momenta) + print( + "[$i/16] Calculating for spin/pol config: $in_pol, $in_spin -> $out_pol, $out_spin... Preparing inputs... ", + ) + process = parse_process("k"^photons * "e->ke", QEDModel(), in_pol, in_spin, out_pol, out_spin) + inputs = build_psp.(Ref(process), input_momenta) - print("Preparing graph... ") - # prepare function - graph = gen_graph(process) - optimize_to_fixpoint!(ReductionOptimizer(), graph) - print("Preparing function... ") - func = get_compute_function(graph, process, mock_machine()) + print("Preparing graph... ") + # prepare function + graph = gen_graph(process) + optimize_to_fixpoint!(ReductionOptimizer(), graph) + print("Preparing function... ") + func = get_compute_function(graph, process, mock_machine()) - print("Calculating... ") - for i in 1:n - results[i] += func(inputs[i])^2 - end - println("Done.") -end + print("Calculating... ") + Threads.@threads for i in 1:n + results[i] += abs2(func(inputs[i])) + end + println("Done.") + i += 1 + end -println("Writing results") -open("$(photons)_congruent_photons_omega_$(omega).csv", "w") do f - println(f, "out_photon_momentum;out_electron_momentum;result") - for (momentum, result) in Iterators.zip(input_momenta, results) - println(f, "$(momentum[2][1]);$(momentum[2][2]);$(result)") + println("Writing results") + + out_ph_moms = getindex.(getindex.(input_momenta, 2), 1) + out_el_moms = getindex.(getindex.(input_momenta, 2), 2) + + @save "$(photons)_congruent_photons_omega_$(omega).jld2" out_ph_moms out_el_moms results end end diff --git a/src/code_gen/function.jl b/src/code_gen/function.jl index fc9c2ba..50b3df6 100644 --- a/src/code_gen/function.jl +++ b/src/code_gen/function.jl @@ -7,7 +7,7 @@ function get_compute_function(graph::DAG, instance, machine::Machine) tape = gen_tape(graph, instance, machine) initCaches = Expr(:block, tape.initCachesCode...) - assignInputs = Expr(:block, expr_from_fc.(tape.inputAssignCode)...) + assignInputs = Expr(:block, tape.inputAssignCode...) code = Expr(:block, expr_from_fc.(tape.computeCode)...) functionId = to_var_name(UUIDs.uuid1(rng[1])) @@ -30,7 +30,7 @@ function get_cuda_kernel(graph::DAG, instance, machine::Machine) tape = gen_tape(graph, instance, machine) initCaches = Expr(:block, tape.initCachesCode...) - assignInputs = Expr(:block, expr_from_fc.(tape.inputAssignCode)...) + assignInputs = Expr(:block, tape.inputAssignCode...) code = Expr(:block, expr_from_fc.(tape.computeCode)...) functionId = to_var_name(UUIDs.uuid1(rng[1])) diff --git a/src/code_gen/tape_machine.jl b/src/code_gen/tape_machine.jl index db8ade7..010fd6f 100644 --- a/src/code_gen/tape_machine.jl +++ b/src/code_gen/tape_machine.jl @@ -75,7 +75,7 @@ end inputSymbols::Dict{String, Vector{Symbol}}, instance::AbstractProblemInstance, machine::Machine, - problemInputSymbol::Symbol = :input, + problemInputSymbol::Symbol = :data_input, ) Return a `Vector{Expr}` doing the input assignments from the given `problemInputSymbol` onto the `inputSymbols`. @@ -84,9 +84,9 @@ function gen_input_assignment_code( inputSymbols::Dict{String, Vector{Symbol}}, instance, machine::Machine, - problemInputSymbol::Symbol = :input, + problemInputSymbol::Symbol = :data_input, ) - assignInputs = Vector{FunctionCall}() + assignInputs = Vector{Expr}() for (name, symbols) in inputSymbols # make a function for this, since we can't use anonymous functions in the FunctionCall @@ -94,7 +94,9 @@ function gen_input_assignment_code( device = entry_device(machine) push!( assignInputs, - FunctionCall(input_expr, SVector{1, Any}(name), SVector{1, Symbol}(problemInputSymbol), symbol, device), + Meta.parse( + "$(eval(gen_access_expr(device, symbol))) = $(input_expr(instance, name, problemInputSymbol))", + ), ) end end @@ -126,7 +128,7 @@ function gen_tape(graph::DAG, instance, machine::Machine, scheduler::AbstractSch outSym = Symbol(to_var_name(get_exit_node(graph).id)) initCaches = gen_cache_init_code(machine) - assignInputs = gen_input_assignment_code(inputSyms, instance, machine, :input) + assignInputs = gen_input_assignment_code(inputSyms, instance, machine, :data_input) return Tape{input_type(instance)}(initCaches, assignInputs, schedule, inputSyms, outSym, Dict(), instance, machine) end @@ -140,7 +142,7 @@ For implementation reasons, this disregards the set [`CacheStrategy`](@ref) of t """ function execute_tape(tape::Tape, input) cache = Dict{Symbol, Any}() - cache[:input] = input + cache[:data_input] = input # simply execute all the code snippets here @assert typeof(input) == input_type(tape.instance) # TODO: `@assert` that input fits the tape.instance diff --git a/src/code_gen/type.jl b/src/code_gen/type.jl index db420fd..1b079dd 100644 --- a/src/code_gen/type.jl +++ b/src/code_gen/type.jl @@ -11,7 +11,7 @@ TODO: update docs """ struct Tape{INPUT} initCachesCode::Vector{Expr} - inputAssignCode::Vector{FunctionCall} + inputAssignCode::Vector{Expr} computeCode::Vector{FunctionCall} inputSymbols::Dict{String, Vector{Symbol}} outputSymbol::Symbol diff --git a/src/models/interface.jl b/src/models/interface.jl index 8bf7514..8ca05b2 100644 --- a/src/models/interface.jl +++ b/src/models/interface.jl @@ -39,8 +39,8 @@ Generate the [`DAG`](@ref) for the given [`AbstractProblemInstance`](@ref). Ever function graph end """ - input_expr(::AbstractProblemInstance, name::String, input) + input_expr(instance::AbstractProblemInstance, name::String, input_symbol::Symbol) -For a given [`AbstractProblemInstance`](@ref), the problem input (of type `input_type(...)`) and an entry node name, return a that specific input value from the input symbol. +For the given [`AbstractProblemInstance`](@ref), the entry node name, and the symbol of the problem input (where a variable of type `input_type(...)` will exist), return an `Expr` that gets that specific input value from the input symbol. """ function input_expr end diff --git a/src/models/physics_models/qed/compute.jl b/src/models/physics_models/qed/compute.jl index 239c805..01b1459 100644 --- a/src/models/physics_models/qed/compute.jl +++ b/src/models/physics_models/qed/compute.jl @@ -1,12 +1,26 @@ using StaticArrays -function input_expr(name::String, psp::PhaseSpacePoint) +construction_string(::Incoming) = "Incoming()" +construction_string(::Outgoing) = "Outgoing()" + +construction_string(::Electron) = "Electron()" +construction_string(::Positron) = "Positron()" +construction_string(::Photon) = "Photon()" + +construction_string(::PolX) = "PolX()" +construction_string(::PolY) = "PolY()" +construction_string(::SpinUp) = "SpinUp()" +construction_string(::SpinDown) = "SpinDown()" + +function input_expr(instance::GenericQEDProcess, name::String, psp_symbol::Symbol) (type, index) = type_index_from_name(QEDModel(), name) - return ParticleValueSP( - type(momentum(psp, particle_direction(type), particle_species(type), index)), - 0.0im, - spin_or_pol(process(psp), type, index), + return Meta.parse( + "ParticleValueSP( + $type(momentum($psp_symbol, $(construction_string(particle_direction(type))), $(construction_string(particle_species(type))), $index)), + 0.0im, + $(construction_string(spin_or_pol(instance, type, index))), +)", ) end -- 2.49.0 From b784720859c969196fbbef5dff4b7cd03b305c55 Mon Sep 17 00:00:00 2001 From: Anton Reinhard Date: Thu, 4 Jul 2024 20:40:22 +0200 Subject: [PATCH 4/6] Minor fixes and commints --- examples/congruent_in_ph.jl | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/examples/congruent_in_ph.jl b/examples/congruent_in_ph.jl index 1f8367d..24b430e 100644 --- a/examples/congruent_in_ph.jl +++ b/examples/congruent_in_ph.jl @@ -26,15 +26,12 @@ end function congruent_input_momenta(processDescription::GenericQEDProcess, omega::Number) # generate an input sample for given e + nk -> e' + k' process, where the nk are equal - massSum = 0 inputMasses = Vector{Float64}() for particle in incoming_particles(processDescription) - massSum += mass(particle) push!(inputMasses, mass(particle)) end outputMasses = Vector{Float64}() for particle in outgoing_particles(processDescription) - massSum += mass(particle) push!(outputMasses, mass(particle)) end @@ -49,7 +46,7 @@ function congruent_input_momenta(processDescription::GenericQEDProcess, omega::N return (tuple(initial_momenta...), tuple(final_momenta...)) end -# cos_theta ∈ [-1, 1] and phi ∈ [0, 2π] +# theta ∈ [0, 2π] and phi ∈ [0, 2π] function congruent_input_momenta_scenario_2( processDescription::GenericQEDProcess, omega::Number, @@ -60,15 +57,12 @@ function congruent_input_momenta_scenario_2( # same as above # generate an input sample for given e + nk -> e' + k' process, where the nk are equal - massSum = 0 inputMasses = Vector{Float64}() for particle in incoming_particles(processDescription) - massSum += mass(particle) push!(inputMasses, mass(particle)) end outputMasses = Vector{Float64}() for particle in outgoing_particles(processDescription) - massSum += mass(particle) push!(outputMasses, mass(particle)) end @@ -163,7 +157,7 @@ for photons in [6] end end - +# scenario n = 1000000 # n is the number of incoming photons -- 2.49.0 From dee44dad66332603d1706f65263a20cfabd2af16 Mon Sep 17 00:00:00 2001 From: Anton Reinhard Date: Fri, 5 Jul 2024 02:15:03 +0200 Subject: [PATCH 5/6] Small fixes --- examples/congruent_in_ph.jl | 27 +++++++++++++++-------- src/models/physics_models/qed/diagrams.jl | 26 +++++++++------------- 2 files changed, 29 insertions(+), 24 deletions(-) diff --git a/examples/congruent_in_ph.jl b/examples/congruent_in_ph.jl index 24b430e..19acc24 100644 --- a/examples/congruent_in_ph.jl +++ b/examples/congruent_in_ph.jl @@ -1,3 +1,5 @@ +ENV["UCX_ERROR_SIGNALS"] = "SIGILL,SIGBUS,SIGFPE" + using MetagraphOptimization using QEDbase using QEDcore @@ -102,7 +104,8 @@ function build_psp(processDescription::GenericQEDProcess, momenta) ) end -return 0 +# hack to fix stacksize for threading +with_stacksize(f, n) = fetch(schedule(Task(f, n))) # scenario 2 N = 1000 @@ -111,7 +114,7 @@ M = 1000 thetas = collect(LinRange(0, 2π, N)) phis = collect(LinRange(0, 2π, M)) -for photons in [6] +for photons in 1:6 # temp process to generate momenta for omega in [2e-3, 2e-6] println("Generating $(N*M) inputs for $photons photons (Scenario 2 grid walk)...") @@ -121,7 +124,8 @@ for photons in [6] congruent_input_momenta_scenario_2(temp_process, omega, theta, phi) for (theta, phi) in Iterators.product(thetas, phis) ] - results = [0.0 for _ in 1:length(input_momenta)] + results = Array{Float64}(undef, size(input_momenta)) + fill!(results, 0.0) i = 1 for (in_pol, in_spin, out_pol, out_spin) in @@ -134,15 +138,17 @@ for photons in [6] inputs = build_psp.(Ref(process), input_momenta) print("Preparing graph... ") - # prepare function graph = gen_graph(process) optimize_to_fixpoint!(ReductionOptimizer(), graph) print("Preparing function... ") func = get_compute_function(graph, process, mock_machine()) + func(inputs[1]) print("Calculating... ") - Threads.@threads for i in 1:length(inputs) - results[i] += abs2(func(inputs[i])) + Threads.@threads for i in 1:N + Threads.@threads for j in 1:M + return results[i, j] += abs2(func(inputs[i, j])) + end end println("Done.") i += 1 @@ -157,20 +163,23 @@ for photons in [6] end end -# scenario +exit(0) + +# scenario 1 (disabled) n = 1000000 # n is the number of incoming photons # omega is the number -for photons in [6] +for photons in 1:6 # temp process to generate momenta for omega in [2e-3, 2e-6] println("Generating $n inputs for $photons photons...") temp_process = parse_process("k"^photons * "e->ke", QEDModel(), PolX(), SpinUp(), PolX(), SpinUp()) input_momenta = [congruent_input_momenta(temp_process, omega) for _ in 1:n] - results = [0.0 for _ in 1:length(input_momenta)] + results = Array{Float64}(undef, size(input_momenta)) + fill!(results, 0.0) i = 1 for (in_pol, in_spin, out_pol, out_spin) in diff --git a/src/models/physics_models/qed/diagrams.jl b/src/models/physics_models/qed/diagrams.jl index 27082e5..0b79920 100644 --- a/src/models/physics_models/qed/diagrams.jl +++ b/src/models/physics_models/qed/diagrams.jl @@ -504,18 +504,15 @@ function gen_compton_diagram_from_order(order::Vector{Int}, inFerm, outFerm, n:: return new_diagram end -#= """ gen_compton_diagram_from_order_one_side(order::Vector{Int}, inFerm, outFerm, n::Int, m::Int) Helper function for [`gen_compton_diagrams`](@Ref). Generates a single diagram for the given order and n input and m output photons. """ -function gen_compton_diagram_from_order_one_side( - order::Vector{Int}, inFerm, outFerm, n::Int, m::Int -) +function gen_compton_diagram_from_order_one_side(order::Vector{Int}, inFerm, outFerm, n::Int, m::Int) photons = vcat( - [FeynmanParticle(ParticleStateful{Incoming, Photon}, i) for i in 1:n], - [FeynmanParticle(ParticleStateful{Outgoing, Photon}, i) for i in 1:m], + [FeynmanParticle(ParticleStateful{Incoming, Photon, SFourMomentum}, i) for i in 1:n], + [FeynmanParticle(ParticleStateful{Outgoing, Photon, SFourMomentum}, i) for i in 1:m], ) new_diagram = FeynmanDiagram( @@ -523,10 +520,10 @@ function gen_compton_diagram_from_order_one_side( missing, [inFerm, outFerm, photons...], Dict{Type, Int64}( - ParticleStateful{Incoming, Electron} => 1, - ParticleStateful{Outgoing, Electron} => 1, - ParticleStateful{Incoming, Photon} => n, - ParticleStateful{Outgoing, Photon} => m, + ParticleStateful{Incoming, Electron, SFourMomentum} => 1, + ParticleStateful{Outgoing, Electron, SFourMomentum} => 1, + ParticleStateful{Incoming, Photon, SFourMomentum} => n, + ParticleStateful{Outgoing, Photon, SFourMomentum} => m, ), ) @@ -538,9 +535,9 @@ function gen_compton_diagram_from_order_one_side( while left_index <= right_index # left side v_left = FeynmanVertex( - FeynmanParticle(ParticleStateful{Incoming, Electron}, iterations), + FeynmanParticle(ParticleStateful{Incoming, Electron, SFourMomentum}, iterations), photons[order[left_index]], - FeynmanParticle(ParticleStateful{Incoming, Electron}, iterations + 1), + FeynmanParticle(ParticleStateful{Incoming, Electron, SFourMomentum}, iterations + 1), ) left_index += 1 add_vertex!(new_diagram, v_left) @@ -553,9 +550,9 @@ function gen_compton_diagram_from_order_one_side( if (iterations == 1) # right side v_right = FeynmanVertex( - FeynmanParticle(ParticleStateful{Outgoing, Electron}, iterations), + FeynmanParticle(ParticleStateful{Outgoing, Electron, SFourMomentum}, iterations), photons[order[right_index]], - FeynmanParticle(ParticleStateful{Outgoing, Electron}, iterations + 1), + FeynmanParticle(ParticleStateful{Outgoing, Electron, SFourMomentum}, iterations + 1), ) right_index -= 1 add_vertex!(new_diagram, v_right) @@ -569,7 +566,6 @@ function gen_compton_diagram_from_order_one_side( add_tie!(new_diagram, FeynmanTie(ps[1], ps[2])) return new_diagram end -=# """ gen_compton_diagrams(n::Int, m::Int) -- 2.49.0 From 4eee23f08121fd383bae437901f5389cf4a1fa2f Mon Sep 17 00:00:00 2001 From: Anton Reinhard Date: Tue, 9 Jul 2024 23:19:25 +0200 Subject: [PATCH 6/6] Run on GPU --- examples/congruent_in_ph.jl | 152 ++++++++++------------- src/MetagraphOptimization.jl | 3 - src/QEDprocesses_patch.jl | 46 ------- src/models/physics_models/qed/compute.jl | 2 +- 4 files changed, 66 insertions(+), 137 deletions(-) diff --git a/examples/congruent_in_ph.jl b/examples/congruent_in_ph.jl index 19acc24..8b805b7 100644 --- a/examples/congruent_in_ph.jl +++ b/examples/congruent_in_ph.jl @@ -6,8 +6,13 @@ using QEDcore using QEDprocesses using Random using UUIDs + +using CUDA + +using NamedDims using CSV using JLD2 +using FlexiMaps RNG = Random.MersenneTwister(123) @@ -79,7 +84,7 @@ function congruent_input_momenta_scenario_2( # ---------- # now calculate the final_momenta from omega, cos_theta and phi - n = number_particles(processDescription, ParticleStateful{Incoming, Photon, SFourMomentum}) + n = number_particles(processDescription, Incoming(), Photon()) cos_theta = cos(theta) omega_prime = (n * omega) / (1 + n * omega * (1 - cos_theta)) @@ -108,109 +113,82 @@ end with_stacksize(f, n) = fetch(schedule(Task(f, n))) # scenario 2 -N = 1000 -M = 1000 +N = 1024 # thetas +M = 1024 # phis +K = 64 # omegas thetas = collect(LinRange(0, 2π, N)) phis = collect(LinRange(0, 2π, M)) +omegas = collect(maprange(log, 2e-2, 2e-7, K)) -for photons in 1:6 +for photons in 1:5 # temp process to generate momenta - for omega in [2e-3, 2e-6] - println("Generating $(N*M) inputs for $photons photons (Scenario 2 grid walk)...") - temp_process = parse_process("k"^photons * "e->ke", QEDModel(), PolX(), SpinUp(), PolX(), SpinUp()) + println("Generating $(K*N*M) inputs for $photons photons (Scenario 2 grid walk)...") + temp_process = parse_process("k"^photons * "e->ke", QEDModel(), PolX(), SpinUp(), PolX(), SpinUp()) - input_momenta = [ - congruent_input_momenta_scenario_2(temp_process, omega, theta, phi) for - (theta, phi) in Iterators.product(thetas, phis) - ] - results = Array{Float64}(undef, size(input_momenta)) - fill!(results, 0.0) + input_momenta = + Array{typeof(congruent_input_momenta_scenario_2(temp_process, omegas[1], thetas[1], phis[1]))}(undef, (K, N, M)) - i = 1 - for (in_pol, in_spin, out_pol, out_spin) in - Iterators.product([PolX(), PolY()], [SpinUp(), SpinDown()], [PolX(), PolY()], [SpinUp(), SpinDown()]) + Threads.@threads for k in 1:K + Threads.@threads for i in 1:N + Threads.@threads for j in 1:M + input_momenta[k, i, j] = congruent_input_momenta_scenario_2(temp_process, omegas[k], thetas[i], phis[j]) + end + end + end - print( - "[$i/16] Calculating for spin/pol config: $in_pol, $in_spin -> $out_pol, $out_spin... Preparing inputs... ", - ) - process = parse_process("k"^photons * "e->ke", QEDModel(), in_pol, in_spin, out_pol, out_spin) - inputs = build_psp.(Ref(process), input_momenta) - print("Preparing graph... ") - graph = gen_graph(process) - optimize_to_fixpoint!(ReductionOptimizer(), graph) - print("Preparing function... ") - func = get_compute_function(graph, process, mock_machine()) - func(inputs[1]) + cu_results = CuArray{Float64}(undef, size(input_momenta)) + fill!(cu_results, 0.0) - print("Calculating... ") + i = 1 + for (in_pol, in_spin, out_pol, out_spin) in + Iterators.product([PolX(), PolY()], [SpinUp(), SpinDown()], [PolX(), PolY()], [SpinUp(), SpinDown()]) + + print( + "[$i/16] Calculating for spin/pol config: $in_pol, $in_spin -> $out_pol, $out_spin... Preparing inputs... ", + ) + process = parse_process("k"^photons * "e->ke", QEDModel(), in_pol, in_spin, out_pol, out_spin) + + inputs = Array{typeof(build_psp(process, input_momenta[1, 1, 1]))}(undef, (K, N, M)) + #println("input_momenta: $input_momenta") + Threads.@threads for k in 1:K Threads.@threads for i in 1:N Threads.@threads for j in 1:M - return results[i, j] += abs2(func(inputs[i, j])) + inputs[k, i, j] = build_psp(process, input_momenta[k, i, j]) end end - println("Done.") - i += 1 end + cu_inputs = CuArray(inputs) - println("Writing results") + print("Preparing graph... ") + graph = gen_graph(process) + optimize_to_fixpoint!(ReductionOptimizer(), graph) + print("Preparing function... ") + kernel! = get_cuda_kernel(graph, process, mock_machine()) + #func = get_compute_function(graph, process, mock_machine()) - out_ph_moms = getindex.(getindex.(input_momenta, 2), 1) - out_el_moms = getindex.(getindex.(input_momenta, 2), 2) + print("Calculating... ") + ts = 32 + bs = Int64(length(cu_inputs) / 32) - @save "$(photons)_congruent_photons_omega_$(omega)_grid.jld2" out_ph_moms out_el_moms results - end -end - -exit(0) - -# scenario 1 (disabled) -n = 1000000 - -# n is the number of incoming photons -# omega is the number - -for photons in 1:6 - # temp process to generate momenta - for omega in [2e-3, 2e-6] - println("Generating $n inputs for $photons photons...") - temp_process = parse_process("k"^photons * "e->ke", QEDModel(), PolX(), SpinUp(), PolX(), SpinUp()) - - input_momenta = [congruent_input_momenta(temp_process, omega) for _ in 1:n] - results = Array{Float64}(undef, size(input_momenta)) - fill!(results, 0.0) - - i = 1 - for (in_pol, in_spin, out_pol, out_spin) in - Iterators.product([PolX(), PolY()], [SpinUp(), SpinDown()], [PolX(), PolY()], [SpinUp(), SpinDown()]) - - print( - "[$i/16] Calculating for spin/pol config: $in_pol, $in_spin -> $out_pol, $out_spin... Preparing inputs... ", - ) - process = parse_process("k"^photons * "e->ke", QEDModel(), in_pol, in_spin, out_pol, out_spin) - inputs = build_psp.(Ref(process), input_momenta) - - print("Preparing graph... ") - # prepare function - graph = gen_graph(process) - optimize_to_fixpoint!(ReductionOptimizer(), graph) - print("Preparing function... ") - func = get_compute_function(graph, process, mock_machine()) - - print("Calculating... ") - Threads.@threads for i in 1:n - results[i] += abs2(func(inputs[i])) - end - println("Done.") - i += 1 - end - - println("Writing results") - - out_ph_moms = getindex.(getindex.(input_momenta, 2), 1) - out_el_moms = getindex.(getindex.(input_momenta, 2), 2) - - @save "$(photons)_congruent_photons_omega_$(omega).jld2" out_ph_moms out_el_moms results + outputs = CuArray{ComplexF64}(undef, size(cu_inputs)) + + @cuda threads = ts blocks = bs always_inline = true kernel!(cu_inputs, outputs, length(cu_inputs)) + CUDA.device_synchronize() + cu_results += abs2.(outputs) + + println("Done.") + i += 1 end + + println("Writing results") + + out_ph_moms = getindex.(getindex.(input_momenta, 2), 1) + out_el_moms = getindex.(getindex.(input_momenta, 2), 2) + + results = NamedDimsArray{(:omegas, :thetas, :phis)}(Array(cu_results)) + println("Named results array: $(typeof(results))") + + @save "$(photons)_congruent_photons_grid.jld2" omegas thetas phis results end diff --git a/src/MetagraphOptimization.jl b/src/MetagraphOptimization.jl index 77a85eb..874de89 100644 --- a/src/MetagraphOptimization.jl +++ b/src/MetagraphOptimization.jl @@ -100,9 +100,6 @@ export ==, in, show, isempty, delete!, length export bytes_to_human_readable -# TODO: this is probably not good -import QEDprocesses.compute - import Base.length import Base.show import Base.== diff --git a/src/QEDprocesses_patch.jl b/src/QEDprocesses_patch.jl index 1eba98c..282d713 100644 --- a/src/QEDprocesses_patch.jl +++ b/src/QEDprocesses_patch.jl @@ -1,25 +1,5 @@ # 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}, @@ -43,29 +23,3 @@ end ) 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 diff --git a/src/models/physics_models/qed/compute.jl b/src/models/physics_models/qed/compute.jl index 01b1459..3650955 100644 --- a/src/models/physics_models/qed/compute.jl +++ b/src/models/physics_models/qed/compute.jl @@ -17,7 +17,7 @@ function input_expr(instance::GenericQEDProcess, name::String, psp_symbol::Symbo return Meta.parse( "ParticleValueSP( - $type(momentum($psp_symbol, $(construction_string(particle_direction(type))), $(construction_string(particle_species(type))), $index)), + $type(momentum($psp_symbol, $(construction_string(particle_direction(type))), $(construction_string(particle_species(type))), Val($index))), 0.0im, $(construction_string(spin_or_pol(instance, type, index))), )", -- 2.49.0