Add more doc test and examples (#40)

* fix a typo

* polish layouts

* add more examples

* add more docs

* add more doc test

* update measure layout

src/layout.jl

@@ -161,6 +161,7 @@ color(::Type{<:MathGate}) = :red
 color(::Type{<:PutBlock}) = :cyan
 color(::Type{T}) where {T <: PauliString} = :cyan
 color(::Type{<:RepeatedBlock}) = :cyan
+color(::Type{<:GeneralMatrixBlock}) = :red
 
 # color(::Type{T}) where {T <: PauliString} = :cyan
 # color(::Type{T}) where {T <: Sequential} = :blue
@@ -172,12 +173,41 @@ print_block(io::IO, swap::Swap) = printstyled(io, "swap", swap.locs; bold=true,
 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::Roller) = printstyled(io, "roller"; bold=true, color=color(Roller))
-print_block(io::IO, x::ReflectGate{N}) where N = print(io, "reflect: nqubits=$N")
+print_block(io::IO, x::ReflectGate{N}) where N = print(io, "reflect($(summary(x.psi)))")
 print_block(io::IO, c::Concentrator) = print(io, "Concentrator: ", occupied_locs(c))
 print_block(io::IO, c::CachedBlock) = print_block(io, content(c))
 print_block(io::IO, c::Prod) = printstyled(io, "prod"; bold=true, color=color(ChainBlock))
 print_block(io::IO, c::Sum) = printstyled(io, "sum"; bold=true, color=color(ChainBlock))
 print_block(io::IO, c::TagBlock) = nothing
+print_block(io::IO, c::GeneralMatrixBlock) = printstyled(io, "matblock(...)"; color=color(GeneralMatrixBlock))
+
+function print_block(io::IO, c::Measure{N, K, OT}) where {N, K, OT}
+    strs = String[]
+    if c.operator != ComputationalBasis()
+        push!(strs, "operator=$(repr(c.operator))")
+    end
+
+    if c.locations != AllLocs()
+        push!(strs, "locs=$(repr(c.locations))")
+    end
+
+    if c.collapseto !== nothing
+        push!(strs, "collapseto=$(c.collapseto)")
+    end
+
+    if c.remove
+        push!(strs, "remove=true")
+    end
+
+    out = join(strs, ", ")
+    if !isempty(strs)
+        out = "Measure($N;" * out
+    else
+        out = "Measure($N" * out
+    end
+
+    return print(io, out, ")")
+end
 
 # TODO: use OhMyREPL's default syntax highlighting for functions
 function print_block(io::IO, m::MathGate{N, <:LegibleLambda}) where N

src/primitive/general_matrix_gate.jl

@@ -30,7 +30,8 @@ Create a [`GeneralMatrixBlock`](@ref) with a matrix `m`.
 # Example
 
 ```jldoctest
-
+julia> matblock(ComplexF64[0 1;1 0])
+matblock(...)
 ```
 """
 matblock(m::AbstractMatrix) = GeneralMatrixBlock(m)

src/primitive/measure.jl

@@ -1,9 +1,8 @@
 using YaoBase, YaoArrayRegister
-export Measure
+export Measure, AllLocs, ComputationalBasis
 
 """
     Measure{N, K, OT} <: PrimitiveBlock{N, Bool}
-    Measure(n::Int; operator=ComputationalBasis(), locs=1:n, collapseto=nothing, remove=false)
 
 Measure operator.
 """
@@ -24,7 +23,71 @@ end
 
 @interface nqubits_measured(::Measure{N, K}) where {N, K} = K
 
-function Measure(n::Int; operator::OT=ComputationalBasis(), locs, collapseto=nothing, remove=false) where OT
+"""
+    Measure(n::Int; operator=ComputationalBasis(), locs=AllLocs(), collapseto=nothing, remove=false)
+
+Create a `Measure` block with number of qubits `n`.
+
+# Example
+
+You can create a `Measure` block on given basis (default is the computational basis).
+
+```jldoctest
+julia> Measure(4)
+Measure(4)
+```
+
+Or you could specify which qubits you are going to measure
+
+```jldoctest
+julia> Measure(4; locs=1:3)
+Measure(4; locs=(1, 2, 3))
+```
+
+by default this will collapse the current register to measure results.
+
+```julia
+julia> r = rand_state(3)
+ArrayReg{1, Complex{Float64}, Array...}
+    active qubits: 3/3
+
+julia> state(r)
+8×1 Array{Complex{Float64},2}:
+  0.19864933724343375 - 0.4740335956912438im
+  -0.2057912765333517 - 0.2262668486124923im
+ -0.41680007712245676 - 0.13759187422609387im
+ -0.24336704548326407 + 0.27343538360398184im
+ -0.09308092255704317 - 0.005308959093704435im
+  0.24555464152683212 + 0.02737969837364506im
+  -0.3828287267256825 + 0.02401578941643196im
+ 0.048647936794205926 + 0.31047610497928607im
+
+julia> r |> Measure(3)
+Measure(3)
+
+julia> state(r)
+8×1 Array{Complex{Float64},2}:
+                 0.0 + 0.0im
+                 0.0 + 0.0im
+ -0.9495962023170939 - 0.31347576069762273im
+                 0.0 + 0.0im
+                 0.0 + 0.0im
+                 0.0 + 0.0im
+                 0.0 + 0.0im
+                 0.0 + 0.0im
+```
+
+But you can also specify the target bit configuration you want to collapse to with keyword `collapseto`.
+
+```jldoctest
+julia> Measure(4; collapseto=0b101)
+Measure(4;collapseto=5)
+
+julia> m.collapseto
+5
+```
+"""
+function Measure(n::Int; operator::OT=ComputationalBasis(), locs=AllLocs(), collapseto=nothing, remove=false) where OT
     if locs isa AllLocs
         Measure{n, n, OT}(operator, locs, collapseto, remove)
     else

src/primitive/phase_gate.jl

@@ -14,7 +14,20 @@ end
 """
     phase(theta)
 
-Returns a global phase gate.
+Returns a global phase gate. Defined with following matrix form:
+
+```math
+exp(iθ) \\mathbf{I}
+```
+
+# Example
+
+You can create a global phase gate with a phase (a real number).
+
+```jldoctest
+julia> phase(0.1)
+phase(0.1)
+```
 """
 phase(θ::AbstractFloat) = PhaseGate(θ)
 phase(θ::Real) = phase(Float64(θ))

src/primitive/reflect_gate.jl

@@ -42,6 +42,13 @@ reflect(r::ArrayReg) = reflect(statevec(r))
     reflect(v::AbstractVector{<:Complex})
 
 Create a [`ReflectGate`](@ref) with an quantum state vector `v`.
+
+# Example
+
+```jldoctest
+julia> reflect(rand_state(3))
+reflect(ArrayReg{1, Complex{Float64}, Array...})
+```
 """
 reflect(v::AbstractVector{<:Complex}) = ReflectGate(v)
 

src/primitive/rotation_gate.jl

@@ -6,6 +6,12 @@ export RotationGate, Rx, Ry, Rz, rot
     RotationGate{N, T, GT <: AbstractBlock{N, Complex{T}}} <: PrimitiveBlock{N, Complex{T}}
 
 RotationGate, with GT both hermitian and isreflexive.
+
+# Definition
+
+```math
+\\mathbf{I} cos(θ / 2) - im sin(θ / 2) * mat(U)
+```
 """
 mutable struct RotationGate{N, T, GT <: AbstractBlock{N, Complex{T}}} <: PrimitiveBlock{N, Complex{T}}
     block::GT
@@ -24,20 +30,41 @@ RotationGate(block::GT, theta) where {N, T, GT<:AbstractBlock{N, Complex{T}}} =
     Rx(theta)
 
 Return a [`RotationGate`](@ref) on X axis.
+
+# Example
+
+```jldoctest
+julia> Rx(0.1)
+rot(X gate, 0.1)
+```
 """
 Rx(theta::T) where T <: AbstractFloat = RotationGate(X(Complex{T}), theta)
 
 """
     Ry(theta)
 
 Return a [`RotationGate`](@ref) on Y axis.
+
+# Example
+
+```jldoctest
+julia> Ry(0.1)
+rot(Y gate, 0.1)
+```
 """
 Ry(theta::T) where T <: AbstractFloat = RotationGate(Y(Complex{T}), theta)
 
 """
     Rz(theta)
 
 Return a [`RotationGate`](@ref) on Z axis.
+
+# Example
+
+```jldoctest
+julia> Rz(0.1)
+rot(Z gate, 0.1)
+```
 """
 Rz(theta::T) where T <: AbstractFloat = RotationGate(Z(Complex{T}), theta)
 

src/primitive/shift_gate.jl

@@ -6,6 +6,15 @@ export ShiftGate, shift
     ShiftGate <: PrimitiveBlock
 
 Phase shift gate.
+
+# Definition
+
+```math
+\\begin{pmatrix}
+1 & 0
+0 & e^(im θ)
+\\end{pmatrix}
+```
 """
 mutable struct ShiftGate{T} <: PrimitiveBlock{1, Complex{T}}
     theta::T
@@ -14,7 +23,14 @@ end
 """
     shift(θ)
 
-Returns a shift gate.
+Create a [`ShiftGate`](@ref) with phase `θ`.
+
+# Example
+
+```jldoctest
+julia> shift(0.1)
+shift(0.1)
+```
 """
 shift(θ::AbstractFloat) = ShiftGate(θ)
 shift(θ::Real) = shift(Float64(θ))

src/primitive/swap_gate.jl

@@ -3,6 +3,11 @@ using YaoArrayRegister: swaprows!
 
 export Swap, swap
 
+"""
+    Swap{N, T} <: PrimitiveBlock{N, T}
+
+Swap block, which will swap two locations given.
+"""
 struct Swap{N, T} <: PrimitiveBlock{N, T}
     locs::Tuple{Int, Int}
 
@@ -20,15 +25,29 @@ swap(::Type{T}, n::Int, loc1::Int, loc2::Int) where T = Swap{n, T}((loc1, loc2))
 """
     swap([T=ComplexF64], n, loc1, loc2)
 
-Return a `n`-qubit [`Swap`](@ref) gate which swap `loc1` and `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(ComplexF64, n, loc1, loc2)
 
 """
     swap(loc1, loc2) -> f(n)
 
-Return a lambda that takes the total number of active qubits as input. See also
-[`swap`](@ref), [`Swap`](@ref).
+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))