From 813d40cd30b8623e39f13260adbab770f4c49853 Mon Sep 17 00:00:00 2001 From: Anton Reinhard Date: Thu, 4 Jul 2024 20:24:39 +0200 Subject: [PATCH] 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