Congruent incoming photons implementation (#11)
Co-authored-by: Rubydragon <anton.reinhard@proton.me> Reviewed-on: #11
This commit is contained in:
194
examples/congruent_in_ph.jl
Normal file
194
examples/congruent_in_ph.jl
Normal file
@@ -0,0 +1,194 @@
|
||||
ENV["UCX_ERROR_SIGNALS"] = "SIGILL,SIGBUS,SIGFPE"
|
||||
|
||||
using MetagraphOptimization
|
||||
using QEDbase
|
||||
using QEDcore
|
||||
using QEDprocesses
|
||||
using Random
|
||||
using UUIDs
|
||||
|
||||
using CUDA
|
||||
|
||||
using NamedDims
|
||||
using CSV
|
||||
using JLD2
|
||||
using FlexiMaps
|
||||
|
||||
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_momenta(processDescription::GenericQEDProcess, omega::Number)
|
||||
# generate an input sample for given e + nk -> e' + k' process, where the nk are equal
|
||||
inputMasses = Vector{Float64}()
|
||||
for particle in incoming_particles(processDescription)
|
||||
push!(inputMasses, mass(particle))
|
||||
end
|
||||
outputMasses = Vector{Float64}()
|
||||
for particle in outgoing_particles(processDescription)
|
||||
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))
|
||||
final_momenta = MetagraphOptimization.generate_physical_massive_moms(RNG, ss, outputMasses)
|
||||
|
||||
return (tuple(initial_momenta...), tuple(final_momenta...))
|
||||
end
|
||||
|
||||
# theta ∈ [0, 2π] 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
|
||||
inputMasses = Vector{Float64}()
|
||||
for particle in incoming_particles(processDescription)
|
||||
push!(inputMasses, mass(particle))
|
||||
end
|
||||
outputMasses = Vector{Float64}()
|
||||
for particle in outgoing_particles(processDescription)
|
||||
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, Incoming(), Photon())
|
||||
|
||||
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,
|
||||
PerturbativeQED(),
|
||||
PhasespaceDefinition(SphericalCoordinateSystem(), ElectronRestFrame()),
|
||||
momenta[1],
|
||||
momenta[2],
|
||||
)
|
||||
end
|
||||
|
||||
# hack to fix stacksize for threading
|
||||
with_stacksize(f, n) = fetch(schedule(Task(f, n)))
|
||||
|
||||
# scenario 2
|
||||
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:5
|
||||
# temp process to generate momenta
|
||||
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 =
|
||||
Array{typeof(congruent_input_momenta_scenario_2(temp_process, omegas[1], thetas[1], phis[1]))}(undef, (K, N, M))
|
||||
|
||||
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
|
||||
|
||||
|
||||
cu_results = CuArray{Float64}(undef, size(input_momenta))
|
||||
fill!(cu_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 = 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
|
||||
inputs[k, i, j] = build_psp(process, input_momenta[k, i, j])
|
||||
end
|
||||
end
|
||||
end
|
||||
cu_inputs = CuArray(inputs)
|
||||
|
||||
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())
|
||||
|
||||
print("Calculating... ")
|
||||
ts = 32
|
||||
bs = Int64(length(cu_inputs) / 32)
|
||||
|
||||
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
|
||||
Reference in New Issue
Block a user