use new bit_str literal as interface & printing, add more features fo…

…r mixed state simulation (#25)

* fix-reg type

* update measure

* fix collapseto

* update for BitStr

* two new APIs: purify and exchange_sysenv

Project.toml

@@ -13,7 +13,7 @@ TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 YaoBase = "a8f54c17-34bc-5a9d-b050-f522fe3f755f"
 
 [compat]
-BitBasis = "^0.5.0"
+BitBasis = "^0.6.0"
 LuxurySparse = "~0.4.0"
 YaoBase = "~0.11.0"
 julia = "^1.0.0"

src/density_matrix.jl

@@ -1,4 +1,5 @@
 export DensityMatrix, density_matrix, ρ
+export purify
 
 """
     DensityMatrix{B, T, MT}
@@ -60,3 +61,22 @@ function YaoBase.probs(m::DensityMatrix{B, T}) where {B, T}
 end
 
 YaoBase.probs(m::DensityMatrix{1})= diag(view(m.state, :,:,1))
+
+"""
+    purify(r::DensityMatrix{B}; nbit_env::Int=nactive(r)) -> ArrayReg
+
+Get a purification of target density matrix.
+"""
+function purify(r::DensityMatrix{B}; nbit_env::Int=nactive(r)) where B
+    Ne = 1<<nbit_env
+    Ns = size(r.state,1)
+    state = similar(r.state,Ns,Ne,B)
+    for ib in 1:B
+        R, U = eigen!(r.state[:,:,ib])
+        state[:,:,ib] .= view(U,:,Ns-Ne+1:Ns) .* sqrt.(abs.(view(R,Ns-Ne+1:Ns)'))
+    end
+    return ArrayReg(state)
+end
+
+# obtaining matrix from Yao.DensityMatrix{1}, `1` is the batch size.
+LinearAlgebra.Matrix(d::DensityMatrix{1}) = dropdims(d.state, dims=3)

src/focus.jl

@@ -2,7 +2,8 @@ using YaoBase, TupleTools
 
 export focus!,
     relax!,
-    partial_tr
+    partial_tr,
+    exchange_sysenv
 
 
 """
@@ -125,3 +126,12 @@ function YaoBase.partial_tr(r::ArrayReg{B}, locs) where B
     relax!(r, orders)
     return normalize!(ArrayReg(state))
 end
+
+"""
+    exchange_sysenv(reg::ArrayReg) -> ArrayReg
+
+Exchange system (focused qubits) and environment (remaining qubits).
+"""
+function exchange_sysenv(reg::ArrayReg{B}) where B
+    ArrayReg{B}(reshape(permutedims(rank3(reg), (2,1,3)), :,size(reg.state, 1)*B))
+end

src/instruct.jl

@@ -17,11 +17,6 @@ A list of symbol for specialized gates/operators.
 """
 const SPECIALIZATION_LIST = Symbol[:X, :Y, :Z, :S, :T, :Sdag, :Tdag]
 
-# to avoid potential ambiguity, we limit them to tuple for now
-# but they only has to be an iterator over integers
-const Locations{T} = NTuple{N, T} where N
-const BitConfigs{T} = NTuple{N, T} where N
-
 function YaoBase.instruct!(
     state::AbstractVecOrMat{T1},
     operator::AbstractMatrix{T2},

src/measure.jl

@@ -6,10 +6,14 @@ export measure,
     select,
     select!
 
-_measure(rng::AbstractRNG, pl::AbstractVector, nshots::Int) = sample(rng, 0:length(pl)-1, Weights(pl), nshots)
+function _measure(rng::AbstractRNG, pl::AbstractVector, nshots::Int)
+    N = log2i(length(pl))
+    sample(rng, basis(BitStr64{N}), Weights(pl), nshots)
+end
+
 function _measure(rng::AbstractRNG, pl::AbstractMatrix, nshots::Int)
     B = size(pl, 2)
-    res = Matrix{Int}(undef, nshots, B)
+    res = Matrix{BitStr64{log2i(length(pl))}}(undef, nshots, B)
     for ib=1:B
         @inbounds res[:,ib] = _measure(rng, view(pl,:,ib), nshots)
     end
@@ -27,10 +31,11 @@ function YaoBase.measure_remove!(rng::AbstractRNG, ::ComputationalBasis, reg::Ar
     state = reg |> rank3
     nstate = similar(reg.state, 1<<nremain(reg), B)
     pl = dropdims(sum(state .|> abs2, dims=2), dims=2)
-    res = Vector{Int}(undef, B)
+    res = Vector{BitStr64{nactive(reg)}}(undef, B)
     @inbounds for ib = 1:B
         ires = _measure(rng, view(pl, :, ib), 1)[]
-        nstate[:,ib] = view(state, ires+1,:,ib)./sqrt(pl[ires+1, ib])
+        # notice ires is `BitStr` type, can be use as indices directly.
+        nstate[:,ib] = view(state, Int64(ires)+1,:,ib)./sqrt(pl[Int64(ires)+1, ib])
         res[ib] = ires
     end
     reg.state = reshape(nstate,1,:)
@@ -43,7 +48,7 @@ function YaoBase.measure!(rng::AbstractRNG, ::ComputationalBasis, reg::ArrayReg{
     res = measure_remove!(rng, reg)
     _nstate = reshape(reg.state, :, B)
     for ib in 1:B
-        @inbounds nstate[res[ib]+1, :, ib] .= view(_nstate, :,ib)
+        @inbounds nstate[Int64(res[ib])+1, :, ib] .= view(_nstate, :,ib)
     end
     reg.state = reshape(nstate, size(state, 1), :)
     return res
@@ -54,22 +59,18 @@ function YaoBase.measure_collapseto!(rng::AbstractRNG, ::ComputationalBasis, reg
     M, N, B1 = size(state)
     nstate = zero(state)
     res = measure_remove!(rng, reg)
-    nstate[config+1, :, :] = reshape(reg.state, :, B)
+    nstate[Int(config)+1, :, :] = reshape(reg.state, :, B)
     reg.state = reshape(nstate, M, N*B)
     return res
 end
 
-YaoBase.select(r::ArrayReg{B}, bits) where B = ArrayReg{B}(r.state[map(to_location, bits), :])
-YaoBase.select(r::ArrayReg{B}, bit::Union{Integer, BitStr}) where B = select(r, [bit])
-
-function YaoBase.select!(r::ArrayReg, bits)
-    r.state = r.state[map(to_location, bits), :]
-    return r
-end
+import YaoBase: select, select!
+select(r::ArrayReg{B}, bits::AbstractVector{T}) where {B, T<:Integer} = ArrayReg{B}(r.state[Int64.(bits) .+ 1, :])
+select(r::ArrayReg{B}, bit::Integer) where B = select(r, [bit])
 
-function YaoBase.select!(r::ArrayReg, bits::Integer)
-    r.state = reshape(r.state[bits+1, :], 1, :)
+function select!(r::ArrayReg, bits::AbstractVector{T}) where T<:Integer
+    r.state = r.state[Int64.(bits) .+ 1, :]
     return r
 end
 
-YaoBase.select!(r::ArrayReg, bits::BitStr) = select!(r, bits.val)
+select!(r::ArrayReg, bit::Integer) = select!(r, [bit])

src/register.jl

@@ -1,5 +1,5 @@
 using YaoBase, BitBasis
-import BitBasis: BitStr
+import BitBasis: BitStr, BitStr64
 
 export ArrayReg,
     AdjointArrayReg,
@@ -188,9 +188,9 @@ function YaoBase.reorder!(r::ArrayReg, orders)
     return r
 end
 
-function YaoBase.collapseto!(r::ArrayReg, bit_config::Integer=0)
+function YaoBase.collapseto!(r::ArrayReg, bit_config::BitStr=0)
     fill!(r.state, 0)
-    @inbounds r.state[bit_config+1,:] .= 1
+    r.state[Int64(bit_config)+1,:] .= 1
     return r
 end
 

test/density_matrix.jl

@@ -39,3 +39,21 @@ end
     @test isapprox(tracedist(dm, dm_)[], tracedist(dm4, dm5)[], atol=1e-5)
     @test isapprox.(tracedist(dm, dm_)[], tracedist(repeat(reg4, 3)|>density_matrix, repeat(reg5, 3)|>density_matrix), atol=1e-5) |> all
 end
+
+@testset "purify" begin
+    reg = rand_state(6)
+    reg_p = purify(reg |> ρ)
+    @test reg_p |> isnormalized
+    @test reg_p |> exchange_sysenv |> probs |> maximum ≈ 1
+    reg_p = purify(reg |> ρ; nbit_env=0)
+    @test Yao.fidelity(reg, reg_p) ≈ [1]
+
+    reg = rand_state(6; nbatch=10)
+    reg_p = purify(reg |> ρ)
+    @test reg_p |> isnormalized
+    @test reg_p |> exchange_sysenv |> probs |> maximum ≈ 1
+    reg_p = purify(reg |> ρ; nbit_env=0)
+    @test Yao.fidelity(reg, reg_p) ≈ ones(10)
+    reg_p = purify(reg |> ρ; nbit_env=2)
+    @test reg_p |> nqubits == 8
+end