optimize projection gate

Project.toml

@@ -16,7 +16,7 @@ YaoToEinsumCUDAExt = "CUDA"
 
 [compat]
 CUDA = "4, 5"
-OMEinsum = "0.7"
+OMEinsum = "0.8"
 Yao = "0.8"
 julia = "1.9"
 

src/circuitmap.jl

@@ -26,6 +26,19 @@ function add_matrix!(eb::EinBuilder{T}, k::Int, m::AbstractMatrix, locs::Vector)
         nlabels = [newlabel!(eb) for _=1:k]
         add_tensor!(eb, reshape(Matrix{T}(m), fill(2, 2k)...), [nlabels..., eb.slots[locs]...])
         eb.slots[locs] .= nlabels
+    elseif m isa Yao.OuterProduct  # low rank
+        nlabels = [newlabel!(eb) for _=1:k]
+        K = rank(m)
+        if K == 1  # projector
+            add_tensor!(eb, reshape(Vector{T}(m.right), fill(2, k)...), [eb.slots[locs]...])
+            add_tensor!(eb, reshape(Vector{T}(m.left), fill(2, k)...), [nlabels...])
+            eb.slots[locs] .= nlabels
+        else
+            midlabel = newlabel!(eb)
+            add_tensor!(eb, reshape(Matrix{T}(m.right), fill(2, k)..., K), [eb.slots[locs]..., midlabel])
+            add_tensor!(eb, reshape(Matrix{T}(m.left), fill(2, k)..., K), [nlabels..., midlabel])
+            eb.slots[locs] .= nlabels
+        end
     else
         add_tensor!(eb, reshape(Vector{T}(diag(m)), fill(2, k)...), eb.slots[locs])
     end
@@ -39,6 +52,19 @@ function add_gate!(eb::EinBuilder{T}, b::PutBlock{2,2,ConstGate.SWAPGate}) where
     return eb
 end
 
+# projection gate, todo: generalize to arbitrary low rank gate
+function add_gate!(eb::EinBuilder{T}, b::PutBlock{2,1,ConstGate.P0Gate}) where {T}
+    add_matrix!(eb, 1, Yao.OuterProduct(T[1, 0], T[1, 0]), collect(b.locs))
+    return eb
+end
+
+# projection gate, todo: generalize to arbitrary low rank gate
+function add_gate!(eb::EinBuilder{T}, b::PutBlock{2,1,ConstGate.P1Gate}) where {T}
+    add_matrix!(eb, 1, Yao.OuterProduct(T[0, 1], T[0, 1]), collect(b.locs))
+    return eb
+end
+
+
 # control gates
 function add_gate!(eb::EinBuilder{T}, b::ControlBlock{BT,C,M}) where {T, BT,C,M}
     return add_controlled_matrix!(eb, M, mat(T, b.content), collect(b.locs), collect(b.ctrl_locs), collect(b.ctrl_config))

test/circuitmap.jl

@@ -6,7 +6,7 @@ using SymEngine
 
 @testset "YaoToEinsum.jl" begin
     n = 5
-    for c in [put(n, 2=>Y), put(n, (5,3)=>SWAP), put(n, (4,2)=>ConstGate.CNOT), put(n, (2,3,1)=>kron(ConstGate.CNOT, X)),
+    for c in [put(n, 2=>Y), put(n, 2=>ConstGate.P0), put(n, 2=>ConstGate.P1), put(n, (5,3)=>SWAP), put(n, (4,2)=>ConstGate.CNOT), put(n, (2,3,1)=>kron(ConstGate.CNOT, X)),
             put(n, 2=>Z), control(n, -3, 2=>X), control(n, 3, 2=>X), control(n, (2, -1), 3=>Y), control(n, (4,1,-2), 5=>Z)]
         @show c
         C = chain([put(n, i=>Rx(rand()*2π)) for i=1:n]..., c)