fix bug and add test

Author: cailixun

src/QuantumToolbox.jl

@@ -84,6 +84,7 @@ include("linear_maps.jl")
 
 # Quantum Object
 include("qobj/space.jl")
+include("qobj/energy_restricted.jl")
 include("qobj/dimensions.jl")
 include("qobj/quantum_object_base.jl")
 include("qobj/quantum_object.jl")

src/qobj/dimensions.jl

@@ -78,7 +78,6 @@ dimensions_to_dims(dimensions::GeneralDimensions) =
     SVector{2}(dimensions_to_dims(dimensions.to), dimensions_to_dims(dimensions.from))
 
 dimensions_to_dims(::Nothing) = nothing # for EigsolveResult.dimensions = nothing
-dimensions_to_dims(s_enr::EnrSpace) = s_enr.dims
 
 Base.length(::AbstractDimensions{N}) where {N} = N
 
@@ -99,3 +98,8 @@ function _get_dims_string(dimensions::GeneralDimensions)
     return "[$(string(dims[1])), $(string(dims[2]))]"
 end
 _get_dims_string(::Nothing) = "nothing" # for EigsolveResult.dimensions = nothing
+
+Base.:(==)(dim1::Dimensions, dim2::Dimensions) = (dim1.to == dim2.to)
+Base.:(==)(dim1::GeneralDimensions, dim2::GeneralDimensions) = (dim1.to == dim2.to) && (dim1.from == dim2.from)
+Base.:(==)(dim1::Dimensions, dim2::GeneralDimensions) = false
+Base.:(==)(dim1::GeneralDimensions, dim2::Dimensions) = false

src/qobj/energy_restricted.jl

@@ -1,74 +1,86 @@
+#=
+This file defines the energy restricted space structure.
+=#
+
+export EnrSpace, enr_state_dictionaries
+export enr_identity, enr_fock, enr_destroy, enr_thermal_dm
+
 struct EnrSpace{N} <: AbstractSpace
     size::Int
-    dims::SVector{N,Int}
+    dims::NTuple{N,Int}
     n_excitations::Int
-    state2idx::Dict{}
-    idx2state::Dict{}
+    state2idx::Dict{SVector{N,Int},Int}
+    idx2state::Dict{Int,SVector{N,Int}}
 
-    function EnrSpace(dims::Union{Tuple, AbstractVector}, excitations::Int)
+    function EnrSpace(dims::Union{Tuple,AbstractVector}, excitations::Int)
         _non_static_array_warning("dims", dims)
         dim_len = length(dims)
         dims_T = NTuple{dim_len}(dims)
 
         size, state2idx, idx2state = enr_state_dictionaries(dims, excitations)
-        
 
         return new{dim_len}(size, dims_T, excitations, state2idx, idx2state)
     end
 end
 
-function enr_state_dictionaries(dims::Union{Tuple, AbstractVector}, excitations::Int)
+function Base.show(io::IO, s::EnrSpace)
+    print(io, "EnrSpace($(s.dims), $(s.n_excitations))")
+    return nothing
+end
+
+Base.:(==)(s_enr1::EnrSpace, s_enr2::EnrSpace) = (all([s_enr1.size, s_enr1.dims] .== [s_enr2.size, s_enr2.dims]))
+
+dimensions_to_dims(s_enr::EnrSpace) = s_enr.dims
+
+function enr_state_dictionaries(dims::Union{Tuple,AbstractVector}, excitations::Int)
     len = length(dims)
-    nvec = MVector{len}(zeros(Int, len))
-    result = [copy(nvec)]
+    nvec = zeros(Int, len)
+    result = SVector{len,Int}[nvec] # in the following, all nvec will first be converted (copied) to SVector and then push to result 
     nexc = 0
 
     while true
         idx = len
         nvec[end] += 1
         nexc += 1
         if nvec[idx] < dims[idx]
-            push!(result, copy(nvec))
+            push!(result, nvec)
         end
         while (nexc == excitations) || (nvec[idx] == dims[idx])
             #nvec[idx] = 0
             idx -= 1
             if idx < 1
                 enr_size = length(result)
-                return (
-                    enr_size,
-                    Dict(zip(result, 1:enr_size)),
-                    Dict(zip(1:enr_size, result))
-                )
+                return (enr_size, Dict(zip(result, 1:enr_size)), Dict(zip(1:enr_size, result)))
             end
 
             nexc -= nvec[idx+1] - 1
             nvec[idx+1] = 0
             nvec[idx] += 1
             if nvec[idx] < dims[idx]
-                push!(result, copy(nvec))
+                push!(result, nvec)
             end
         end
     end
-
 end
 
