Implement mean occupation and get coherence functions

src/qobj/arithmetic_and_attributes.jl

@@ -6,7 +6,9 @@ Arithmetic and Attributes for QuantumObject
 
 export proj, ptrace, purity, permute
 export tidyup, tidyup!
-export get_data, get_coherence
+export get_data, get_coherence, remove_coherence, mean_occupation
+
+import Base: _ind2sub
 
 #    Broadcasting
 Base.broadcastable(x::QuantumObject) = x.data
@@ -695,26 +697,156 @@ Returns the data of a [`AbstractQuantumObject`](@ref).
 get_data(A::AbstractQuantumObject) = A.data
 
 @doc raw"""
-    get_coherence(ψ::QuantumObject)
+    get_coherence(ψ::QuantumObject; idx)
+
+Returns the coherence value ``\alpha`` by measuring the expectation value of the destruction operator ``\hat{a}`` on a state ``\ket{\psi}``, density matrix ``\hat{\rho}`` or vectorized state ``|\hat{\rho}\rangle\rangle``.
+
+If the state is a composite state, the coherence values of each subsystems are returned as an array of complex numbers. It is possible to specify the index of the sybsystem to get the coherence value by fixing `idx`.
+"""
+function get_coherence(ψ::QuantumObject{T,KetQuantumObject,N}; idx::Union{Nothing,Int} = nothing) where {T,N}
+    dims = reverse(Tuple(ψ.dims))
+    offsets = cumprod([1; dims[1:end-1]...])
+
+    expectation_values = zeros(ComplexF64, N)
+
+    for J in 2:length(ψ.data)
+        indices = Base._ind2sub(dims, J)
+        for n in 1:N
+            if indices[n] > 1
+                expectation_values[n] += conj(ψ.data[J-offsets[n]]) * sqrt(indices[n] - 1) * ψ.data[J]
+            end
+        end
+    end
+
+    reverse!(expectation_values)
+
+    return isnothing(idx) ? expectation_values : expectation_values[idx]
+end
+
+function get_coherence(ρ::QuantumObject{T,OperatorQuantumObject,N}; idx::Union{Int,Nothing} = nothing) where {T,N}
+    dims = reverse(Tuple(ρ.dims))
+    offsets = cumprod([1; dims[1:end-1]...])
+
+    expectation_values = zeros(ComplexF64, N)
+
+    @inbounds for J in 2:size(ρ.data, 1)
+        indices = Base._ind2sub(dims, J)
+        for n in 1:N
+            if indices[n] > 1
+                expectation_values[n] += sqrt(indices[n] - 1) * ρ.data[J, J-offsets[n]]
+            end
+        end
+    end
+
+    reverse!(expectation_values)
+
+    return isnothing(idx) ? expectation_values : expectation_values[idx]
+end
+
+function get_coherence(v::QuantumObject{T,OperatorKetQuantumObject,N}; idx::Union{Int,Nothing} = nothing) where {T,N}
+    dims = reverse(Tuple(v.dims))
+
+    offsets = cumprod([1; dims[1:end-1]...])
+    D = prod(dims)
+
+    expectation_values = zeros(ComplexF64, N)
+
+    @inbounds for J in 2:D
+        indices = Base._ind2sub(dims, J)
+        for n in 1:N
+            if indices[n] > 1
+                expectation_values[n] += sqrt(indices[n] - 1) * v.data[(J-offsets[n]-1)*D+J]
+            end
+        end
+    end
+
+    reverse!(expectation_values)
+
+    return isnothing(idx) ? expectation_values : expectation_values[idx]
+end
+
+get_coherence(ψ::QuantumObject{T,KetQuantumObject,1}) where {T} =
+    mapreduce(n -> sqrt(n - 1) * ψ.data[n] * conj(ψ.data[n-1]), +, 2:ψ.dims[1])
+get_coherence(ρ::QuantumObject{T,OperatorQuantumObject,1}) where {T} =
+    mapreduce(n -> sqrt(n - 1) * ρ.data[n, n-1], +, 2:ρ.dims[1])
+get_coherence(v::QuantumObject{T,OperatorKetQuantumObject,1}) where {T} =
+    mapreduce(n -> sqrt(n - 1) * v.data[(n-2)*v.dims[1]+n], +, 2:v.dims[1])
+
+@doc raw"""
+    remove_coherence(ψ::QuantumObject)
+
+Returns the corresponding state with the coherence removed: ``\ket{\delta_\alpha} = \exp ( \alpha^* \hat{a} - \alpha \hat{a}^\dagger ) \ket{\psi}`` for a pure state, and ``\hat{\rho_\alpha} = \exp ( \alpha^* \hat{a} - \alpha \hat{a}^\dagger ) \hat{\rho} \exp ( -\bar{\alpha} \hat{a} + \alpha \hat{a}^\dagger )`` for a density matrix. These states correspond to the quantum fluctuations around the coherent state ``\ket{\alpha}`` or ``|\alpha\rangle\langle\alpha|``.
+"""
+function remove_coherence(ψ::QuantumObject{T,KetQuantumObject}) where {T}
+    αs = get_coherence(ψ)
+    D = mapreduce(k -> displace(ψ.dims[k], αs[k]), ⊗, eachindex(ψ.dims))
+
+    return D' * ψ
+end
 
