Small fixes

.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

@@ -6,13 +6,8 @@ using QEDcore
 using QEDprocesses
 using Random
 using UUIDs
-
-using CUDA
-
-using NamedDims
 using CSV
 using JLD2
-using FlexiMaps
 
 RNG = Random.MersenneTwister(123)
 
@@ -84,7 +79,7 @@ function congruent_input_momenta_scenario_2(
     # ----------
 
     # now calculate the final_momenta from omega, cos_theta and phi
-    n = number_particles(processDescription, Incoming(), Photon())
+    n = number_particles(processDescription, ParticleStateful{Incoming, Photon, SFourMomentum})
 
     cos_theta = cos(theta)
     omega_prime = (n * omega) / (1 + n * omega * (1 - cos_theta))
@@ -113,82 +108,109 @@ end
 with_stacksize(f, n) = fetch(schedule(Task(f, n)))
 
 # scenario 2
-N = 1024    # thetas
-M = 1024    # phis
-K = 64      # omegas
+N = 1000
+M = 1000
 
 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
+for photons in 1:6
     # 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())
+    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 =
-        Array{typeof(congruent_input_momenta_scenario_2(temp_process, omegas[1], thetas[1], phis[1]))}(undef, (K, N, M))
+        input_momenta = [
+            congruent_input_momenta_scenario_2(temp_process, omega, theta, phi) for
+            (theta, phi) in Iterators.product(thetas, phis)
+        ]
+        results = Array{Float64}(undef, size(input_momenta))
+        fill!(results, 0.0)
 
-    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
+        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)
 
-    cu_results = CuArray{Float64}(undef, size(input_momenta))
-    fill!(cu_results, 0.0)
+            print("Preparing graph... ")
+            graph = gen_graph(process)
+            optimize_to_fixpoint!(ReductionOptimizer(), graph)
+            print("Preparing function... ")
+            func = get_compute_function(graph, process, mock_machine())
+            func(inputs[1])
 
-    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
+            print("Calculating... ")
             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])
+                    return results[i, j] += abs2(func(inputs[i, j]))
                 end
             end
+            println("Done.")
+            i += 1
         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())
+        println("Writing results")
 
-        print("Calculating... ")
-        ts = 32
-        bs = Int64(length(cu_inputs) / 32)
+        out_ph_moms = getindex.(getindex.(input_momenta, 2), 1)
+        out_el_moms = getindex.(getindex.(input_momenta, 2), 2)
 
-        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
+        @save "$(photons)_congruent_photons_omega_$(omega)_grid.jld2" out_ph_moms out_el_moms results
+    end
+end
+
+exit(0)
+
+# scenario 1 (disabled)
+n = 1000000
+
+# n is the number of incoming photons
+# omega is the number
+
+for photons in 1: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())
+
+        input_momenta = [congruent_input_momenta(temp_process, omega) for _ in 1:n]
+        results = Array{Float64}(undef, size(input_momenta))
+        fill!(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 = 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:n
+                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).jld2" out_ph_moms out_el_moms results
     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,12 +1,71 @@
 # patch QEDprocesses
 # see issue https://github.com/QEDjl-project/QEDprocesses.jl/issues/77
 @inline function QEDprocesses.number_particles(
-    proc_def::QEDbase.AbstractProcessDefinition, ::Type{PS}
+    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},
+    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},
+) 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 QEDcore.ParticleStateful{DIR, SPECIES}(
+    mom::AbstractFourMomentum,
+) where {DIR <: ParticleDirection, SPECIES <: AbstractParticleType}
+    return ParticleStateful(DIR(), SPECIES(), mom)
+end
+
+@inline function QEDcore.ParticleStateful{DIR, SPECIES, EL}(
+    mom::EL,
+) 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

@@ -16,7 +16,9 @@ function get_compute_function(graph::DAG, instance, machine::Machine)
         "function compute_$(functionId)(data_input::$(input_type(instance))) $(initCaches); $(assignInputs); $code; return $resSym; end",
     )
 
-    return expr
+    func = eval(expr)
+
+    return func
 end
 
 """
@@ -33,22 +35,22 @@ function get_cuda_kernel(graph::DAG, instance, machine::Machine)
 
     functionId = to_var_name(UUIDs.uuid1(rng[1]))
     resSym = eval(gen_access_expr(entry_device(tape.machine), tape.outputSymbol))
-    expr = Meta.parse(
-        "function compute_$(functionId)(input_vector, output_vector, n::Int64)
-            id = (blockIdx().x - 1) * blockDim().x + threadIdx().x
-            if (id > n)
-                return
-            end
-            @inline data_input = input_vector[id]
-            $(initCaches)
-            $(assignInputs)
-            $code
-            @inline output_vector[id] = $resSym
-            return nothing
-        end"
-    )
+    expr = Meta.parse("function compute_$(functionId)(input_vector, output_vector, n::Int64)
+                          id = (blockIdx().x - 1) * blockDim().x + threadIdx().x
+                          if (id > n)
+                              return
+                          end
+                          @inline data_input = input_vector[id]
+                          $(initCaches)
+                          $(assignInputs)
+                          $code
+                          @inline output_vector[id] = $resSym
+                          return nothing
+                      end")
 
-    return expr
+    func = eval(expr)
+
+    return func
 end
 
 """

src/models/physics_models/qed/compute.jl

@@ -17,7 +17,7 @@ function input_expr(instance::GenericQEDProcess, name::String, psp_symbol::Symbo
 
     return Meta.parse(
         "ParticleValueSP(
-    $type(momentum($psp_symbol, $(construction_string(particle_direction(type))), $(construction_string(particle_species(type))), Val($index))),
+    $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))),
 )",