Merge pull request #157 from ytdHuang/fix/states

Fix incorrect inputs for functions that generate states

src/qobj/states.jl

@@ -8,46 +8,49 @@ export rand_dm
 @doc raw"""
     fock(N::Int, pos::Int=0; dims::Vector{Int}=[N], sparse::Bool=false)
 
-Generates a fock state ``\ket{\psi}`` of dimension `N`. It is also possible
-to specify the list of dimensions `dims` if different subsystems are present.
+Generates a fock state ``\ket{\psi}`` of dimension `N`. 
+
+It is also possible to specify the list of dimensions `dims` if different subsystems are present.
 """
 function fock(N::Int, pos::Int = 0; dims::Vector{Int} = [N], sparse::Bool = false)
     if sparse
-        return QuantumObject(sparsevec([pos + 1], [1.0 + 0im], N), Ket, dims)
+        array = sparsevec([pos + 1], [1.0 + 0im], N)
     else
         array = zeros(ComplexF64, N)
         array[pos+1] = 1
-        return QuantumObject(array, Ket, dims)
     end
+    return QuantumObject(array; type = Ket, dims = dims)
 end
 
-"""
+@doc raw"""
     basis(N::Int, pos::Int = 0; dims::Vector{Int}=[N])
 
 Generates a fock state like [`fock`](@ref).
+
+It is also possible to specify the list of dimensions `dims` if different subsystems are present.
 """
 basis(N::Int, pos::Int = 0; dims::Vector{Int} = [N]) = fock(N, pos, dims = dims)
 
 @doc raw"""
-    coherent(N::Real, α::T)
+    coherent(N::Real, α::Number)
 
-Generates a coherent state ``\ket{\alpha}``, which is defined as an eigenvector of the
-bosonic annihilation operator ``\hat{a} \ket{\alpha} = \alpha \ket{\alpha}``.
+Generates a coherent state ``\ket{\alpha}``, which is defined as an eigenvector of the bosonic annihilation operator ``\hat{a} \ket{\alpha} = \alpha \ket{\alpha}``.
 """
 function coherent(N::Real, α::T) where {T<:Number}
     a = destroy(N)
     return exp(α * a' - α' * a) * fock(N, 0)
 end
 
 @doc raw"""
-    rand_dm(N::Integer; kwargs...)
+    rand_dm(N::Integer; dims::Vector{Int}=[N])
+
+Generates a random density matrix ``\hat{\rho}``, with the property to be positive semi-definite and ``\textrm{Tr} \left[ \hat{\rho} \right] = 1``.
 
-Generates a random density matrix ``\hat{\rho}``, with the property to be positive definite,
-and that ``\Tr \left[ \hat{\rho} \right] = 1``.
+It is also possible to specify the list of dimensions `dims` if different subsystems are present.
 """
-function rand_dm(N::Integer; kwargs...)
+function rand_dm(N::Integer; dims::Vector{Int} = [N])
     ρ = rand(ComplexF64, N, N)
     ρ *= ρ'
     ρ /= tr(ρ)
-    return QuantumObject(ρ; kwargs...)
+    return QuantumObject(ρ; type = Operator, dims = dims)
 end

test/states_and_operators.jl

@@ -1,4 +1,23 @@
 @testset "States and Operators" begin
+    # fock, basis
+    @test fock(4; dims = [2, 2], sparse = true) ≈ basis(4; dims = [2, 2])
+    @test_throws DimensionMismatch fock(4; dims = [2])
+
+    # rand_dm
+    ρ_AB = rand_dm(4, dims = [2, 2])
+    ρ_A = ptrace(ρ_AB, 1)
+    ρ_B = ptrace(ρ_AB, 2)
+    @test tr(ρ_AB) ≈ 1.0
+    @test tr(ρ_A) ≈ 1.0
+    @test tr(ρ_B) ≈ 1.0
+    @test ishermitian(ρ_AB) == true
+    @test ishermitian(ρ_A) == true
+    @test ishermitian(ρ_B) == true
+    @test all(eigenenergies(ρ_AB) .>= 0)
+    @test all(eigenenergies(ρ_A) .>= 0)
+    @test all(eigenenergies(ρ_B) .>= 0)
+    @test_throws DimensionMismatch rand_dm(4, dims = [2])
+
     # test commutation relations for fermionic creation and annihilation operators
     sites = 4
     SIZE = 2^sites