non-inplace version apply, dispatch and setiparams (#162)

* non-inplace version apply, dispatch and setiparams

* clean up methods

Author: GiggleLiu

src/YaoBlocks.jl

@@ -48,11 +48,13 @@ export AbstractBlock,
     ResetTo,
     TagBlock,
     apply!,
+    apply,
     apply_back!,
     chcontent,
     chsubblocks,
     content,
     dispatch!,
+    dispatch,
     expect,
     getiparams,
     iparams_eltype,
@@ -70,6 +72,7 @@ export AbstractBlock,
     print_block,
     render_params,
     setiparams!,
+    setiparams,
     subblocks,
     ishermitian,
     nparameters

src/abstract_block.jl

@@ -5,7 +5,7 @@ using YaoBase, YaoArrayRegister, SimpleTraits
 
 Apply a block (of quantum circuit) to a quantum register.
 """
-function apply!(r::AbstractRegister, @nospecialize(b::AbstractBlock))
+function apply!(r::AbstractRegister, b::AbstractBlock)
     _check_size(r, b)
     _apply!(r, b)
 end
@@ -131,34 +131,41 @@ Returns the intrinsic parameters of node `block`, default is an empty tuple.
 getiparams(x::AbstractBlock) = ()
 
 """
-    setiparams!(block, itr)
-    setiparams!(block, params...)
+    setiparams!([f], block, itr)
+    setiparams!([f], block, params...)
 
 Set the parameters of `block`.
+When `f` is provided, set parameters of `block` to the value in `collection` mapped by `f`.
+`iter` can be an iterator or a symbol, the symbol can be `:zero`, `:random`.
 """
-setiparams!(x::AbstractBlock, args...) =
-    niparams(x) == length(args) == 0 ? x : throw(NotImplementedError(:setiparams!, (x, args...)))
-
-setiparams!(x::AbstractBlock, it::Union{Tuple,AbstractArray,Base.Generator}) = setiparams!(x, it...)
-setiparams!(x::AbstractBlock, a::Number, xs::Number...) =
-    error("setparams!(x, θ...) is not implemented")
-setiparams!(x::AbstractBlock, it::Symbol) = setiparams!(x, render_params(x, it))
+function setiparams! end
 
 """
-    setiparams(f, block, collection)
+    setiparams([f], block, itr)
+    setiparams([f], block, params...)
 
-Set parameters of `block` to the value in `collection` mapped by `f`.
+Set the parameters of `block`, the non-inplace version.
+When `f` is provided, set parameters of `block` to the value in `collection` mapped by `f`.
+`iter` can be an iterator or a symbol, the symbol can be `:zero`, `:random`.
 """
-setiparams!(f::Function, x::AbstractBlock, it) =
-    setiparams!(x, map(x -> f(x...), zip(getiparams(x), it)))
-setiparams!(f::Nothing, x::AbstractBlock, it) = setiparams!(x, it)
+function setiparams end
 
-"""
-    setiparams(f, block, symbol)
+for F in [:setiparams!, :setiparams]
+    @eval begin
+        $F(x::AbstractBlock, args...) =
+            niparams(x) == length(args) == 0 ? x : throw(NotImplementedError($(QuoteNode(F)), (x, args...)))
 
-Set the parameters to a given symbol, which can be :zero, :random.
-"""
-setiparams!(f::Function, x::AbstractBlock, it::Symbol) = setiparams!(f, x, render_params(x, it))
+        $F(x::AbstractBlock, it::Union{Tuple,AbstractArray,Base.Generator}) = $F(x, it...)
+        $F(x::AbstractBlock, a::Number, xs::Number...) =
+            error("setparams!(x, θ...) is not implemented")
+        $F(x::AbstractBlock, it::Symbol) = $F(x, render_params(x, it))
+
+        $F(f::Function, x::AbstractBlock, it) =
+            $F(x, map(x -> f(x...), zip(getiparams(x), it)))
+        $F(f::Nothing, x::AbstractBlock, it) = $F(x, it)
+        $F(f::Function, x::AbstractBlock, it::Symbol) = $F(f, x, render_params(x, it))
+    end
+end
 
 """
     parameters(block)
@@ -363,3 +370,37 @@ function parameters_range!(out::Vector{Tuple{T,T}}, block::AbstractBlock) where
         parameters_range!(out, subblock)
     end
 end
+
+# non-inplace versions
+"""
+    apply(register, block)
+
+The non-inplace version of applying a block (of quantum circuit) to a quantum register.
+Check `apply!` for the faster inplace version.
+"""
+apply(r::AbstractRegister, b) = apply!(copy(r), b)
+
+function generic_dispatch!(f::Union{Function,Nothing}, x::AbstractBlock, it::Dispatcher)
+    x = setiparams(f, x, consume!(it, niparams(x)))
+    chsubblocks(x, map(subblocks(x)) do each
+        generic_dispatch!(f, each, it)
+    end)
+end
+
+"""
+    dispatch(x::AbstractBlock, collection)
+
+Dispatch parameters in collection to block tree `x`, the generic non-inplace version.
+
+!!! note
+
+    it will try to dispatch the parameters in collection first.
+"""
+function dispatch(f::Union{Function,Nothing}, x::AbstractBlock, it)
+    dp = Dispatcher(it)
+    res = generic_dispatch!(f, x, dp)
+    @assert (it isa Symbol || length(it) == dp.loc) "expect $(dp.loc) parameters, got $(length(it))"
+    return res
+end
+
+dispatch(x::AbstractBlock, it) = dispatch(nothing, x, it)

src/primitive/phase_gate.jl

@@ -37,6 +37,7 @@ mat(::Type{T}, gate::PhaseGate) where {T} = exp(T(im * gate.theta)) * IMatrix{2,
 niparams(::Type{<:PhaseGate}) = 1
 getiparams(x::PhaseGate) = x.theta
 setiparams!(r::PhaseGate, param::Number) = (r.theta = param; r)
+setiparams(r::PhaseGate, param::Number) = PhaseGate(param)
 
 # fallback to matrix method if it is not real
 YaoBase.isunitary(r::PhaseGate{<:Real}) = true

src/primitive/rotation_gate.jl

@@ -108,6 +108,7 @@ niparams(::Type{<:RotationGate}) = 1
 getiparams(x::RotationGate) = x.theta
 # no need to specify the type of param, Julia will try to do the conversion
 setiparams!(r::RotationGate, param::Number) where {N,T} = (r.theta = param; r)
+setiparams(r::RotationGate, param::Number) where {N,T} = RotationGate(r.block, param)
 
 # fallback to matrix methods if it is not real
 YaoBase.isunitary(r::RotationGate{N,<:Real}) where {N} = true

src/primitive/shift_gate.jl

@@ -43,6 +43,7 @@ cache_key(gate::ShiftGate) = gate.theta
 niparams(::Type{<:ShiftGate}) = 1
 getiparams(x::ShiftGate) = x.theta
 setiparams!(r::ShiftGate, param::Number) = (r.theta = param; r)
+setiparams(r::ShiftGate, param::Number) = ShiftGate(param)
 
 
 Base.adjoint(blk::ShiftGate) = ShiftGate(-blk.theta)

src/primitive/time_evolution.jl

@@ -81,6 +81,7 @@ cache_key(te::TimeEvolution) = (te.dt, cache_key(te.H))
 niparams(::Type{<:TimeEvolution}) = 1
 getiparams(x::TimeEvolution) = x.dt
 setiparams!(r::TimeEvolution, param::Number) = (r.dt = param; r)
+setiparams(r::TimeEvolution, param::Number) = TimeEvolution(r.H, param; tol=r.tol, check_hermicity=false)
 
 function Base.:(==)(lhs::TimeEvolution, rhs::TimeEvolution)
     return lhs.H == rhs.H && lhs.dt == rhs.dt

src/routines.jl

@@ -34,7 +34,7 @@ set specific col of a CSC matrix
 end
 
 """
-    num_nonzero(nbits, nctrls, U)
+    num_nonzero(nbits, nctrls, U, [N])
 
 Return number of nonzero entries of the matrix form of control-U gate. `nbits`
 is the number of qubits, and `nctrls` is the number of control qubits.
@@ -43,8 +43,8 @@ is the number of qubits, and `nctrls` is the number of control qubits.
     return N + (1 << (nbits - nctrls - log2dim1(U))) * (length(U) - size(U, 2))
 end
 
-@inline function num_nonzero(nbits::Int, U, N::Int = 1 << nbits)
-    return N + (1 << (nbits - log2dim1(U))) * (length(U) - size(U, 2))
+@inline function num_nonzero(nbits::Int, nctrls::Int, U::SDSparseMatrixCSC, N::Int = 1 << nbits)
+    return N + (1 << (nbits - nctrls - log2dim1(U))) * (nnz(U) - size(U, 2))
 end
 
 """
@@ -121,7 +121,7 @@ function u1mat(nbits::Int, U1::SDMatrix, ibit::Int)
     a, c, b, d = U1
     step = 1 << (ibit - 1)
     step_2 = 1 << ibit
-    NNZ = num_nonzero(nbits, U1, N)
+    NNZ = num_nonzero(nbits, 0, U1, N)
 
     colptr = Vector{Int}(1:2:2*N+1)
     rowval = Vector{Int}(undef, NNZ)

test/abstract_blocks.jl

@@ -5,6 +5,7 @@ using YaoBase
     @test getiparams(phase(0.1)) == 0.1
     @test getiparams(2 * phase(0.1)) == ()
     @test_throws NotImplementedError setiparams!(rot(X, 0.5), :nothing)
+    @test_throws NotImplementedError setiparams(rot(X, 0.5), :nothing)
 end
 
 @testset "block to matrix conversion" begin
@@ -18,6 +19,9 @@ end
 @testset "apply lambda" begin
     r = rand_state(3)
     @test apply!(copy(r), put(1 => X)) ≈ apply!(copy(r), put(3, 1 => X))
+    r2 = copy(r)
+    @test apply(r, put(1 => X)) ≈ apply!(copy(r), put(3, 1 => X))
+    @test r2.state == r.state
     f(x::Float32) = x
     @test_throws ErrorException apply!(r, f)
 end

test/composite/swap_gate.jl

@@ -28,7 +28,13 @@ end
     dispatch!(pb, π)
     @test copy(reg) |> pb ≈ -im * (copy(reg) |> swap(6, 2, 4))
     @test copy(reg) |> pb |> isnormalized
+    pb = dispatch(pb, π)
+    @test copy(reg) |> pb ≈ -im * (copy(reg) |> swap(6, 2, 4))
+    @test copy(reg) |> pb |> isnormalized
 
+    pb = pswap(6, 2, 4, 0.0)
     dispatch!(pb, :random)
     @test copy(reg) |> pb ≈ invoke(apply!, Tuple{ArrayReg,PutBlock}, copy(reg), pb)
+    pb = dispatch(pb, :random)
+    @test copy(reg) |> pb ≈ invoke(apply!, Tuple{ArrayReg,PutBlock}, copy(reg), pb)
 end

test/primitive/phase_gate.jl

@@ -30,6 +30,7 @@ end
     @test nparameters(phase(0.1)) == 1
     @test adjoint(phase(0.1)) == phase(-0.1)
     @test dispatch!(phase(0.1), 0.3) == phase(0.3)
+    @test dispatch(phase(0.1), 3) == phase(3) && eltype(getiparams(dispatch(phase(0.1), 3))) == Int
 
     @testset "test $op" for op in [+, -, *, /]
         @test dispatch!(op, phase(0.1), π) == phase(op(0.1, π))