First try of ENR 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

@@ -98,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

@@ -0,0 +1,133 @@
+#=
+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::NTuple{N,Int}
+    n_excitations::Int
+    state2idx::Dict{SVector{N,Int},Int}
+    idx2state::Dict{Int,SVector{N,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 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 = 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, 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)))
+            end
+
+            nexc -= nvec[idx+1] - 1
+            nvec[idx+1] = 0
+            nvec[idx] += 1
+            if nvec[idx] < dims[idx]
+                push!(result, nvec)
+            end
+        end
+    end
+end
+
+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),
+)
+    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
+    end
+
+    return QuantumObject(array, Ket(), s_enr)
+end
+
+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 = ntuple(i -> QuantumObject(spzeros(ComplexF64, N, N), Operator(), s_enr), length(dims))
+
+    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{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
+
+    diags = [prod((nvec ./ (1 .+ nvec)) .^ idx2state[idx]) for idx in 1:N]
+
+    diags /= sum(diags)
+
+    return QuantumObject(Diagonal(diags), Operator(), s_enr)
+end

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