Work on QED Model execution

test/Project.toml

@@ -3,4 +3,5 @@ AccurateArithmetic = "22286c92-06ac-501d-9306-4abd417d9753"
 QEDbase = "10e22c08-3ccb-4172-bfcf-7d7aa3d04d93"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/unit_tests_qedmodel.jl

@@ -1,5 +1,6 @@
 using MetagraphOptimization
 using QEDbase
+using StatsBase # for countmap
 
 import MetagraphOptimization.caninteract
 import MetagraphOptimization.issame
@@ -44,3 +45,142 @@ end
         @test direction(propagation_result(p)(def_momentum)) != direction(p(def_momentum))
     end
 end
+
+@testset "Parse Process" begin
+    @testset "Order invariance" begin
+        @test parse_process("ke->ke", QEDModel()) == parse_process("ek->ke", QEDModel())
+        @test parse_process("ke->ke", QEDModel()) == parse_process("ek->ek", QEDModel())
+        @test parse_process("ke->ke", QEDModel()) == parse_process("ke->ek", QEDModel())
+
+        @test parse_process("kkke->eep", QEDModel()) == parse_process("kkek->epe", QEDModel())
+    end
+
+    @testset "Known processes" begin
+        compton_process = QEDProcessDescription(
+            Dict{Type, Int}(PhotonStateful{Incoming} => 1, FermionStateful{Incoming} => 1),
+            Dict{Type, Int}(PhotonStateful{Outgoing} => 1, FermionStateful{Outgoing} => 1),
+        )
+
+        @test parse_process("ke->ke", QEDModel()) == compton_process
+
+        positron_compton_process = QEDProcessDescription(
+            Dict{Type, Int}(PhotonStateful{Incoming} => 1, AntiFermionStateful{Incoming} => 1),
+            Dict{Type, Int}(PhotonStateful{Outgoing} => 1, AntiFermionStateful{Outgoing} => 1),
+        )
+
+        @test parse_process("kp->kp", QEDModel()) == positron_compton_process
+
+        trident_process = QEDProcessDescription(
+            Dict{Type, Int}(PhotonStateful{Incoming} => 1, FermionStateful{Incoming} => 1),
+            Dict{Type, Int}(FermionStateful{Outgoing} => 2, AntiFermionStateful{Outgoing} => 1),
+        )
+
+        @test parse_process("ke->eep", QEDModel()) == trident_process
+
+        pair_production_process = QEDProcessDescription(
+            Dict{Type, Int}(PhotonStateful{Incoming} => 2),
+            Dict{Type, Int}(FermionStateful{Outgoing} => 1, AntiFermionStateful{Outgoing} => 1),
+        )
+
+        @test parse_process("kk->pe", QEDModel()) == pair_production_process
+
+        pair_annihilation_process = QEDProcessDescription(
+            Dict{Type, Int}(FermionStateful{Incoming} => 1, AntiFermionStateful{Incoming} => 1),
+            Dict{Type, Int}(PhotonStateful{Outgoing} => 2),
+        )
+
+        @test parse_process("pe->kk", QEDModel()) == pair_annihilation_process
+    end
+end
+
+@testset "Generate Process Inputs" begin
+    @testset "Process $proc_str" for proc_str in ["ke->ke", "kp->kp", "kk->ep", "ep->kk"]
+        # currently can only generate for 2->2 processes
+        process = parse_process(proc_str, QEDModel())
+
+        for i in 1:100
+            input = gen_process_input(process)
+            @test countmap(typeof.(input.inParticles)) == process.inParticles
+            @test countmap(typeof.(input.outParticles)) == process.outParticles
+
+            @test isapprox(
+                sum(getfield.(input.inParticles, :momentum)) + sum(getfield.(input.outParticles, :momentum)),
+                SFourMomentum(0.0, 0.0, 0.0, 0.0);
+                atol = sqrt(eps()),
+            )
+        end
+    end
+end
+
+@testset "Compton" begin
+    import MetagraphOptimization.insert_node!
+    import MetagraphOptimization.insert_edge!
+    import MetagraphOptimization.make_node
+
+    model = QEDModel()
+    process = parse_process("ke->ke", model)
+    machine = get_machine_info()
+
+    graph = MetagraphOptimization.DAG()
+
+    # manually build a graph for compton
+    graph = DAG()
+
+    # s to output (exit node)
+    d_exit = insert_node!(graph, make_node(DataTask(16)), track = false)
+
+    # final s compute
+    s0 = insert_node!(graph, make_node(ComputeTaskQED_S2()), track = false)
+
+    # data from v0 and v1 to s0
+    d_v0_s0 = insert_node!(graph, make_node(DataTask(96)), track = false)
+    d_v1_s0 = insert_node!(graph, make_node(DataTask(96)), track = false)
+
+    # v0 and v1 compute
+    v0 = insert_node!(graph, make_node(ComputeTaskQED_V()), track = false)
+    v1 = insert_node!(graph, make_node(ComputeTaskQED_V()), track = false)
+
+    # data from uPhIn, uPhOut, uElIn, uElOut
+    d_uPhIn_v0 = insert_node!(graph, make_node(DataTask(96)), track = false)
+    d_uElIn_v0 = insert_node!(graph, make_node(DataTask(96)), track = false)
+    d_uPhOut_v1 = insert_node!(graph, make_node(DataTask(96)), track = false)
+    d_uElOut_v1 = insert_node!(graph, make_node(DataTask(96)), track = false)
+
+    # uPhIn, uPhOut, uElIn and uElOut computes
+    uPhIn = insert_node!(graph, make_node(ComputeTaskQED_U()), track = false)
+    uPhOut = insert_node!(graph, make_node(ComputeTaskQED_U()), track = false)
+    uElIn = insert_node!(graph, make_node(ComputeTaskQED_U()), track = false)
+    uElOut = insert_node!(graph, make_node(ComputeTaskQED_U()), track = false)
+
+    # data into U
+    d_uPhIn = insert_node!(graph, make_node(DataTask(16), "ki1"), track = false)
+    d_uPhOut = insert_node!(graph, make_node(DataTask(16), "ko1"), track = false)
+    d_uElIn = insert_node!(graph, make_node(DataTask(16), "ei1"), track = false)
+    d_uElOut = insert_node!(graph, make_node(DataTask(16), "eo1"), track = false)
+
+    # now for all the edges
+    insert_edge!(graph, d_uPhIn, uPhIn, track = false)
+    insert_edge!(graph, d_uPhOut, uPhOut, track = false)
+    insert_edge!(graph, d_uElIn, uElIn, track = false)
+    insert_edge!(graph, d_uElOut, uElOut, track = false)
+
+    insert_edge!(graph, uPhIn, d_uPhIn_v0, track = false)
+    insert_edge!(graph, uPhOut, d_uPhOut_v1, track = false)
+    insert_edge!(graph, uElIn, d_uElIn_v0, track = false)
+    insert_edge!(graph, uElOut, d_uElOut_v1, track = false)
+
+    insert_edge!(graph, d_uPhIn_v0, v0, track = false)
+    insert_edge!(graph, d_uPhOut_v1, v1, track = false)
+    insert_edge!(graph, d_uElIn_v0, v0, track = false)
+    insert_edge!(graph, d_uElOut_v1, v1, track = false)
+
+    insert_edge!(graph, v0, d_v0_s0, track = false)
+    insert_edge!(graph, v1, d_v1_s0, track = false)
+
+    insert_edge!(graph, d_v0_s0, s0, track = false)
+    insert_edge!(graph, d_v1_s0, s0, track = false)
+
+    insert_edge!(graph, s0, d_exit, track = false)
+
+    compton_function = get_compute_function(graph, process, machine)
+end