Fix congruent ph script

.gitattributes

@@ -1,5 +1,3 @@
 input/AB->ABBBBBBBBB.txt filter=lfs diff=lfs merge=lfs -text
 input/AB->ABBBBBBB.txt filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs
-*.gif filter=lfs diff=lfs merge=lfs
-*.jld2 filter=lfs diff=lfs merge=lfs
+*.zip filter=lfs diff=lfs merge=lfs -text

examples/congruent_in_ph.jl

@@ -1,18 +1,10 @@
-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)
 
@@ -33,72 +25,28 @@ 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
 
-    initial_momenta = [
-        i == length(inputMasses) ? 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(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,
@@ -109,86 +57,45 @@ function build_psp(processDescription::GenericQEDProcess, momenta)
     )
 end
 
-# hack to fix stacksize for threading
-with_stacksize(f, n) = fetch(schedule(Task(f, n)))
+n = 1000000
+photons = 4
+omega = 2e-3
 
-# scenario 2
-N = 1024    # thetas
-M = 1024    # phis
-K = 64      # omegas
+# n is the number of incoming photons
+# omega is the number
 
-thetas = collect(LinRange(0, 2π, N))
-phis = collect(LinRange(0, 2π, M))
-omegas = collect(maprange(log, 2e-2, 2e-7, K))
+println("Generating $n inputs for $photons photons, omega=$omega...")
 
-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())
+# 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 =
-        Array{typeof(congruent_input_momenta_scenario_2(temp_process, omegas[1], thetas[1], phis[1]))}(undef, (K, N, M))
+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
+    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
-
-
-    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

src/MetagraphOptimization.jl

@@ -100,6 +100,9 @@ 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.==

src/QEDprocesses_patch.jl

@@ -1,5 +1,25 @@
 # 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},
@@ -23,3 +43,29 @@ 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

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, tape.inputAssignCode...)
+    assignInputs = Expr(:block, expr_from_fc.(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, tape.inputAssignCode...)
+    assignInputs = Expr(:block, expr_from_fc.(tape.inputAssignCode)...)
     code = Expr(:block, expr_from_fc.(tape.computeCode)...)
 
     functionId = to_var_name(UUIDs.uuid1(rng[1]))

src/code_gen/tape_machine.jl

@@ -75,7 +75,7 @@ end
         inputSymbols::Dict{String, Vector{Symbol}},
         instance::AbstractProblemInstance,
         machine::Machine,
-        problemInputSymbol::Symbol = :data_input,
+        problemInputSymbol::Symbol = :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 = :data_input,
+    problemInputSymbol::Symbol = :input,
 )
-    assignInputs = Vector{Expr}()
+    assignInputs = Vector{FunctionCall}()
     for (name, symbols) in inputSymbols
         # make a function for this, since we can't use anonymous functions in the FunctionCall
 
@@ -94,9 +94,7 @@ function gen_input_assignment_code(
             device = entry_device(machine)
             push!(
                 assignInputs,
-                Meta.parse(
-                    "$(eval(gen_access_expr(device, symbol))) = $(input_expr(instance, name, problemInputSymbol))",
-                ),
+                FunctionCall(input_expr, SVector{1, Any}(name), SVector{1, Symbol}(problemInputSymbol), symbol, device),
             )
         end
     end
@@ -128,7 +126,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, :data_input)
+    assignInputs = gen_input_assignment_code(inputSyms, instance, machine, :input)
 
     return Tape{input_type(instance)}(initCaches, assignInputs, schedule, inputSyms, outSym, Dict(), instance, machine)
 end
@@ -142,7 +140,7 @@ For implementation reasons, this disregards the set [`CacheStrategy`](@ref) of t
 """
 function execute_tape(tape::Tape, input)
     cache = Dict{Symbol, Any}()
-    cache[:data_input] = input
+    cache[:input] = input
     # simply execute all the code snippets here
     @assert typeof(input) == input_type(tape.instance)
     # TODO: `@assert` that input fits the tape.instance

src/code_gen/type.jl

@@ -11,7 +11,7 @@ TODO: update docs
 """
 struct Tape{INPUT}
     initCachesCode::Vector{Expr}
