fix tensor and GeneralDimensions for EnrSpace

Author: ytdHuang

src/qobj/dimensions.jl

@@ -43,7 +43,6 @@ Dimensions(dims::Any) = throw(
 A structure that describes the left-hand side (`to`) and right-hand side (`from`) Hilbert [`Space`](@ref) of an [`Operator`](@ref).
 """
 struct GeneralDimensions{M,N,T1<:Tuple,T2<:Tuple} <: AbstractDimensions{M,N}
-    # note that the number `N` should be the same for both `to` and `from`
     to::T1   # space acting on the left
     from::T2 # space acting on the right
 

src/qobj/energy_restricted.jl

@@ -10,19 +10,17 @@ export enr_identity, enr_fock, enr_destroy, enr_thermal_dm
         size::Int
         dims::NTuple{N,Int}
         n_excitations::Int
-        n_subspace::Int
         state2idx::Dict{SVector{N,Int},Int}
         idx2state::Dict{Int,SVector{N,Int}}
     end
 
-A structure that describes an excitation-number restricted (ENR) state space.
+A structure that describes an excitation-number restricted (ENR) state space, where `N` is the number of sub-systems.
 
 # Fields
 
 - `size`: Number of states in the excitation-number restricted state space
 - `dims`: A list of the number of states in each sub-system
 - `n_excitations`: Maximum number of excitations
-- `n_subspace`: The number of sub-systems
 - `state2idx`: A dictionary for looking up a state index from a state (`SVector`)
 - `idx2state`: A dictionary for looking up state (`SVector`) from a state index
 
@@ -43,7 +41,6 @@ struct EnrSpace{N} <: AbstractSpace
     size::Int
     dims::SVector{N,Int}
     n_excitations::Int
-    n_subspace::Int
     state2idx::Dict{SVector{N,Int},Int}
     idx2state::Dict{Int,SVector{N,Int}}
 
@@ -52,7 +49,7 @@ struct EnrSpace{N} <: AbstractSpace
         size, state2idx, idx2state = enr_state_dictionaries(dims, excitations)
 
         L = length(dims)
-        return new{L}(size, SVector{L}(dims), excitations, L, state2idx, idx2state)
+        return new{L}(size, SVector{L}(dims), excitations, state2idx, idx2state)
     end
 end
 
@@ -116,8 +113,9 @@ end
 
 function enr_identity(dims::Union{AbstractVector{T},NTuple{N,T}}, excitations::Int) where {T<:Integer,N}
     s_enr = EnrSpace(dims, excitations)
-    return QuantumObject(Diagonal(ones(ComplexF64, s_enr.size)), Operator(), Dimensions(s_enr))
+    return enr_identity(s_enr)
 end
+enr_identity(s_enr::EnrSpace) = QuantumObject(Diagonal(ones(ComplexF64, s_enr.size)), Operator(), Dimensions(s_enr))
 
 function enr_fock(
     dims::Union{AbstractVector{T},NTuple{N,T}}, 
@@ -126,6 +124,13 @@ function enr_fock(
     sparse::Union{Bool,Val} = Val(false),
 ) where {T<:Integer,N}
     s_enr = EnrSpace(dims, excitations)
+    return enr_fock(s_enr, state, sparse = sparse)
+end
+function enr_fock(
+    s_enr::EnrSpace,
+    state::AbstractVector{T};
+    sparse::Union{Bool,Val} = Val(false),
+) where {T<:Integer}
     if getVal(sparse)
         array = sparsevec([s_enr.state2idx[[state...]]], [1.0 + 0im], s_enr.size)
     else
@@ -136,13 +141,19 @@ function enr_fock(
     return QuantumObject(array, Ket(), s_enr)
 end
 
-function enr_destroy(dims::Union{AbstractVector{T},NTuple{N,T}}, excitations::Int) where {T<:Integer,N}
+function enr_destroy(
+    dims::Union{AbstractVector{T},NTuple{N,T}}, 
+    excitations::Int
+) where {T<:Integer,N}
     s_enr = EnrSpace(dims, excitations)
+    return enr_destroy(s_enr)
+end
+function enr_destroy(s_enr::EnrSpace{N}) where {N}
     D = s_enr.size
     idx2state = s_enr.idx2state
     state2idx = s_enr.state2idx
 
-    a_ops = ntuple(i -> QuantumObject(spzeros(ComplexF64, D, D), Operator(), s_enr), s_enr.n_subspace)
+    a_ops = ntuple(i -> QuantumObject(spzeros(ComplexF64, D, D), Operator(), s_enr), N)
 
     for (n1, state1) in idx2state
         for (idx, s) in pairs(state1)
@@ -162,11 +173,13 @@ end
 
 function enr_thermal_dm(dims::Union{AbstractVector{T1},NTuple{N,T1}}, excitations::Int, n::Union{T2,AbstractVector{T2}}) where {T1<:Integer,T2<:Real,N}
     s_enr = EnrSpace(dims, excitations)
-
+    return enr_thermal_dm(s_enr, n)
+end
+function enr_thermal_dm(s_enr::EnrSpace{N}, n::Union{T,AbstractVector{T}}) where {N,T<:Real}
     if n isa Real
-        nvec = fill(n, s_enr.n_subspace)
+        nvec = fill(n, N)
     else
-        (length(n) == s_enr.n_subspace) || throw(ArgumentError("The length of the vector `n` should be the same as `dims`."))
+        (length(n) == N) || throw(ArgumentError("The length of the vector `n` should be the same as `dims`."))
         nvec = n
     end
 

src/qobj/quantum_object_base.jl

@@ -222,22 +222,27 @@ end
 
 # this returns `to` in GeneralDimensions representation
 get_dimensions_to(A::AbstractQuantumObject{Ket,<:Dimensions}) = A.dimensions.to
-get_dimensions_to(A::AbstractQuantumObject{Bra,<:Dimensions{N}}) where {N} = space_one_list(N)
+get_dimensions_to(A::AbstractQuantumObject{Bra,<:Dimensions}) = space_one_list(A.dimensions.to)
 get_dimensions_to(A::AbstractQuantumObject{Operator,<:Dimensions}) = A.dimensions.to
 get_dimensions_to(A::AbstractQuantumObject{Operator,<:GeneralDimensions}) = A.dimensions.to
 get_dimensions_to(
     A::AbstractQuantumObject{ObjType,<:Dimensions},
 ) where {ObjType<:Union{SuperOperator,OperatorBra,OperatorKet}} = A.dimensions.to
 
 # this returns `from` in GeneralDimensions representation
-get_dimensions_from(A::AbstractQuantumObject{Ket,<:Dimensions{N}}) where {N} = space_one_list(N)
+get_dimensions_from(A::AbstractQuantumObject{Ket,<:Dimensions}) = space_one_list(A.dimensions.to)
 get_dimensions_from(A::AbstractQuantumObject{Bra,<:Dimensions}) = A.dimensions.to
 get_dimensions_from(A::AbstractQuantumObject{Operator,<:Dimensions}) = A.dimensions.to
 get_dimensions_from(A::AbstractQuantumObject{Operator,<:GeneralDimensions}) = A.dimensions.from
 get_dimensions_from(
     A::AbstractQuantumObject{ObjType,<:Dimensions},
 ) where {ObjType<:Union{SuperOperator,OperatorBra,OperatorKet}} = A.dimensions.to
 
+# this creates a list of Space(1), it is used to generate `from` for Ket, and `to` for Bra
+space_one_list(dimensions::NTuple{N,AbstractSpace}) where {N} = ntuple(i -> Space(1), Val(sum(_get_dims_length, dimensions)))
+_get_dims_length(::Space) = 1
+_get_dims_length(::EnrSpace{N}) where {N} = N
+
 # functions for getting Float or Complex element type
 _FType(A::AbstractQuantumObject) = _FType(eltype(A))
 _CType(A::AbstractQuantumObject) = _CType(eltype(A))

src/qobj/space.jl

@@ -23,6 +23,3 @@ struct Space <: AbstractSpace
 end
 
 dimensions_to_dims(s::Space) = SVector{1,Int}(s.size)
-
-# this creates a list of Space(1), it is used to generate `from` for Ket, and `to` for Bra)
-space_one_list(N::Int) = ntuple(i -> Space(1), Val(N))

test/core-test/enr_state_operator.jl

@@ -1,4 +1,49 @@
 @testitem "Excitation number restricted state space" begin
+    @testset "kron" begin
+        # normal Space
+        D1 = 4
+        D2 = 5
+        dims_s = (D1, D2)
+        ρ_s = rand_dm(dims_s)
+        I_s = qeye(D1) ⊗ qeye(D2)
+        size_s = prod(dims_s)
+        space_s = (Space(D1), Space(D2))
+        
+        # EnrSpace
+        dims_enr = (2, 2, 3)
+        excitations = 3
+        space_enr = EnrSpace(dims_enr, excitations)
+        ρ_enr = enr_thermal_dm(space_enr, rand(3))
+        I_enr = enr_identity(space_enr)
+        size_enr = space_enr.size
+
+        # tensor between normal and ENR space
+        ρ_tot = tensor(ρ_s, ρ_enr)
+        opstring = sprint((t, s) -> show(t, "text/plain", s), ρ_tot)
+        datastring = sprint((t, s) -> show(t, "text/plain", s), ρ_tot.data)
+        ρ_tot_dims = [dims_s..., dims_enr...]
+        ρ_tot_size = size_s * size_enr
+        ρ_tot_isherm = isherm(ρ_tot)
+        @test opstring ==
+              "\nQuantum Object:   type=Operator()   dims=$ρ_tot_dims   size=$((ρ_tot_size, ρ_tot_size))   ishermitian=$ρ_tot_isherm\n$datastring"
+
+        # use GeneralDimensions to do partial trace
+        new_dims1 = GeneralDimensions((Space(1), Space(1), space_enr), (Space(1), Space(1), space_enr))
+        ρ_enr_compound = Qobj(zeros(ComplexF64, size_enr, size_enr), dims = new_dims1)
+        basis_list = [tensor(basis(D1, i), basis(D2, j)) for i in 0:(D1-1) for j in 0:(D2-1)]
+        for b in basis_list
+            ρ_enr_compound += tensor(b', I_enr) * ρ_tot * tensor(b, I_enr)
+        end
+        new_dims2 = GeneralDimensions((space_s..., Space(1), Space(1), Space(1)), (space_s..., Space(1), Space(1), Space(1)))
+        ρ_s_compound = Qobj(zeros(ComplexF64, size_s, size_s), dims = new_dims2)
+        basis_list = [enr_fock(space_enr, space_enr.idx2state[idx]) for idx in 1:space_enr.size]
+        for b in basis_list
+            ρ_s_compound += tensor(I_s, b') * ρ_tot * tensor(I_s, b)
+        end
+        @test ρ_enr.data ≈ ρ_enr_compound.data
+        @test ρ_s.data ≈ ρ_s_compound.data
+    end
+
     @testset "mesolve, steadystate, and eigenstates" begin
         ε = 2π
         ωc = 2π
@@ -37,7 +82,15 @@
         # check eigenstates
         λ, v = eigenstates(H_enr)
         @test all([H_enr * v[k] ≈ λ[k] * v[k] for k in eachindex(λ)])
+    end
 
-        qeye(3) ⊗ H_enr
+    @testset "Type Inference" begin
+        N = 3
+        dims = (2, 2, 3)
+        excitations = 3
+        @inferred enr_identity(dims, excitations)
+        @inferred enr_fock(dims, excitations, zeros(Int, N))
+        @inferred enr_destroy(dims, excitations)
+        @inferred enr_thermal_dm(dims, excitations, rand(N))
     end
 end