return list of a and n if multiple subsystems and tests

Author: Adria Labay

src/qobj/arithmetic_and_attributes.jl

@@ -663,64 +663,96 @@ 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}`` or a density matrix ``\hat{\rho}``.
-"""
-function get_coherence(ψ::QuantumObject{<:AbstractArray,KetQuantumObject})
-    if length(ψ.dims) == 1
-        return mapreduce(n -> sqrt(n - 1) * ψ.data[n] * conj(ψ.data[n-1]), +, 2:ψ.dims[1])
-    else
-        off = sum(cumprod(reverse(ψ.dims[2:end]))) + 1
-        t = Tuple(reverse(ψ.dims))
+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``.
 
-        x = 0.0im
-        for J in off+2:length(ψ.data)
-            x += ψ[J] * conj(ψ[J-off]) * prod(sqrt.(_ind2sub(t, J) .- 1))
+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
-        return x
     end
+
+    reverse!(expectation_values)
+
+    return isnothing(idx) ? expectation_values : expectation_values[idx]
 end
 
-function get_coherence(ρ::QuantumObject{<:AbstractArray,OperatorQuantumObject})
-    if length(ρ.dims) == 1
-        return mapreduce(n -> sqrt(n - 1) * ρ.data[n, n-1], +, 2:ρ.dims[1])
-    else
-        off = sum(cumprod(reverse(ρ.dims[2:end]))) + 1
-        t = Tuple(reverse(ρ.dims))
+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)
 
-        x = 0.0im
-        for J in off+2:length(ρ.data[1, :])
-            x += ρ[J, J-off] * prod(sqrt.(_ind2sub(t, J) .- 1))
+    @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
-        return x
     end
+
+    reverse!(expectation_values)
+
+    return isnothing(idx) ? expectation_values : expectation_values[idx]
 end
 
-function get_coherence(v::QuantumObject{T,OperatorKetQuantumObject}) where {T}
-    if length(v.dims) > 1
-        throw(ArgumentError("Mean photon number not implemented for composite OPeratorKetQuantumObject"))
+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
 
-    d = v.dims[1]
-    return mapreduce(n -> sqrt(n - 1) * v.data[(n-1)*d+n-1], +, 2:d)
+    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{<:AbstractArray,KetQuantumObject,1})
-    α = get_coherence(ψ)
-    D = displace(ψ.dims[1], α)
+function remove_coherence(ψ::QuantumObject{T,KetQuantumObject}) where {T}
+    αs = get_coherence(ψ)
+    D = mapreduce(k -> displace(ψ.dims[k], αs[k]), ⊗, eachindex(ψ.dims))
 
     return D' * ψ
 end
 
-function remove_coherence(ρ::QuantumObject{<:AbstractArray,OperatorQuantumObject,1})
-    α = get_coherence(ρ)
-    D = displace(ρ.dims[1], α)
+function remove_coherence(ρ::QuantumObject{T,OperatorQuantumObject}) where {T}
+    αs = get_coherence(ρ)
+    D = mapreduce(k -> displace(ρ.dims[k], αs[k]), ⊗, eachindex(ρ.dims))
 
     return D' * ρ * D
 end
@@ -730,7 +762,7 @@ end
 
 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}}``.
 
-It returns the expectation value of the number operator.
+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 mean_occupation(ψ::QuantumObject{T,KetQuantumObject}; idx::Union{Int,Nothing} = nothing) where {T}
     t = Tuple(reverse(ψ.dims))
@@ -751,7 +783,6 @@ function mean_occupation(ρ::QuantumObject{T,OperatorQuantumObject}; idx::Union{
     t = Tuple(reverse(ρ.dims))
     mean_occ = zeros(eltype(ρ.data), length(ρ.dims))
 
-    x = 0.0im
     for J in eachindex(ρ.data[:, 1])
         sub_indices = _ind2sub(t, J) .- 1
         mean_occ .+= ρ[J, J] .* sub_indices
@@ -764,15 +795,26 @@ end
 mean_occupation(ρ::QuantumObject{T,OperatorQuantumObject,1}) where {T} =
     mapreduce(k -> ρ[k, k] * (k - 1), +, 1:ρ.dims[1])
 
-function mean_occupation(v::QuantumObject{T,OperatorKetQuantumObject}) where {T}
-    if length(v.dims) > 1
-        throw(ArgumentError("Mean photon number not implemented for composite OPeratorKetQuantumObject"))
-    end
-
+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 isnothing(idx) ? real.(mean_occ) : real(mean_occ[idx])
+end
+
 @doc raw"""
     permute(A::QuantumObject, order::Union{AbstractVector{Int},Tuple})
 

test/core-test/quantum_objects.jl

@@ -443,20 +443,93 @@
     end
 
     @testset "get and remove coherence" begin
-        ψ = coherent(30, 3)
-        α = get_coherence(ψ)
-        @test isapprox(abs(α), 3, atol = 1e-5)
+        α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(ψ)
-        α = get_coherence(ρ)
-        @test isapprox(abs(α), 3, atol = 1e-5)
         δρ = remove_coherence(ρ)
         @test abs2(δρ[1, 1]) > 0.999
 
-        @testset "Type Inference (get_coherence)" begin
-            @inferred get_coherence(ψ)
-            @inferred get_coherence(ρ)
+        @testset "Type Inference (remove_coherence)" begin
+            @inferred remove_coherence(ψ1)
+            @inferred remove_coherence(ρ1)
             @inferred remove_coherence(ψ)
             @inferred remove_coherence(ρ)
         end
@@ -489,17 +562,24 @@
 
         @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