From e1fbab2a812e7307f071a372fbd800ef70e78ab9 Mon Sep 17 00:00:00 2001 From: cailixun Date: Fri, 27 Jun 2025 15:48:36 +0800 Subject: [PATCH 1/2] basic enr objects --- src/qobj/dimensions.jl | 1 + src/qobj/energy_restricted.jl | 123 ++++++++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+) create mode 100644 src/qobj/energy_restricted.jl diff --git a/src/qobj/dimensions.jl b/src/qobj/dimensions.jl index ba72659e1..0b4617b42 100644 --- a/src/qobj/dimensions.jl +++ b/src/qobj/dimensions.jl @@ -78,6 +78,7 @@ 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 diff --git a/src/qobj/energy_restricted.jl b/src/qobj/energy_restricted.jl new file mode 100644 index 000000000..e60fc214e --- /dev/null +++ b/src/qobj/energy_restricted.jl @@ -0,0 +1,123 @@ +struct EnrSpace{N} <: AbstractSpace + size::Int + dims::SVector{N,Int} + n_excitations::Int + state2idx::Dict{} + idx2state::Dict{} + + 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) + len = length(dims) + nvec = MVector{len}(zeros(Int, len)) + result = [copy(nvec)] + nexc = 0 + + while true + idx = len + nvec[end] += 1 + nexc += 1 + if nvec[idx] < dims[idx] + push!(result, copy(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, copy(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 = [zeros(ComplexF64, N, N) for _ in 1: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 + end + end + + return [QuantumObject(array, Operator(), s_enr) for array in 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 From 3e6545989287c798a5daf8cb6fae43965ceb171e Mon Sep 17 00:00:00 2001 From: cailixun Date: Fri, 27 Jun 2025 22:04:59 +0800 Subject: [PATCH 2/2] fix bug and add test --- src/QuantumToolbox.jl | 1 + src/qobj/dimensions.jl | 6 +- src/qobj/energy_restricted.jl | 82 ++++++++++++++++------------ src/qobj/space.jl | 13 ----- test/core-test/enr_state_operator.jl | 28 ++++++++++ 5 files changed, 80 insertions(+), 50 deletions(-) create mode 100644 test/core-test/enr_state_operator.jl diff --git a/src/QuantumToolbox.jl b/src/QuantumToolbox.jl index 915c4abe8..a4d0b7841 100644 --- a/src/QuantumToolbox.jl +++ b/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") diff --git a/src/qobj/dimensions.jl b/src/qobj/dimensions.jl index 0b4617b42..92dfc0bd9 100644 --- a/src/qobj/dimensions.jl +++ b/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 diff --git a/src/qobj/energy_restricted.jl b/src/qobj/energy_restricted.jl index e60fc214e..4a9b5d2d1 100644 --- a/src/qobj/energy_restricted.jl +++ b/src/qobj/energy_restricted.jl @@ -1,26 +1,41 @@ +#= +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 @@ -28,47 +43,44 @@ function enr_state_dictionaries(dims::Union{Tuple, AbstractVector}, excitations: 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 diff --git a/src/qobj/space.jl b/src/qobj/space.jl index 0b90ecf55..f5cb1cd8e 100644 --- a/src/qobj/space.jl +++ b/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 -=# diff --git a/test/core-test/enr_state_operator.jl b/test/core-test/enr_state_operator.jl new file mode 100644 index 000000000..eb5e0a949 --- /dev/null +++ b/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