new-parametrized SWAP gate, remove Swap type (#52)

* new-parametrized SWAP gate, remove Swap type

* fix printing for swap, export swap correctly

* update benchmarks

benchmark/benchmarks.jl

@@ -25,6 +25,8 @@ end
 for n in 5:5:25
     SUITE["composite"]["kron(rand_const)"] = bench(kron(rand([X, Y, Z, H]) for _ in 1:n))
     SUITE["composite"]["kron(sparse_const)"] = bench(kron(n, k=>rand([X, Y, Z, H]) for k in randperm(n)[1:n÷5]))
+    SUITE["composite"]["swap"] = bench(swap(n,2,4))
+    SUITE["composite"]["swap"] = bench(pswap(n,2,4,0.5))
 end
 
 function heisenberg(n::Int; periodic::Bool=true)

src/composite/put_block.jl

@@ -1,4 +1,4 @@
-export PutBlock, put
+export PutBlock, put, Swap, swap, PSwap, pswap
 
 """
     PutBlock <: AbstractContainer
@@ -104,3 +104,60 @@ function YaoBase.iscommute(x::PutBlock{N}, y::PutBlock{N}) where N
         return iscommute_fallback(x, y)
     end
 end
+
+const Swap{N} = PutBlock{N,2,G} where G<:ConstGate.SWAPGate
+const PSwap{N, T} = PutBlock{N,2,RotationGate{2,T,G}} where G<:ConstGate.SWAPGate
+Swap{N}(locs::Tuple{Int, Int}) where N = PutBlock{N}(ConstGate.SWAPGate(), locs)
+PSwap{N}(locs::Tuple{Int, Int}, θ::Real) where N = PutBlock{N}(rot(ConstGate.SWAPGate(), θ), locs)
+
+"""
+    swap(n, loc1, loc2)
+
+Create a `n`-qubit [`Swap`](@ref) gate which swap `loc1` and `loc2`.
+
+# Example
+
+```jldoctest
+julia> swap(4, 1, 2)
+swap(1, 2)
+```
+"""
+swap(n::Int, loc1::Int, loc2::Int) = Swap{n}((loc1, loc2))
+
+"""
+    swap(loc1, loc2) -> f(n)
+
+Create a lambda that takes the total number of active qubits as input. Lazy curried
+version of `swap(n, loc1, loc2)`. See also [`Swap`](@ref).
+
+# Example
+
+```jldoctest
+julia> swap(1, 2)
+(n -> swap(n, 1, 2))
+```
+"""
+swap(loc1::Int, loc2::Int) = @λ(n -> swap(n, loc1, loc2))
+
+function mat(::Type{T}, g::Swap{N}) where {T, N}
+    mask = bmask(g.locs[1], g.locs[2])
+    orders = map(b->swapbits(b, mask) + 1, basis(N))
+    return PermMatrix(orders, ones(T, 1<<N))
+end
+
+apply!(r::ArrayReg, g::Swap) = (instruct!(matvec(state(r)), Val(:SWAP), g.locs); r)
+occupied_locs(g::Swap) = g.locs
+
+"""
+    pswap(n::Int, i::Int, j::Int, α::Real)
+    pswap(i::Int, j::Int, α::Real) -> f(n)
+
+parametrized swap gate.
+"""
+pswap(n::Int, i::Int, j::Int, α::Real) = PSwap{n}((i,j), α)
+pswap(i::Int, j::Int, α::Real) = n->pswap(n,i,j,α)
+
+function apply!(reg::ArrayReg, g::PSwap{N, T}) where {N,T}
+    instruct!(matvec(state(reg)), Val(:PSWAP), g.locs, g.content.theta)
+    return reg
+end

src/layout.jl

@@ -153,7 +153,6 @@ end
 
 # Custom layouts
 color(m::AbstractBlock) = color(typeof(m))
-color(::Type{<:Swap}) = :magenta
 color(::Type{<:ControlBlock}) = :red
 color(::Type{<:ChainBlock}) = :blue
 color(::Type{<:MathGate}) = :red
@@ -168,7 +167,6 @@ color(::Type{<:GeneralMatrixBlock}) = :red
 print_block(io::IO, g::PhaseGate) = print(io, "phase(", g.theta, ")")
 print_block(io::IO, S::ShiftGate) = print(io, "shift(", S.theta, ")")
 print_block(io::IO, R::RotationGate) = print(io, "rot(", content(R), ", ", R.theta, ")")
-print_block(io::IO, swap::Swap) = printstyled(io, "swap", swap.locs; bold=true, color=color(Swap))
 print_block(io::IO, x::KronBlock) = printstyled(io, "kron"; bold=true, color=color(KronBlock))
 print_block(io::IO, x::ChainBlock) = printstyled(io, "chain"; bold=true, color=color(ChainBlock))
 print_block(io::IO, x::ReflectGate{N}) where N = print(io, "reflect($(summary(x.psi)))")

src/primitive/const_gate.jl

@@ -1,6 +1,6 @@
 export @const_gate, ConstantGate, ConstGate,
-    X, Y, Z, H, I2,
-    XGate, YGate, ZGate, HGate, I2Gate
+    X, Y, Z, H, I2, SWAP,
+    XGate, YGate, ZGate, HGate, I2Gate, SWAPGate
 
 """
     ConstGate
@@ -36,5 +36,5 @@ end
 
 # import some frequently-used objects
 import .ConstGate:
-    XGate, YGate, ZGate, HGate, I2Gate, X, Y, Z, H, I2,
+    XGate, YGate, ZGate, HGate, I2Gate, SWAPGate, X, Y, Z, H, I2, SWAP,
     @const_gate, ConstantGate, PauliGate

src/primitive/primitive.jl

@@ -25,7 +25,6 @@ include("const_gate.jl")
 include("phase_gate.jl")
 include("shift_gate.jl")
 include("rotation_gate.jl")
-include("swap_gate.jl")
 include("time_evolution.jl")
 include("reflect_gate.jl")
 include("general_matrix_gate.jl")

test/composite/put.jl

@@ -38,3 +38,9 @@ end
     @test pb |> isunitary
     @test pb |> isreflexive
 end
+
+@testset "test swap gate" begin
+    include("swap_gate.jl")
+end
+
+

test/primitive/swap_gate.jl → test/composite/swap_gate.jl

@@ -16,3 +16,19 @@ end
 @test isunitary(swap(3, 1, 2))
 @test ishermitian(swap(3, 1, 2))
 @test isreflexive(swap(3, 1, 2))
+
+@testset "pswap gate" begin
+    pb = pswap(6, 2, 4, 0.0)
+    @test pb isa PSwap{6, Float64}
+    @test pb == pswap(2, 4, 0.0)(6)
+    reg = rand_state(6)
+    @test copy(reg) |> pb ≈ invoke(apply!, Tuple{ArrayReg, PutBlock}, copy(reg), pb)
+    @test copy(reg) |> pb ≈ reg
+
+    dispatch!(pb, π)
+    @test copy(reg) |> pb ≈ -im*(copy(reg) |> swap(6, 2, 4))
+    @test copy(reg) |> pb |> isnormalized
+
+    dispatch!(pb, :random)
+    @test copy(reg) |> pb ≈ invoke(apply!, Tuple{ArrayReg, PutBlock}, copy(reg), pb)
+end

test/primitive/primitive.jl

@@ -16,10 +16,6 @@ end
     include("rotation_gate.jl")
 end
 
-@testset "test swap gate" begin
-    include("swap_gate.jl")
-end
-
 @testset "test time evolution" begin
     include("time_evolution.jl")
 end
@@ -41,4 +37,4 @@ end
 end
 
 # it does nothing
-@test chsubblocks(X, Y) === X
+@test chsubblocks(X, Y) === X