new product state and density matrix interfaces (#89)

* new product state and density matrix interceas

* add reduced density matrix test

* fix tests

Project.toml

@@ -1,6 +1,6 @@
 name = "YaoArrayRegister"
 uuid = "e600142f-9330-5003-8abb-0ebd767abc51"
-version = "0.7.9"
+version = "0.7.10"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/density_matrix.jl

@@ -19,6 +19,10 @@ state(ρ::DensityMatrix) = ρ.state
 YaoBase.nqubits(ρ::DensityMatrix) = log2dim1(state(ρ))
 YaoBase.nactive(ρ::DensityMatrix) = nqubits(ρ)
 YaoBase.nbatch(dm::DensityMatrix{B}) where {B} = B
+function YaoBase.density_matrix(reg::ArrayReg, qubits)
+    freg = focus!(copy(reg), qubits)
+    return density_matrix(freg)
+end
 YaoBase.density_matrix(reg::ArrayReg{1}) = DensityMatrix(reg.state * reg.state')
 function YaoBase.density_matrix(reg::ArrayReg{B}) where {B}
     M = size(reg.state, 1)

src/register.jl

@@ -314,14 +314,24 @@ julia> product_state(bit"100"; nbatch=2)
 ArrayReg{2, Complex{Float64}, Array...}
     active qubits: 3/3
 
-julia> product_state(ComplexF32, bit"101"; nbatch=2)
+julia> r1 = product_state(ComplexF32, bit"100"; nbatch=2)
 ArrayReg{2, Complex{Float32}, Array...}
     active qubits: 3/3
+
+julia> r2 = product_state(ComplexF32, [0, 0, 1]; nbatch=2)
+ArrayReg{2, Complex{Float32}, Array...}
+    active qubits: 3/3
+
+julia> r1 ≈ r2   # because we read bit strings from right to left, vectors from left to right.
+true
 ```
 """
 product_state(bit_str::BitStr; nbatch::Int = 1) =
     product_state(ComplexF64, bit_str; nbatch = nbatch)
 
+product_state(bit_str::AbstractVector; nbatch::Int = 1) =
+    product_state(ComplexF64, bit_str; nbatch = nbatch)
+
 """
     product_state([T=ComplexF64], total::Int, bit_config::Integer; nbatch=1, no_transpose_storage=false)
 
@@ -355,6 +365,9 @@ product_state(total::Int, bit_config::Integer; kwargs...) =
 product_state(::Type{T}, bit_str::BitStr{N}; kwargs...) where {T,N} =
     product_state(T, N, buffer(bit_str); kwargs...)
 
+product_state(::Type{T}, bit_configs::AbstractVector; kwargs...) where T = 
+    product_state(T, bit_literal(bit_configs...); kwargs...)
+
 function product_state(
     ::Type{T},
     total::Int,

test/density_matrix.jl

@@ -69,14 +69,23 @@ end
     @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]
+    @test 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)
+    @test fidelity(reg, reg_p) ≈ ones(10)
     reg_p = purify(reg |> ρ; nbit_env = 2)
     @test reg_p |> nqubits == 8
 end
+
+@testset "reduce density matrix" begin
+    reg = (product_state(bit"00000") + product_state(bit"11111")) / sqrt(2)
+    rdm = density_matrix(reg, (1,))
+    @test Matrix(rdm) ≈ [1/2 0; 0 1/2]
+    reg = product_state([1,0,0])
+    rdm = density_matrix(reg, (1,2))
+    @test Matrix(rdm) ≈ [0 0 0 0; 0 1 0 0; 0 0 0 0; 0 0 0 0]
+end

test/register.jl

@@ -32,6 +32,8 @@ end
     @testset "test product state" begin
         st = state(product_state(T, bit"100"; nbatch = 1))
         @test !(st isa Transpose)
+        st2 = state(product_state(T, [0, 0, 1]; nbatch = 1))
+        @test st2 ≈ st
         st = state(product_state(T, bit"100"; nbatch = 2, no_transpose_storage = true))
         @test !(st isa Transpose)
         st = state(product_state(T, bit"100"; nbatch = 2))

test/runtests.jl

@@ -18,3 +18,7 @@ end
 @testset "test measure" begin
     include("measure.jl")
 end
+
+@testset "test density matrix" begin
+    include("density_matrix.jl")
+end