-    inputAssignCode::Vector{Expr}
+    inputAssignCode::Vector{FunctionCall}
     computeCode::Vector{FunctionCall}
     inputSymbols::Dict{String, Vector{Symbol}}
     outputSymbol::Symbol

src/models/interface.jl

@@ -39,8 +39,8 @@ Generate the [`DAG`](@ref) for the given [`AbstractProblemInstance`](@ref). Ever
 function graph end
 
 """
-    input_expr(instance::AbstractProblemInstance, name::String, input_symbol::Symbol)
+    input_expr(::AbstractProblemInstance, name::String, input)
 
-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.
+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.
 """
 function input_expr end

src/models/physics_models/qed/compute.jl

@@ -1,26 +1,12 @@
 using StaticArrays
 
-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)
+function input_expr(name::String, psp::PhaseSpacePoint)
     (type, index) = type_index_from_name(QEDModel(), name)
 
-    return Meta.parse(
-        "ParticleValueSP(
-    $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))),
-)",
+    return ParticleValueSP(
+        type(momentum(psp, particle_direction(type), particle_species(type), index)),
+        0.0im,
+        spin_or_pol(process(psp), type, index),
     )
 end
 

src/models/physics_models/qed/diagrams.jl

@@ -504,15 +504,18 @@ 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, SFourMomentum}, i) for i in 1:n],
-        [FeynmanParticle(ParticleStateful{Outgoing, Photon, SFourMomentum}, i) for i in 1:m],
+        [FeynmanParticle(ParticleStateful{Incoming, Photon}, i) for i in 1:n],
+        [FeynmanParticle(ParticleStateful{Outgoing, Photon}, i) for i in 1:m],
     )
 
     new_diagram = FeynmanDiagram(
@@ -520,10 +523,10 @@ function gen_compton_diagram_from_order_one_side(order::Vector{Int}, inFerm, out
         missing,
         [inFerm, outFerm, photons...],
         Dict{Type, Int64}(
-            ParticleStateful{Incoming, Electron, SFourMomentum} => 1,
-            ParticleStateful{Outgoing, Electron, SFourMomentum} => 1,
-            ParticleStateful{Incoming, Photon, SFourMomentum} => n,
-            ParticleStateful{Outgoing, Photon, SFourMomentum} => m,
+            ParticleStateful{Incoming, Electron} => 1,
+            ParticleStateful{Outgoing, Electron} => 1,
+            ParticleStateful{Incoming, Photon} => n,
+            ParticleStateful{Outgoing, Photon} => m,
         ),
     )
 
@@ -535,9 +538,9 @@ function gen_compton_diagram_from_order_one_side(order::Vector{Int}, inFerm, out
     while left_index <= right_index
         # left side
         v_left = FeynmanVertex(
-            FeynmanParticle(ParticleStateful{Incoming, Electron, SFourMomentum}, iterations),
+            FeynmanParticle(ParticleStateful{Incoming, Electron}, iterations),
             photons[order[left_index]],
-            FeynmanParticle(ParticleStateful{Incoming, Electron, SFourMomentum}, iterations + 1),
+            FeynmanParticle(ParticleStateful{Incoming, Electron}, iterations + 1),
         )
         left_index += 1
         add_vertex!(new_diagram, v_left)
@@ -550,9 +553,9 @@ function gen_compton_diagram_from_order_one_side(order::Vector{Int}, inFerm, out
         if (iterations == 1)
             # right side
             v_right = FeynmanVertex(
-                FeynmanParticle(ParticleStateful{Outgoing, Electron, SFourMomentum}, iterations),
+                FeynmanParticle(ParticleStateful{Outgoing, Electron}, iterations),
                 photons[order[right_index]],
-                FeynmanParticle(ParticleStateful{Outgoing, Electron, SFourMomentum}, iterations + 1),
+                FeynmanParticle(ParticleStateful{Outgoing, Electron}, iterations + 1),
             )
             right_index -= 1
             add_vertex!(new_diagram, v_right)
@@ -566,6 +569,7 @@ function gen_compton_diagram_from_order_one_side(order::Vector{Int}, inFerm, out
     add_tie!(new_diagram, FeynmanTie(ps[1], ps[2]))
     return new_diagram
 end
+=#
 
 """
     gen_compton_diagrams(n::Int, m::Int)