rubydragon/metagraphoptimization.jl
1
0
Fork 0
Code Issues 6 Pull Requests Actions Packages Projects Releases Wiki Activity

Congruent incoming photons implementation #11

Merged
rubydragon merged 9 commits from congruent_in_ph into refactor 2024-07-10 14:04:13 +02:00
Conversation 0 Commits 9 Files Changed 23 +176 -48
6 changed files with 176 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit 813d40cd30 - Show all commits

176
examples/congruent_in_ph.jl
View File

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

4
src/code_gen/function.jl
View File

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

14
src/code_gen/tape_machine.jl
View File

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

2
src/code_gen/type.jl
View File

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

4
src/models/interface.jl
View File

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

24
src/models/physics_models/qed/compute.jl
View File

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