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

Working state
Some checks failed
MetagraphOptimization_CI / docs (push) Failing after 1m40s
Details
MetagraphOptimization_CI / test (push) Failing after 1m47s
Details

Browse Source
This commit is contained in:
Anton Reinhard
2024-07-04 20:24:39 +02:00
parent 92f534f6bf
commit 813d40cd30
6 changed files with 176 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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