From 6c383afce2a1eae7a01249897f05d606c9696b4f Mon Sep 17 00:00:00 2001
From: Leo <cacate0129@gmail.com>
Date: Tue, 28 May 2019 01:49:58 +0800
Subject: [PATCH] fix scale and several other issues (#46)

* seperate kron, export ConstGate

* fix scale

* fix expect for scale
---
 src/blocktools.jl              |  2 +-
 src/composite/kron.jl          | 10 +++++++---
 src/composite/tag/scale.jl     | 14 +++++++++-----
 src/primitive/const_gate.jl    |  2 +-
 src/primitive/rotation_gate.jl |  8 ++++----
 test/blocktools.jl             |  2 ++
 test/composite/tag.jl          | 30 +++++++++++++++++++++++++++++-
 7 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/src/blocktools.jl b/src/blocktools.jl
index 16284e9..aa5ca60 100644
--- a/src/blocktools.jl
+++ b/src/blocktools.jl
@@ -76,7 +76,7 @@ function expect(op::Sum, reg::AbstractRegister)
 end
 
 function expect(op::Scale, reg::AbstractRegister)
-    op.alpha*expect(content(op), reg)
+    factor(op)*expect(content(op), reg)
 end
 
 expect(op::Sum, reg::AbstractRegister{1}) = invoke(expect, Tuple{Sum, AbstractRegister}, op, reg)
diff --git a/src/composite/kron.jl b/src/composite/kron.jl
index d03f86e..8a88eec 100644
--- a/src/composite/kron.jl
+++ b/src/composite/kron.jl
@@ -141,16 +141,20 @@ chsubblocks(pb::KronBlock{N}, it) where N = KronBlock{N}(pb.locs, collect(it))
 cache_key(x::KronBlock) = [cache_key(each) for each in x.blocks]
 color(::Type{T}) where {T <: KronBlock} = :cyan
 
-
-function mat(::Type{T}, k::KronBlock{N}) where {T, N}
+function _prepair_kronmat(k::KronBlock{N}) where N
     sizes = map(nqubits, subblocks(k))
     start_locs = @. N - $(k.locs) - sizes + 1
 
     order = sortperm(start_locs)
     sorted_start_locs = start_locs[order]
     num_bit_list = vcat(diff(push!(sorted_start_locs, N)) .- sizes[order])
+    return order, num_bit_list, sorted_start_locs
+end
 
-    return reduce(zip(subblocks(k)[order], num_bit_list), init=IMatrix{1 << sorted_start_locs[1], T}()) do x, y
+function mat(::Type{T}, k::KronBlock{N}) where {T, N}
+    order, num_bit_list, sorted_start_locs = _prepair_kronmat(k)
+    blocks = subblocks(k)[order]
+    return reduce(zip(blocks, num_bit_list), init=IMatrix{1 << sorted_start_locs[1], T}()) do x, y
         kron(x, mat(T, y[1]), IMatrix(1<<y[2]))
     end
 end
diff --git a/src/composite/tag/scale.jl b/src/composite/tag/scale.jl
index f1fe4aa..4ac6e06 100644
--- a/src/composite/tag/scale.jl
+++ b/src/composite/tag/scale.jl
@@ -1,6 +1,6 @@
 using LinearAlgebra
 
-export Scale
+export Scale, factor
 
 """
     Scale{S <: Union{Number, Val}, N, BT <: AbstractBlock{N}} <: TagBlock{BT, N}
@@ -29,19 +29,23 @@ struct Scale{S <: Union{Number, Val}, N, BT <: AbstractBlock{N}} <: TagBlock{BT,
 end
 
 content(x::Scale) = x.content
+factor(x::Scale{<:Number}) = x.alpha
+factor(x::Scale{Val{X}}) where X = X
+
 Base.copy(x::Scale) = Scale(x.alpha, copy(x.content))
-Base.adjoint(x::Scale) = Scale(adjoint(x.alpha), adjoint(content(x)))
+Base.adjoint(x::Scale{<:Number}) = Scale(adjoint(x.alpha), adjoint(content(x)))
+Base.adjoint(x::Scale{Val{X}}) where X = Scale(Val(adjoint(X)), adjoint(content(x)))
 
-Base.:(==)(x::Scale, y::Scale) = (x.alpha == y.alpha) && (content(x) == content(y))
+Base.:(==)(x::Scale, y::Scale) = (factor(x) == factor(y)) && (content(x) == content(y))
 
 chsubblocks(x::Scale, blk::AbstractBlock) = Scale(x.alpha, blk)
-cache_key(x::Scale) = (x.alpha, cache_key(content(x)))
+cache_key(x::Scale) = (factor(x), cache_key(content(x)))
 
 mat(::Type{T}, x::Scale) where T = T(x.alpha) * mat(T, content(x))
 mat(::Type{T}, x::Scale{Val{S}}) where {T, S} = T(S) * mat(T, content(x))
 
 function apply!(r::ArrayReg{B}, x::Scale{S, N}) where {S, B, N}
     apply!(r, content(x))
-    r.state .*= x.alpha
+    r.state .*= factor(x)
     return r
 end
diff --git a/src/primitive/const_gate.jl b/src/primitive/const_gate.jl
index 6df9c89..50c7f31 100644
--- a/src/primitive/const_gate.jl
+++ b/src/primitive/const_gate.jl
@@ -1,4 +1,4 @@
-export @const_gate, ConstantGate,
+export @const_gate, ConstantGate, ConstGate,
     X, Y, Z, H, I2,
     XGate, YGate, ZGate, HGate, I2Gate
 
diff --git a/src/primitive/rotation_gate.jl b/src/primitive/rotation_gate.jl
index 25e23bf..f343c89 100644
--- a/src/primitive/rotation_gate.jl
+++ b/src/primitive/rotation_gate.jl
@@ -85,10 +85,10 @@ function mat(::Type{T}, R::RotationGate{N}) where {N, T}
 end
 
 # Specialized
-# mat(R::RotationGate{1, T, XGate{Complex{T}}}) where T =
-#     SMatrix{2, 2, Complex{T}}(cos(R.theta/2), -im * sin(R.theta/2), -im * sin(R.theta/2), cos(R.theta/2))
-# mat(R::RotationGate{1, T, YGate{Complex{T}}}) where T =
-#     SMatrix{2, 2, Complex{T}}(cos(R.theta/2), sin(R.theta/2), -sin(R.theta/2), cos(R.theta/2))
+mat(::Type{T}, R::RotationGate{1, <:Any, <:XGate}) where T =
+    T[cos(R.theta/2) -im * sin(R.theta/2); -im * sin(R.theta/2) cos(R.theta/2)]
+mat(::Type{T}, R::RotationGate{1, <:Any, <:YGate}) where T =
+     T[cos(R.theta/2) -sin(R.theta/2); sin(R.theta/2) cos(R.theta/2)]
 # mat(R::RotationGate{1, T, ZGate{Complex{T}}}) where T =
 #     SMatrix{2, 2, Complex{T}}(cos(R.theta/2)-im*sin(R.theta/2), 0, 0, cos(R.theta/2)+im*sin(R.theta/2))
 
diff --git a/test/blocktools.jl b/test/blocktools.jl
index 3781b9d..e0431f8 100644
--- a/test/blocktools.jl
+++ b/test/blocktools.jl
@@ -47,9 +47,11 @@ end
     e2 = expect(put(2, 2=>X), reg |> copy |> focus!(1,2))
     e3 = expect(put(3, 2=>X), reg |> ρ)
     e4 = expect(put(3, 2=>X), reg)
+    e5 = expect(put(3, 2=>-X), reg)
     @test e1 ≈ e2
     @test e1 ≈ e3
     @test e1 ≈ e4
+    @test e5 ≈ -e4
 end
 
 @testset "viewbatch apply" begin
diff --git a/test/composite/tag.jl b/test/composite/tag.jl
index 0ead03d..ed9598f 100644
--- a/test/composite/tag.jl
+++ b/test/composite/tag.jl
@@ -1,4 +1,4 @@
-using Test, YaoBlocks
+using Test, YaoBlocks, YaoArrayRegister
 
 struct MockedTag{BT, N} <: TagBlock{BT, N}
     content::BT
@@ -18,3 +18,31 @@ end
 @test parameters(MockedTag(Rx(0.1))) == [0.1]
 
 @test occupied_locs(MockedTag(chain(3, put(1=>X), put(3=>X)))) == occupied_locs(chain(3, put(1=>X), put(3=>X)))
+
+@testset "scale apply" begin
+    # factor
+    @test factor(-X) == -1
+    @test factor(-2*X) == -2
+
+    # type
+    @test -X isa Scale{<:Val}
+    @test -im*X isa Scale{<:Val}
+    @test -2*X isa Scale{<:Number}
+
+    # apply!
+    reg = rand_state(1)
+    @test apply!(copy(reg), -X) ≈ -apply!(copy(reg), X)
+    @test apply!(copy(reg), (-im*Y)') ≈ apply!(copy(reg), im*Y)
+    @test apply!(copy(reg), (-2im*Y)') ≈ apply!(copy(reg), 2im*Y)
+
+    # mat
+    @test mat(-X) ≈ -mat(X)
+    @test mat((-im*Y)') ≈ (-im*mat(Y))'
+    @test mat((-2im*Y)') ≈ (-2im*mat(Y))'
+
+    # other
+    @test chsubblocks(-X, Y) == -Y
+    @test cache_key(Scale(Val(-1), X)) == cache_key(Scale(-1, X))
+    @test copy(-X) == -X
+    @test copy(-X) isa Scale{<:Val}
+end