-Get the coherence value ``\alpha`` by measuring the expectation value of the destruction operator ``\hat{a}`` on a state ``\ket{\psi}`` or a density matrix ``\hat{\rho}``.
+function remove_coherence(ρ::QuantumObject{T,OperatorQuantumObject}) where {T}
+    αs = get_coherence(ρ)
+    D = mapreduce(k -> displace(ρ.dims[k], αs[k]), ⊗, eachindex(ρ.dims))
 
-It returns both ``\alpha`` and the corresponding state with the coherence removed: ``\ket{\delta_\alpha} = \exp ( \alpha^* \hat{a} - \alpha \hat{a}^\dagger ) \ket{\psi}`` for a pure state, and ``\hat{\rho_\alpha} = \exp ( \alpha^* \hat{a} - \alpha \hat{a}^\dagger ) \hat{\rho} \exp ( -\bar{\alpha} \hat{a} + \alpha \hat{a}^\dagger )`` for a density matrix. These states correspond to the quantum fluctuations around the coherent state ``\ket{\alpha}`` or ``|\alpha\rangle\langle\alpha|``.
+    return D' * ρ * D
+end
+
+@doc raw"""
+    mean_occupation(ψ::QuantumObject)
+
+Get the mean occupation number ``n`` by measuring the expectation value of the number operator ``\hat{n}`` on a state ``\ket{\psi}``, a density matrix ``\hat{\rho}`` or a vectorized density matrix ``\ket{\hat{\rho}}``.
+
+If the state is a composite state, the mean occupation numbers of each subsystems are returned as an array of real numbers. It is possible to specify the index of the sybsystem to get the mean occupation number by fixing `idx`.
 """
