From a2866396a726afd88c6e1f139cc5afea87428040 Mon Sep 17 00:00:00 2001
From: Leo <cacate0129@gmail.com>
Date: Sat, 15 Dec 2018 19:51:50 +0800
Subject: [PATCH 1/8] fix measure

---
 docs/src/man/blocks.md     | 7 -------
 docs/src/tutorial/QFT.md   | 2 +-
 src/Blocks/Diff.jl         | 6 +++++-
 src/Interfaces/TagBlock.jl | 2 +-
 src/Registers/measure.jl   | 6 ++++--
 test/Registers/measure.jl  | 4 ++++
 6 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/docs/src/man/blocks.md b/docs/src/man/blocks.md
index 118268e7d..68942dee6 100644
--- a/docs/src/man/blocks.md
+++ b/docs/src/man/blocks.md
@@ -16,13 +16,6 @@ a quantum circuit into serveral kinds of blocks. The uppermost abstract type for
 
 ## Composite Blocks
 
-### Roller
-
-[`Roller`](@ref) is a special pattern of quantum circuits. Usually is equivalent to a [`KronBlock`](@ref), but we can optimize
-the computation by rotate the tensor form of a quantum state and apply each small block on it each time.
-
-![Block-System](../assets/figures/roller.svg)
-
 ## Blocks
 
 ```@autodocs
diff --git a/docs/src/tutorial/QFT.md b/docs/src/tutorial/QFT.md
index 02167f4f3..66a5de352 100644
--- a/docs/src/tutorial/QFT.md
+++ b/docs/src/tutorial/QFT.md
@@ -9,7 +9,7 @@ using Yao
 # Control-R(k) gate in block-A
 A(i::Int, j::Int, k::Int) = control([i, ], j=>shift(2π/(1<<k)))
 # block-B
-B(n::Int, i::Int) = chain(i==j ? kron(i=>H) : A(j, i, j-i+1) for j = i:n)
+B(n::Int, i::Int) = chain(i==j ? put(i=>H) : A(j, i, j-i+1) for j = i:n)
 QFT(n::Int) = chain(n, B(n, i) for i = 1:n)
 
 # define QFT and IQFT block.
diff --git a/src/Blocks/Diff.jl b/src/Blocks/Diff.jl
index 7768e6654..7128c1a3f 100644
--- a/src/Blocks/Diff.jl
+++ b/src/Blocks/Diff.jl
@@ -60,7 +60,11 @@ chblock(cb::BPDiff, blk::MatrixBlock) = BPDiff(blk)
 
 @forward BPDiff.block mat
 function apply!(reg::AbstractRegister, df::BPDiff)
-    df.input = copy(reg)
+    if isdefined(df, :input)
+        copyto!(df.input, reg)
+    else
+        df.input = copy(reg)
+    end
     apply!(reg, parent(df))
     reg
 end
diff --git a/src/Interfaces/TagBlock.jl b/src/Interfaces/TagBlock.jl
index 1722ed635..f368d14d8 100644
--- a/src/Interfaces/TagBlock.jl
+++ b/src/Interfaces/TagBlock.jl
@@ -131,6 +131,6 @@ scale(x::Number) = blk -> scale(blk, x)
 staticscale(blk::MatrixBlock, x::Number) = StaticScale(blk, x)
 staticscale(x::Number) = blk -> staticscale(blk, x)
 
-as_weights(probs::AbstractVector{T}) where T = Weights(probs, T(1))
+as_weights(probs::AbstractVector) = Weights(probs, 1)
 
 include("Cache.jl")
diff --git a/src/Registers/measure.jl b/src/Registers/measure.jl
index a57894ba4..239f33842 100644
--- a/src/Registers/measure.jl
+++ b/src/Registers/measure.jl
@@ -75,17 +75,19 @@ end
 
 for FUNC in [:measure_reset!, :measure!, :measure]
     @eval function $FUNC(reg::AbstractRegister, locs; args...)
+        nbit = nactive(reg)
         focus!(reg, locs)
         res = $FUNC(reg; args...)
-        relax!(reg, locs)
+        relax!(reg, locs; nbit=nbit)
         res
     end
 end
 
 function measure_remove!(reg::AbstractRegister, locs)
+    nbit = nactive(reg)
     focus!(reg, locs)
     res = measure_remove!(reg)
-    relax!(reg)
+    relax!(reg; nbit=nbit)
     res
 end
 
diff --git a/test/Registers/measure.jl b/test/Registers/measure.jl
index 5e87abdc5..8dea7756c 100644
--- a/test/Registers/measure.jl
+++ b/test/Registers/measure.jl
@@ -26,4 +26,8 @@ end
     reg0 = copy(reg)
     res = measure_remove!(reg)
     @test select(reg0, res) |> normalize! ≈ reg
+
+    reg = rand_state(6,5) |> focus!((1:5)...)
+    measure_reset!(reg, 1)
+    @test nactive(reg) == 5
 end

From 31e26cdf0076e545254db4eeab00b2cb8939e300 Mon Sep 17 00:00:00 2001
From: Leo <cacate0129@gmail.com>
Date: Mon, 17 Dec 2018 14:02:37 +0800
Subject: [PATCH 2/8] measure operators, insert qubit

---
 src/Blocks/Blocks.jl                          |  2 +-
 src/Blocks/block_operations.jl                | 23 ++++++++++--
 src/Registers/Default.jl                      |  8 +++++
 src/Registers/Registers.jl                    |  2 +-
 src/Registers/measure.jl                      | 27 ++++++++++++++
 test/Blocks/Blocks.jl                         |  4 +--
 ...blockoperations.jl => block_operations.jl} | 11 ++++++
 test/Registers/Default.jl                     | 11 +++++-
 test/Registers/Registers.jl                   |  8 +++--
 test/Registers/measure.jl                     | 36 +++++++++++++++++--
 test/Registers/register_operations.jl         |  2 +-
 11 files changed, 121 insertions(+), 13 deletions(-)
 rename test/Blocks/{blockoperations.jl => block_operations.jl} (76%)

diff --git a/src/Blocks/Blocks.jl b/src/Blocks/Blocks.jl
index 2913764b2..ea136da9c 100644
--- a/src/Blocks/Blocks.jl
+++ b/src/Blocks/Blocks.jl
@@ -13,7 +13,7 @@ using LuxurySparse
 # import package APIs
 import ..Yao
 import ..Yao: DefaultType, nqubits, nactive, invorder
-import ..Registers: focus!, relax!, datatype
+import ..Registers: focus!, relax!, datatype, measure!, measure_reset!, measure_remove!
 import ..Intrinsics: ishermitian, isunitary, isreflexive, iscommute
 import CacheServers: update!, iscached, clear!, pull, iscacheable
 export clear! # TODO: rm this later
diff --git a/src/Blocks/block_operations.jl b/src/Blocks/block_operations.jl
index 7de596540..776b41016 100644
--- a/src/Blocks/block_operations.jl
+++ b/src/Blocks/block_operations.jl
@@ -103,10 +103,29 @@ function expect end
 
 #expect(op::AbstractBlock, reg::AbstractRegister) = sum(conj(reg |> statevec) .* (apply!(copy(reg), op) |> statevec), dims=1) |> vec
 #expect(op::AbstractBlock, reg::AbstractRegister{1}) = reg'*apply!(copy(reg), op)
-expect(op::AbstractBlock, reg::AbstractRegister) = reg'*apply!(copy(reg), op)
+#expect(op::AbstractBlock, reg::AbstractRegister) = reg'*apply!(copy(reg), op)
 
-expect(op::MatrixBlock, dm::DensityMatrix) = mapslices(x->sum(mat(op).*x)[], dm.state, dims=[1,2]) |> vec
+#expect(op::MatrixBlock, dm::DensityMatrix) = mapslices(x->sum(mat(op).*x)[], dm.state, dims=[1,2]) |> vec
 expect(op::MatrixBlock, dm::DensityMatrix{1}) = sum(mat(op).*dropdims(dm.state, dims=3))
+function expect(op::AbstractBlock, dm::DensityMatrix{B}) where B
+    mop = mat(op)
+    [tr(view(dm.state,:,:,i)*mop) for i=1:B]
+end
+
+expect(op::AbstractBlock, reg::AbstractRegister{1}) = reg'*apply!(copy(reg), op)
+
+function expect(op::AbstractBlock, reg::AbstractRegister{B}) where B
+    ket = apply!(copy(reg), op)
+    C = conj!(reshape(ket.state, :, B))
+    A = reshape(reg.state, :, B)
+    dropdims(sum(A.*C, dims=1), dims=1) |> conj
+end
+
+for FUNC in [:measure!, :measure_reset!, :measure_remove!]
+    @eval function $FUNC(op::AbstractBlock, reg::AbstractRegister; kwargs...) where B
+        $FUNC(eigen!(mat(op) |> Matrix), reg; kwargs...)
+    end
+end
 
 ################### AutoDiff Circuit ###################
 export gradient, backward!
diff --git a/src/Registers/Default.jl b/src/Registers/Default.jl
index 70f7c584c..c7e6a5813 100644
--- a/src/Registers/Default.jl
+++ b/src/Registers/Default.jl
@@ -137,6 +137,14 @@ end
 
 addbit!(n::Int) = r->addbit!(r, n)
 
+function insert_qubit!(reg::DefaultRegister{B}, loc::Int; nbit::Int=1) where B
+    na = nactive(reg)
+    focus!(reg, 1:loc-1)
+    reg2 = join(zero_state(nbit, B), reg) |> relax! |> focus!((1:na+nbit)...)
+    reg.state = reg2.state
+    reg
+end
+
 repeat(reg::DefaultRegister{B}, n::Int) where B = DefaultRegister{B*n}(hcat((reg.state for i=1:n)...,))
 
 ############ ConjDefaultRegister ##############
diff --git a/src/Registers/Registers.jl b/src/Registers/Registers.jl
index 27d87a1b9..089dababb 100644
--- a/src/Registers/Registers.jl
+++ b/src/Registers/Registers.jl
@@ -20,7 +20,7 @@ export nqubits, nactive, nremain, nbatch, basis, state, datatype, viewbatch
 export relaxedvec, statevec, hypercubic, rank3
 export focus!, relax!, oneto, probs, isnormalized, normalize!
 export AbstractRegister, Register, ConjRegister, RegOrConjReg, ConjDefaultRegister
-export invorder!, reorder!, addbit!, reset!
+export invorder!, reorder!, addbit!, reset!, insert_qubit!
 
 # factories
 export register, zero_state, product_state, rand_state, uniform_state
diff --git a/src/Registers/measure.jl b/src/Registers/measure.jl
index 239f33842..3e9aed78d 100644
--- a/src/Registers/measure.jl
+++ b/src/Registers/measure.jl
@@ -91,6 +91,33 @@ function measure_remove!(reg::AbstractRegister, locs)
     res
 end
 
+for FUNC in [:measure!, :measure_reset!, :measure_remove!]
+    rotback = FUNC == :measure! ? :(reg.state = V*reg.state) : :()
+    @eval function $FUNC(op::Eigen, reg::AbstractRegister{B}; kwargs...) where B
+        E, V = op
+        reg.state = V'*reg.state
+        res = $FUNC(reg; kwargs...)
+        $rotback
+        E[res.+1]
+    end
+end
+
+for FUNC in [:measure_reset!, :measure!]
+    @eval function $FUNC(op, reg::AbstractRegister, locs; args...)
+        focus!(reg, locs)
+        res = $FUNC(op, reg; args...)
+        relax!(reg, locs)
+        res
+    end
+end
+
+function measure_remove!(op, reg::AbstractRegister, locs)
+    focus!(reg, locs)
+    res = measure_remove!(op, reg)
+    relax!(reg)
+    res
+end
+
 """
     select!(reg::AbstractRegister, b::Integer) -> AbstractRegister
     select!(b::Integer) -> Function
diff --git a/test/Blocks/Blocks.jl b/test/Blocks/Blocks.jl
index 092093a46..21943896b 100644
--- a/test/Blocks/Blocks.jl
+++ b/test/Blocks/Blocks.jl
@@ -21,8 +21,8 @@ end
     include("Sequential.jl")
 end
 
-@testset "blockoperations" begin
-    include("blockoperations.jl")
+@testset "block operations" begin
+    include("block_operations.jl")
 end
 
 @testset "arithmatics" begin
diff --git a/test/Blocks/blockoperations.jl b/test/Blocks/block_operations.jl
similarity index 76%
rename from test/Blocks/blockoperations.jl
rename to test/Blocks/block_operations.jl
index 4de3aef40..2cdce8711 100644
--- a/test/Blocks/blockoperations.jl
+++ b/test/Blocks/block_operations.jl
@@ -30,4 +30,15 @@ end
     @test expect(op, dm) ≈ expect(op, reg)
 end
 
+@testset "expect" begin
+    reg = rand_state(3,10)
+    e1 = expect(put(2, 2=>X), reg |> copy |> focus!(1,2) |> ρ)
+    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)
+    @test e1 ≈ e2
+    @test e1 ≈ e3
+    @test e1 ≈ e4
+end
+
 include("linalg.jl")
diff --git a/test/Registers/Default.jl b/test/Registers/Default.jl
index 7353dbb6e..4431da466 100644
--- a/test/Registers/Default.jl
+++ b/test/Registers/Default.jl
@@ -1,9 +1,18 @@
 using Test, Random, LinearAlgebra, SparseArrays
+using Statistics: mean
 
 using Yao
 using Yao.Registers
 using Yao.Intrinsics
 
+@testset "insert_qubit!" begin
+    reg = rand_state(5, 10)
+    insert_qubit!(reg, 3, nbit=2)
+    @test reg |> nqubits == 7
+    @test expect(put(7, 3=>Z), reg) .|> tr |> mean ≈ 1
+    @test expect(put(7, 4=>Z), reg) .|> tr |> mean ≈ 1
+end
+
 @testset "Constructors" begin
     test_data = zeros(ComplexF32, 2^5, 3)
     reg = register(test_data)
@@ -81,7 +90,7 @@ end
     @test repeat(register(v1 ⊗ v2 ⊗ v3), 2) |> reorder!(3,2,1) ≈ repeat(register(v3 ⊗ v2 ⊗ v1), 2)
 end
 
-@testset "addbit" begin
+@testset "addbit!" begin
     reg = zero_state(3)
     @test addbit!(copy(reg), 3) == zero_state(6)
     reg = rand_state(3, 2)
diff --git a/test/Registers/Registers.jl b/test/Registers/Registers.jl
index ad9312601..892f5a91b 100644
--- a/test/Registers/Registers.jl
+++ b/test/Registers/Registers.jl
@@ -1,8 +1,5 @@
 using Test, Random, LinearAlgebra, SparseArrays
 
-@testset "default register" begin
-    include("Default.jl")
-end
 @testset "focus" begin
     include("focus.jl")
 end
@@ -18,3 +15,8 @@ end
 @testset "density matrix" begin
     include("DensityMatrix.jl")
 end
+
+@testset "default register" begin
+    include("Default.jl")
+end
+
diff --git a/test/Registers/measure.jl b/test/Registers/measure.jl
index 8dea7756c..cdd29ebff 100644
--- a/test/Registers/measure.jl
+++ b/test/Registers/measure.jl
@@ -1,4 +1,4 @@
-using Test, Random, LinearAlgebra, SparseArrays
+using Test, Random, LinearAlgebra, SparseArrays, Statistics
 
 using Yao
 using Yao.Registers
@@ -19,7 +19,7 @@ end
     reg = rand_state(4)
     res = measure_reset!(reg, (4,))
     result = measure(reg; nshot=10)
-    println(result)
+    #println(result)
     @test all(result .< 8)
 
     reg = rand_state(6) |> focus(1,4,3)
@@ -31,3 +31,35 @@ end
     measure_reset!(reg, 1)
     @test nactive(reg) == 5
 end
+
+@testset "op-measures" begin
+    reg = rand_state(6, 10)
+    op = repeat(3, X)
+
+    # measure!
+    reg2 = reg |> copy
+    res = measure!(op, reg2, 2:4)
+    res2 = measure!(op, reg2, 2:4)
+    @test size(res) == (10,)
+    @test res2 == res
+
+    # measure_reset!
+    reg2 = reg |> copy
+    res = measure_reset!(op, reg2, 2:4)
+    reg2 |> repeat(6, H, 2:4)
+    res2 = measure_reset!(op, reg2, 2:4)
+    @test size(res) == (10,) == size(res2)
+    @test all(res2 .== 1)
+
+    # measure_remove!
+    reg2 = reg |> copy
+    res = measure_remove!(op, reg2, 2:4)
+    reg2 |> repeat(6, H, 2:4)
+    @test size(res) == (10,)
+    @test nqubits(reg2) == 3
+
+    reg = repeat(register([1,-1]/sqrt(2.0)), 10)
+    @test measure!(X, reg) |> mean ≈ -1
+    reg = repeat(register([1.0,0]), 1000)
+    @test abs(measure!(X, reg) |> mean) < 0.1
+end
diff --git a/test/Registers/register_operations.jl b/test/Registers/register_operations.jl
index 7ed0d03ae..3ab84b3d1 100644
--- a/test/Registers/register_operations.jl
+++ b/test/Registers/register_operations.jl
@@ -1,6 +1,6 @@
 using Test, Random, LinearAlgebra, SparseArrays
 
-using Yao.Registers
+using Yao.Registers, Yao
 using Yao.Intrinsics
 
 @testset "reorder" begin

From 10354269ebbb4f2d88dcc315d296efb34c74477f Mon Sep 17 00:00:00 2001
From: Leo <cacate0129@gmail.com>
Date: Tue, 18 Dec 2018 00:25:37 +0800
Subject: [PATCH 3/8] fix two mini bugs

---
 src/Blocks/ChainBlock.jl | 2 +-
 src/Registers/measure.jl | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/Blocks/ChainBlock.jl b/src/Blocks/ChainBlock.jl
index 52515fd54..2df815383 100644
--- a/src/Blocks/ChainBlock.jl
+++ b/src/Blocks/ChainBlock.jl
@@ -51,7 +51,7 @@ eachindex(c::ChainBlock) = eachindex(c.blocks)
 subblocks(c::ChainBlock) = c.blocks
 addrs(c::ChainBlock) = ones(Int, length(c))
 usedbits(c::ChainBlock) = unique(vcat([usedbits(b) for b in subblocks(c)]...))
-chsubblocks(pb::ChainBlock, blocks) = ChainBlock(blocks)
+chsubblocks(pb::ChainBlock{N, T}, blocks) where {N, T} = length(blocks) == 0 ? ChainBlock{N, T}([]) : ChainBlock(blocks)
 @forward ChainBlock.blocks popfirst!, pop!
 
 isunitary(c::ChainBlock) = all(isunitary, c.blocks) || isunitary(mat(c))
diff --git a/src/Registers/measure.jl b/src/Registers/measure.jl
index 3e9aed78d..96d55b1da 100644
--- a/src/Registers/measure.jl
+++ b/src/Registers/measure.jl
@@ -87,7 +87,7 @@ function measure_remove!(reg::AbstractRegister, locs)
     nbit = nactive(reg)
     focus!(reg, locs)
     res = measure_remove!(reg)
-    relax!(reg; nbit=nbit)
+    relax!(reg; nbit=nbit-length(locs))
     res
 end
 

From 92ca1e59a6e6003414ef2f4f11ec281f0280ff4c Mon Sep 17 00:00:00 2001
From: Leo <cacate0129@gmail.com>
Date: Tue, 8 Jan 2019 20:05:53 +0800
Subject: [PATCH 4/8] optimize expect for Addblock

---
 src/Blocks/block_operations.jl  | 6 ++++++
 test/Blocks/block_operations.jl | 3 +++
 2 files changed, 9 insertions(+)

diff --git a/src/Blocks/block_operations.jl b/src/Blocks/block_operations.jl
index 776b41016..f9654cc25 100644
--- a/src/Blocks/block_operations.jl
+++ b/src/Blocks/block_operations.jl
@@ -121,6 +121,12 @@ function expect(op::AbstractBlock, reg::AbstractRegister{B}) where B
     dropdims(sum(A.*C, dims=1), dims=1) |> conj
 end
 
+function expect(op::AddBlock, reg::AbstractRegister)
+    sum(opi->expect(opi, reg), op)
+end
+
+expect(op::AddBlock, reg::AbstractRegister{1}) = invoke(expect, Tuple{AddBlock, AbstractRegister}, op, reg)
+
 for FUNC in [:measure!, :measure_reset!, :measure_remove!]
     @eval function $FUNC(op::AbstractBlock, reg::AbstractRegister; kwargs...) where B
         $FUNC(eigen!(mat(op) |> Matrix), reg; kwargs...)
diff --git a/test/Blocks/block_operations.jl b/test/Blocks/block_operations.jl
index 2cdce8711..27c241579 100644
--- a/test/Blocks/block_operations.jl
+++ b/test/Blocks/block_operations.jl
@@ -14,6 +14,9 @@ using LuxurySparse
     @test expect(obs2, ghz) ≈ 0
     @test expect(obs3, ghz) ≈ 1
 
+    @test expect(obs1+obs2+obs3, ghz) ≈ 1
+    @test expect(obs1+obs2+obs3, repeat(ghz, 3)) ≈ [1,1,1]
+
     @test blockfilter(ishermitian, chain(2, kron(2, X, P0), repeat(2, Rx(0), (1,2)), kron(2, 2=>Rz(0.3)))) == [kron(2, X, P0), X, P0, repeat(2, Rx(0), (1,2)), Rx(0)]
     @test blockfilter(b->ishermitian(b) && b isa PrimitiveBlock, chain(2, kron(2, X, P0), repeat(2, Rx(0), (1,2)), kron(2, 2=>Rz(0.3)))) == [X, P0, Rx(0)]
 end

From 81db877ec0a7b26d528dad0dde55d5da16dd106d Mon Sep 17 00:00:00 2001
From: Leo <cacate0129@gmail.com>
Date: Tue, 8 Jan 2019 21:01:45 +0800
Subject: [PATCH 5/8] update / and -

---
 src/Blocks/linalg.jl      | 3 +++
 test/Blocks/linalg.jl     | 3 +++
 test/Registers/Default.jl | 2 +-
 test/Registers/measure.jl | 3 ++-
 4 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/Blocks/linalg.jl b/src/Blocks/linalg.jl
index 73e3304e4..0288f8a16 100644
--- a/src/Blocks/linalg.jl
+++ b/src/Blocks/linalg.jl
@@ -16,6 +16,7 @@
 -(blk::AbstractScale) = chfactor(blk, -factor(blk))
 -(blk::MatrixBlock) = -1*blk
 -(blk::Neg) = blk.block
+-(A::MatrixBlock, B::MatrixBlock) = A + (-B)
 
 +(a::MatrixBlock{N}, b::MatrixBlock{N}) where N = AddBlock(a, b)
 +(a::AddBlock{N, T1}, b::MatrixBlock{N, T2}) where {N, T1, T2} = AddBlock{N, promote_type(T1, T2)}([a.blocks...; b])
@@ -43,3 +44,5 @@
 *(x::AddBlock{N, T1}, y::ChainBlock{N, T2}) where {N, T1, T2} = AddBlock{N, promote_type(T1, T2)}([b*y for b in subblocks(x)])
 
 Base.:^(blk::MatrixBlock, n::Int) = ChainBlock(fill(blk, n))
+
+/(A::MatrixBlock, x::Number) = (1/x)*A
diff --git a/test/Blocks/linalg.jl b/test/Blocks/linalg.jl
index 7a90548c8..13011ed8a 100644
--- a/test/Blocks/linalg.jl
+++ b/test/Blocks/linalg.jl
@@ -22,4 +22,7 @@ using Yao, Yao.Blocks
     @test g1*g2 isa ChainBlock
     @test (g1*g2)*(g3+g2) |> mat |> Matrix ≈ Matrix(mat(g1)*mat(g2)*(mat(g2)+mat(g3)))
     @test (g1+g3)*(g3+g2) |> mat ≈ (mat(g1)+mat(g3))*(mat(g2)+mat(g3))
+
+    @test mat(g1-g2) ≈ mat(g1+(-g2))
+    @test mat(g1/2) ≈ mat(0.5*g1)
 end
diff --git a/test/Registers/Default.jl b/test/Registers/Default.jl
index 4431da466..7edf7d7ea 100644
--- a/test/Registers/Default.jl
+++ b/test/Registers/Default.jl
@@ -1,5 +1,5 @@
 using Test, Random, LinearAlgebra, SparseArrays
-using Statistics: mean
+using StatsBase: mean
 
 using Yao
 using Yao.Registers
diff --git a/test/Registers/measure.jl b/test/Registers/measure.jl
index cdd29ebff..d0d1e5c39 100644
--- a/test/Registers/measure.jl
+++ b/test/Registers/measure.jl
@@ -1,4 +1,5 @@
-using Test, Random, LinearAlgebra, SparseArrays, Statistics
+using Test, Random, LinearAlgebra, SparseArrays
+using StatsBase: mean
 
 using Yao
 using Yao.Registers

From efaf4cb6d4c20e1328d804d2aae93201967d2e9f Mon Sep 17 00:00:00 2001
From: Leo <cacate0129@gmail.com>
Date: Thu, 10 Jan 2019 00:10:03 +0800
Subject: [PATCH 6/8] fix a trim induced bug

---
 src/Blocks/CachedBlock.jl       | 3 ++-
 src/Blocks/block_operations.jl  | 5 +++++
 test/Blocks/block_operations.jl | 2 ++
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/Blocks/CachedBlock.jl b/src/Blocks/CachedBlock.jl
index 19219d134..2eba7ddfd 100644
--- a/src/Blocks/CachedBlock.jl
+++ b/src/Blocks/CachedBlock.jl
@@ -39,7 +39,8 @@ clear!(c::CachedBlock) = (clear!(c.server, c.block); c)
 # forward methods
 function mat(c::CachedBlock)
     if !iscached(c.server, c.block)
-        m = dropzeros!(mat(c.block))
+        #m = dropzeros!(mat(c.block); trim=1e-16)
+        m = mat(c.block)
         push!(c.server, m, c.block)
         return m
     end
diff --git a/src/Blocks/block_operations.jl b/src/Blocks/block_operations.jl
index f9654cc25..f0b469dfa 100644
--- a/src/Blocks/block_operations.jl
+++ b/src/Blocks/block_operations.jl
@@ -125,7 +125,12 @@ function expect(op::AddBlock, reg::AbstractRegister)
     sum(opi->expect(opi, reg), op)
 end
 
+function expect(op::AbstractScale, reg::AbstractRegister)
+    factor(op)*expect(parent(op), reg)
+end
+
 expect(op::AddBlock, reg::AbstractRegister{1}) = invoke(expect, Tuple{AddBlock, AbstractRegister}, op, reg)
+expect(op::AbstractScale, reg::AbstractRegister{1}) = invoke(expect, Tuple{AbstractScale, AbstractRegister}, op, reg)
 
 for FUNC in [:measure!, :measure_reset!, :measure_remove!]
     @eval function $FUNC(op::AbstractBlock, reg::AbstractRegister; kwargs...) where B
diff --git a/test/Blocks/block_operations.jl b/test/Blocks/block_operations.jl
index 27c241579..59c7b8e3b 100644
--- a/test/Blocks/block_operations.jl
+++ b/test/Blocks/block_operations.jl
@@ -16,6 +16,8 @@ using LuxurySparse
 
     @test expect(obs1+obs2+obs3, ghz) ≈ 1
     @test expect(obs1+obs2+obs3, repeat(ghz, 3)) ≈ [1,1,1]
+    @test expect(2*obs3, ghz) ≈ 2
+    @test expect(2*obs3, repeat(ghz, 3)) ≈ [2,2,2]
 
     @test blockfilter(ishermitian, chain(2, kron(2, X, P0), repeat(2, Rx(0), (1,2)), kron(2, 2=>Rz(0.3)))) == [kron(2, X, P0), X, P0, repeat(2, Rx(0), (1,2)), Rx(0)]
     @test blockfilter(b->ishermitian(b) && b isa PrimitiveBlock, chain(2, kron(2, X, P0), repeat(2, Rx(0), (1,2)), kron(2, 2=>Rz(0.3)))) == [X, P0, Rx(0)]

From c5086a87c8c0960fa64ec1a0cb84c01002220094 Mon Sep 17 00:00:00 2001
From: Leo <cacate0129@gmail.com>
Date: Thu, 10 Jan 2019 13:07:57 +0800
Subject: [PATCH 7/8] update measure test

---
 test/Registers/measure.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/Registers/measure.jl b/test/Registers/measure.jl
index d0d1e5c39..829fe61c6 100644
--- a/test/Registers/measure.jl
+++ b/test/Registers/measure.jl
@@ -55,7 +55,7 @@ end
     # measure_remove!
     reg2 = reg |> copy
     res = measure_remove!(op, reg2, 2:4)
-    reg2 |> repeat(6, H, 2:4)
+    reg2 |> repeat(3, H, 2:3)
     @test size(res) == (10,)
     @test nqubits(reg2) == 3
 

From d2220ed86b4e0a423b699114c560f09beaf403cf Mon Sep 17 00:00:00 2001
From: Leo <cacate0129@gmail.com>
Date: Fri, 1 Feb 2019 02:01:32 +0800
Subject: [PATCH 8/8] fix focus

---
 src/Interfaces/Primitive.jl | 2 +-
 src/Registers/focus.jl      | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/Interfaces/Primitive.jl b/src/Interfaces/Primitive.jl
index ec650f38d..45f28f7ba 100644
--- a/src/Interfaces/Primitive.jl
+++ b/src/Interfaces/Primitive.jl
@@ -90,7 +90,7 @@ Returns a swap gate on `line1` and `line2`
 """
 function swap end
 
-swap(n::Int, ::Type{T}, line1::Int, line2::Int) where T = Swap{n, T}(line1, line2)
+swap(::Type{T}, n::Int, line1::Int, line2::Int) where T = Swap{n, T}(line1, line2)
 swap(::Type{T}, line1::Int, line2::Int) where T = n -> swap(n, T, line1, line2)
 swap(n::Int, line1::Int, line2::Int) = Swap{n, DefaultType}(line1, line2)
 swap(line1::Int, line2::Int) = n->swap(n, line1, line2)
diff --git a/src/Registers/focus.jl b/src/Registers/focus.jl
index 8f82163dd..828695966 100644
--- a/src/Registers/focus.jl
+++ b/src/Registers/focus.jl
@@ -58,6 +58,8 @@ end
 
 function focus!(reg::DefaultRegister{B}, bits) where B
     nbit = nactive(reg)
+    assert_addr_safe(nbit, [bit:bit for bit in bits] |> vec)
+
     if all(bits .== 1:length(bits))
         arr = reg.state
     else