Excitation number restricted (ENR) state space implementation

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

@@ -17,7 +17,7 @@ struct Dimensions{N,T<:Tuple} <: AbstractDimensions{N,N}
     to::T
 
     # make sure the elements in the tuple are all AbstractSpace
-    Dimensions(to::NTuple{N,T}) where {N,T<:AbstractSpace} = new{N,typeof(to)}(to)
+    Dimensions(to::NTuple{N,AbstractSpace}) where {N} = new{N,typeof(to)}(to)
 end
 function Dimensions(dims::Union{AbstractVector{T},NTuple{N,T}}) where {T<:Integer,N}
     _non_static_array_warning("dims", dims)
@@ -43,12 +43,11 @@ 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
 
     # make sure the elements in the tuple are all AbstractSpace
-    GeneralDimensions(to::NTuple{M,T1}, from::NTuple{N,T2}) where {M,N,T1<:AbstractSpace,T2<:AbstractSpace} =
+    GeneralDimensions(to::NTuple{M,AbstractSpace}, from::NTuple{N,AbstractSpace}) where {M,N} =
         new{M,N,typeof(to),typeof(from)}(to, from)
 end
 function GeneralDimensions(dims::Union{AbstractVector{T},NTuple{N,T}}) where {T<:Union{AbstractVector,NTuple},N}
@@ -98,3 +97,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

@@ -0,0 +1,191 @@
+#=
+This file defines the energy restricted space structure.
+=#
+
+export EnrSpace, enr_state_dictionaries
+export enr_identity, enr_fock, enr_destroy, enr_thermal_dm
+
+@doc raw"""
+    struct EnrSpace{N} <: AbstractSpace
+        size::Int
+        dims::NTuple{N,Int}
+        n_excitations::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, 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
+- `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
+
+# Example
+
+To construct an `EnrSpace`, we only need to specify the `dims` and `n_excitations`, namely
+
+```jldoctest
+julia> dims = (2, 2, 3);
+
+julia> n_excitations = 3;
+
+julia> EnrSpace(dims, n_excitations)
+EnrSpace((2, 2, 3), 3)
+```
+"""
+struct EnrSpace{N} <: AbstractSpace
+    size::Int
+    dims::SVector{N,Int}
+    n_excitations::Int
+    state2idx::Dict{SVector{N,Int},Int}
+    idx2state::Dict{Int,SVector{N,Int}}
+
+    function EnrSpace(dims::Union{AbstractVector{T},NTuple{N,T}}, excitations::Int) where {T<:Integer,N}
+        # all arguments will be checked in `enr_state_dictionaries`
+        size, state2idx, idx2state = enr_state_dictionaries(dims, excitations)
+
+        L = length(dims)
+        return new{L}(size, SVector{L}(dims), excitations, state2idx, idx2state)
+    end
+end
+
+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) = (s_enr1.size == s_enr2.size) && (s_enr1.dims == s_enr2.dims)
+
+dimensions_to_dims(s_enr::EnrSpace) = s_enr.dims
+
+@doc raw"""
+    enr_state_dictionaries(dims, excitations)
+
+Return the number of states, and lookup-dictionaries for translating a state (`SVector`) to a state index, and vice versa, for a system with a given number of components and maximum number of excitations.
+
+# Arguments
+- `dims::Union{AbstractVector,Tuple}`: A list of the number of states in each sub-system
+- `excitations::Int`: Maximum number of excitations
+
+# Returns
+- `nstates`: Number of states
+- `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
+"""
+function enr_state_dictionaries(dims::Union{AbstractVector{T},NTuple{N,T}}, excitations::Int) where {T<:Integer,N}
+    # argument checks
+    _non_static_array_warning("dims", dims)
+    L = length(dims)
+    (L > 0) || throw(DomainError(dims, "The argument dims must be of non-zero length"))
+    all(>=(1), dims) || throw(DomainError(dims, "All the elements of dims must be non-zero integers (≥ 1)"))
+    (excitations > 0) || throw(DomainError(excitations, "The argument excitations must be a non-zero integer (≥ 1)"))
+
+    nvec = zeros(Int, L) # Vector
+    nexc = 0
+
+    # in the following, all `nvec` (Vector) will first be converted (copied) to SVector and then push to `result` 
+    result = SVector{L,Int}[nvec]
+    while true
+        idx = L
+        nvec[end] += 1
+        nexc += 1
+        (nvec[idx] < dims[idx]) && push!(result, nvec)
+        while (nexc == excitations) || (nvec[idx] == dims[idx])
+            idx -= 1
+
+            # if idx < 1, break while-loop and return
+            if idx < 1
+                enr_size = length(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
+            (nvec[idx] < dims[idx]) && push!(result, nvec)
+        end
+    end
+end
+
+function enr_identity(dims::Union{AbstractVector{T},NTuple{N,T}}, excitations::Int) where {T<:Integer,N}
+    s_enr = EnrSpace(dims, excitations)
+    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}},
+    excitations::Int,
+    state::AbstractVector{T};
+    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
+        j = s_enr.state2idx[state]
+        array = [i == j ? 1.0 + 0im : 0.0 + 0im for i in 1:(s_enr.size)]
+    end
+
+    return QuantumObject(array, Ket(), s_enr)
+end
+
+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), N)
+
+    for (n1, state1) in idx2state
+        for (idx, s) in pairs(state1)
+            # 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 a_ops
+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, N)
+    else
+        (length(n) == N) || throw(ArgumentError("The length of the vector `n` should be the same as `dims`."))
+        nvec = n
+    end
+
+    D = s_enr.size
+    idx2state = s_enr.idx2state
+
+    diags = ComplexF64[prod((nvec ./ (1 .+ nvec)) .^ idx2state[idx]) for idx in 1:D]
+
+    diags /= sum(diags)
+
+    return QuantumObject(Diagonal(diags), Operator(), s_enr)
+end

src/qobj/quantum_object_base.jl

@@ -222,22 +222,28 @@ 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,19 +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))
-
-# 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,97 @@
+@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π
+        g = 0.1ωc
+        γ = 0.01ωc
+        tlist = range(0, 20, 100)
+        N_cut = 2
+
+        # normal mesolve and steadystate
+        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)
+        c_ops_JC = (√γ * a,)
+        ψ0_JC = basis(2, 0) ⊗ fock(N_cut, 0)
+        sol_JC = mesolve(H_JC, ψ0_JC, tlist, c_ops_JC; e_ops = [sz], progress_bar = Val(false))
+        ρ_ss_JC = steadystate(H_JC, c_ops_JC)
+
+        # ENR mesolve and steadystate
+        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)
+        c_ops_enr = (√γ * a_enr,)
+        sol_enr = mesolve(H_enr, ψ0_enr, tlist, c_ops_enr; e_ops = [sz_enr], progress_bar = Val(false))
+        ρ_ss_enr = steadystate(H_enr, c_ops_enr)
+
+        # check mesolve result
+        @test all(isapprox.(sol_JC.expect, sol_enr.expect, atol = 1e-4))
+
+        # check steadystate result
+        @test expect(sz, ρ_ss_JC) ≈ expect(sz_enr, ρ_ss_enr) atol=1e-4
+
+        # check eigenstates
+        λ, v = eigenstates(H_enr)
+        @test all([H_enr * v[k] ≈ λ[k] * v[k] for k in eachindex(λ)])
+    end
+
+    @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