-function enr_identity(dims::Union{Tuple, AbstractVector}, excitations::Int)
+function enr_identity(dims::Union{Tuple,AbstractVector}, excitations::Int)
     s_enr = EnrSpace(dims, excitations)
     return QuantumObject(Diagonal(ones(ComplexF64, s_enr.size)), Operator(), Dimensions(s_enr))
 end
 
 function enr_fock(
-        dims::Union{Tuple, AbstractVector}, excitations::Int, state::AbstractVector; 
-        sparse::Union{Bool,Val} = Val(false)
-    )
+    dims::Union{Tuple,AbstractVector},
+    excitations::Int,
+    state::AbstractVector;
+    sparse::Union{Bool,Val} = Val(false),
+)
     s_enr = EnrSpace(dims, excitations)
     if getVal(sparse)
         array = sparsevec([s_enr.state2idx[[state...]]], [1.0 + 0im], s_enr.size)
     else
         j = s_enr.state2idx[state]
         array = [i == j ? 1.0 + 0im : 0.0 + 0im for i in 1:(s_enr.size)]
-        
+
         # s = zeros(ComplexF64, s_enr.size)
         # s[s_enr.state2idx[state]] += 1
         # s
@@ -77,40 +89,38 @@ function enr_fock(
     return QuantumObject(array, Ket(), s_enr)
 end
 
-function enr_destroy(dims::Union{Tuple, AbstractVector}, excitations::Int)
+function enr_destroy(dims::Union{Tuple,AbstractVector}, excitations::Int)
     s_enr = EnrSpace(dims, excitations)
     N = s_enr.size
     idx2state = s_enr.idx2state
     state2idx = s_enr.state2idx
 
-    a_ops = [zeros(ComplexF64, N, N) for _ in 1:length(dims)]
+    a_ops = ntuple(i -> QuantumObject(spzeros(ComplexF64, N, N), Operator(), s_enr), length(dims))
 
     for (n1, state1) in idx2state
         for (idx, s) in pairs(state1)
-            s > 0 || continue
-
-            state2 = copy(state1)
-            state2[idx] -= 1
-            n2 = state2idx[state2]
-            a_ops[idx][n2, n1] += √s
+            # if s > 0, the annihilation operator of mode idx has a non-zero
+            # entry with one less excitation in mode idx in the final state
+            if s > 0
+                state2 = Vector(state1)
+                state2[idx] -= 1
+                n2 = state2idx[state2]
+                a_ops[idx][n2, n1] = √s
+            end
         end
     end
 
-    return [QuantumObject(array, Operator(), s_enr) for array in a_ops]
+    return a_ops
 end
 
-function enr_thermal_dm(
-        dims::Union{Tuple, AbstractVector}, 
-        excitations::Int,
-        n::Union{Int, AbstractVector}
-    )
+function enr_thermal_dm(dims::Union{Tuple,AbstractVector}, excitations::Int, n::Union{Int,AbstractVector})
     if n isa Number
         nvec = Vector{typeof(n)}(n, length(dims))
     else
         length(n) == length(dims) || throw(ArgumentError("The Vector `n` has different length to `dims`."))
         nvec = n
     end
-        
+
     s_enr = EnrSpace(dims, excitations)
     N = s_enr.size
     idx2state = s_enr.idx2state

src/qobj/space.jl

@@ -26,16 +26,3 @@ 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))
-
-# TODO: introduce energy restricted space
-#=
-struct EnrSpace{N} <: AbstractSpace
-    size::Int
-    dims::SVector{N,Int}
-    n_excitations::Int
-    state2idx
-    idx2state
-end
-
-dimensions_to_dims(s::EnrSpace) = s.dims
-=#

test/core-test/enr_state_operator.jl

@@ -0,0 +1,28 @@
+@testitem "Excitation number restricted state space" begin
+    @testset "mesolve" begin
+        ε = 2π
+        ωc = 2π
+        g = 0.1ωc
+        γ = 0.01ωc
+        tlist = range(0, 20, 100)
+        N_cut = 2
+
+        sz = sigmaz() ⊗ qeye(N_cut)
+        sm = sigmam() ⊗ qeye(N_cut)
+        a = qeye(2) ⊗ destroy(N_cut)
+
+        H_JC = 0.5ε * sz + ωc * a' * a + g * (sm' * a + a' * sm)
+        ψ0 = basis(2, 0) ⊗ fock(N_cut, 0)
+        sol_JC = mesolve(H_JC, ψ0, tlist, [√γ * a]; e_ops = [sz], progress_bar = Val(false))
+
+        N_exc = 1
+        dims = [2, N_cut]
+        sm_enr, a_enr = enr_destroy(dims, N_exc)
+        sz_enr = 2 * sm_enr' * sm_enr - 1
+        ψ0_enr = enr_fock(dims, N_exc, [1, 0])
+        H_enr = ε * sm_enr' * sm_enr + ωc * a_enr' * a_enr + g * (sm_enr' * a_enr + a_enr' * sm_enr)
+        sol_enr = mesolve(H_enr, ψ0_enr, tlist, [√γ * a_enr]; e_ops = [sz_enr], progress_bar = Val(false))
+
+        @test all(isapprox.(sol_JC.expect, sol_enr.expect, atol = 1e-4))
+    end
+end