-function get_coherence(ψ::QuantumObject{<:AbstractArray,KetQuantumObject})
-    a = destroy(prod(ψ.dims))
-    α = expect(a, ψ)
-    D = exp(α * a' - conj(α) * a)
+function mean_occupation(ψ::QuantumObject{T,KetQuantumObject}; idx::Union{Int,Nothing} = nothing) where {T}
+    t = Tuple(reverse(ψ.dims))
+    mean_occ = zeros(length(ψ.dims))
+
+    for J in eachindex(ψ.data)
+        sub_indices = _ind2sub(t, J) .- 1
+        mean_occ .+= abs2(ψ[J]) .* sub_indices
+    end
+    reverse!(mean_occ)
+
+    return isnothing(idx) ? mean_occ : mean_occ[idx]
+end
+
+mean_occupation(ψ::QuantumObject{T,KetQuantumObject,1}) where {T} = mapreduce(k -> abs2(ψ[k]) * (k - 1), +, 1:ψ.dims[1])
+
+function mean_occupation(ρ::QuantumObject{T,OperatorQuantumObject}; idx::Union{Int,Nothing} = nothing) where {T}
+    t = Tuple(reverse(ρ.dims))
+    mean_occ = zeros(eltype(ρ.data), length(ρ.dims))
+
+    for J in eachindex(ρ.data[:, 1])
+        sub_indices = _ind2sub(t, J) .- 1
+        mean_occ .+= ρ[J, J] .* sub_indices
+    end
+    reverse!(mean_occ)
 
-    return α, D' * ψ
+    return isnothing(idx) ? real.(mean_occ) : real(mean_occ[idx])
 end
 
-function get_coherence(ρ::QuantumObject{<:AbstractArray,OperatorQuantumObject})
-    a = destroy(prod(ρ.dims))
-    α = expect(a, ρ)
-    D = exp(α * a' - conj(α) * a)
+mean_occupation(ρ::QuantumObject{T,OperatorQuantumObject,1}) where {T} =
+    mapreduce(k -> ρ[k, k] * (k - 1), +, 1:ρ.dims[1])
+
+function mean_occupation(v::QuantumObject{T,OperatorKetQuantumObject,1}) where {T}
+    d = v.dims[1]
+    return real(mapreduce(k -> v[(k-1)*d+k] * (k - 1), +, 1:d))
+end
+
+function mean_occupation(v::QuantumObject{T,OperatorKetQuantumObject}; idx::Union{Int,Nothing} = nothing) where {T}
+    dims = reverse(Tuple(v.dims))
+    mean_occ = zeros(eltype(v.data), length(v.dims))
+
+    offset = prod(dims)
+
+    for J in 1:offset
+        sub_indices = _ind2sub(dims, J) .- 1
+        mean_occ .+= v[(J-1)*offset+J] .* sub_indices
+    end
+    reverse!(mean_occ)
 
-    return α, D' * ρ * D
+    return isnothing(idx) ? real.(mean_occ) : real(mean_occ[idx])
 end
 
 @doc raw"""

test/core-test/quantum_objects.jl

@@ -450,17 +450,144 @@
         end
     end
 
-    @testset "get coherence" begin
-        ψ = coherent(30, 3)
-        α, δψ = get_coherence(ψ)
-        @test isapprox(abs(α), 3, atol = 1e-5) && abs2(δψ[1]) > 0.999
-        ρ = ket2dm(ψ)
-        α, δρ = get_coherence(ρ)
-        @test isapprox(abs(α), 3, atol = 1e-5) && abs2(δρ[1, 1]) > 0.999
+    @testset "get and remove coherence" begin
+        α1 = 3.0 + 1.5im
+        ψ1 = coherent(30, α1)
+        β1 = get_coherence(ψ1)
+        @test isapprox(α1, β1, atol = 1e-4)
+
+        ρ1 = ket2dm(ψ1)
+        β1 = get_coherence(ρ1)
+        @test isapprox(α1, β1, atol = 1e-4)
+
+        v1 = mat2vec(ρ1)
+        β1 = get_coherence(v1)
+        @test isapprox(α1, β1, atol = 1e-4)
+
+        α2 = 1.2 + 0.3im
+        ψ2 = coherent(15, α2)
+        β2 = get_coherence(ψ2)
+        @test isapprox(α2, β2, atol = 1e-4)
+
+        ρ2 = ket2dm(ψ2)
+        β2 = get_coherence(ρ2)
+        @test isapprox(α2, β2, atol = 1e-4)
+
+        v2 = mat2vec(ρ2)
+        β2 = get_coherence(v2)
+        @test isapprox(α2, β2, atol = 1e-4)
+
+        ψ = ψ1 ⊗ ψ2
+        βs = get_coherence(ψ)
+        @test length(βs) == 2
+        @test isapprox(βs[1], α1, atol = 1e-4)
+        @test isapprox(βs[2], α2, atol = 1e-4)
+
+        ρ = ρ1 ⊗ ρ2
+        βs = get_coherence(ρ)
+        @test length(βs) == 2
+        @test isapprox(βs[1], α1, atol = 1e-4)
+        @test isapprox(βs[2], α2, atol = 1e-4)
+
+        v = mat2vec(ρ)
+        βs = get_coherence(v)
+        @test length(βs) == 2
+        @test isapprox(βs[1], α1, atol = 1e-4)
+        @test isapprox(βs[2], α2, atol = 1e-4)
+
+        αs = [α1, α2]
+        for idx in 1:2
+            β = get_coherence(ψ, idx = idx)
+            @test isapprox(β, αs[idx], atol = 1e-4)
+            β = get_coherence(ρ, idx = idx)
+            @test isapprox(β, αs[idx], atol = 1e-4)
+            β = get_coherence(v, idx = idx)
+            @test isapprox(β, αs[idx], atol = 1e-4)
+        end
 
         @testset "Type Inference (get_coherence)" begin
+            @inferred get_coherence(ψ1)
+            @inferred get_coherence(ρ1)
+            @inferred get_coherence(v1)
             @inferred get_coherence(ψ)
             @inferred get_coherence(ρ)
+            @inferred get_coherence(v)
+        end
+    end
+
+    @testset "remove coherence" begin
+        α1 = 3.0 + 1.5im
+        ψ1 = coherent(30, α1)
+        ψ2 = coherent(15, 1.2 + 0.3im)
+
+        δψ1 = remove_coherence(ψ1)
+        @test abs2(δψ1[1]) > 0.999
+
+        ρ1 = ket2dm(ψ1)
+        δρ1 = remove_coherence(ρ1)
+        @test abs2(δρ1[1, 1]) > 0.999
+
+        ψ = ψ1 ⊗ ψ2
+        δψ = remove_coherence(ψ)
+        @test abs2(δψ[1]) > 0.999
+
+        ρ = ket2dm(ψ)
+        δρ = remove_coherence(ρ)
+        @test abs2(δρ[1, 1]) > 0.999
+
+        @testset "Type Inference (remove_coherence)" begin
+            @inferred remove_coherence(ψ1)
+            @inferred remove_coherence(ρ1)
+            @inferred remove_coherence(ψ)
+            @inferred remove_coherence(ρ)
+        end
+    end
+
+    @testset "mean occupation" begin
+        N1 = 9.0
+        ψ1 = coherent(50, 3.0)
+        ρ1 = ket2dm(ψ1)
+        v1 = mat2vec(ρ1)
+
+        @test mean_occupation(ψ1) ≈ N1
+        @test mean_occupation(ρ1) ≈ N1
+        @test mean_occupation(v1) ≈ N1
+
+        N2 = 4.0
+        Nc = N1 * N2
+        Ns = [N1, N2]
+        ψ2 = coherent(30, 2.0)
+        ρ2 = ket2dm(ψ2)
+        v2 = mat2vec(ρ2)
+
+        ψc = ψ1 ⊗ ψ2
+        ρc = ket2dm(ψc)
+        vc = mat2vec(ρc)
+
+        @test mean_occupation(ψ2) ≈ N2
+        @test mean_occupation(ρ2) ≈ N2
+        @test mean_occupation(v2) ≈ N2
+
+        @test mean_occupation(ψc) ≈ Ns
+        @test mean_occupation(ρc) ≈ Ns
+        @test mean_occupation(vc) ≈ Ns
+        @test prod(mean_occupation(ψc)) ≈ Nc
+        @test prod(mean_occupation(ρc)) ≈ Nc
+        @test prod(mean_occupation(vc)) ≈ Nc
+        @test mean_occupation(ψc, idx = 1) ≈ N1
+        @test mean_occupation(ρc, idx = 1) ≈ N1
+        @test mean_occupation(ψc, idx = 2) ≈ N2
+        @test mean_occupation(ρc, idx = 2) ≈ N2
+        @test mean_occupation(vc, idx = 1) ≈ N1
+        @test mean_occupation(vc, idx = 2) ≈ N2
+
+        @testset "Type Inference (mean_occupation)" begin
+            @inferred mean_occupation(ψ1)
+            @inferred mean_occupation(ρ1)
+            @inferred mean_occupation(v1)
+            @inferred mean_occupation(ψc)
+            @inferred mean_occupation(ρc)
+            @inferred mean_occupation(vc)
         end
     end