From f1f9bf2907404daf937b5deb861b7e4aee3ca4fb Mon Sep 17 00:00:00 2001
From: Adria Labay <adria.labay@gmail.com>
Date: Sun, 8 Dec 2024 18:47:21 +0100
Subject: [PATCH 1/5] implement mean occupation and get coherence functions

---
 src/qobj/arithmetic_and_attributes.jl | 121 +++++++++++++++++++++++---
 test/core-test/quantum_objects.jl     |  47 ++++++++--
 2 files changed, 152 insertions(+), 16 deletions(-)

diff --git a/src/qobj/arithmetic_and_attributes.jl b/src/qobj/arithmetic_and_attributes.jl
index 52e6859e1..1c161e98c 100644
--- a/src/qobj/arithmetic_and_attributes.jl
+++ b/src/qobj/arithmetic_and_attributes.jl
@@ -697,24 +697,123 @@ get_data(A::AbstractQuantumObject) = A.data
 @doc raw"""
     get_coherence(ψ::QuantumObject)
 
-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}``.
-
-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|``.
+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})
-    a = destroy(prod(ψ.dims))
-    α = expect(a, ψ)
-    D = exp(α * a' - conj(α) * a)
+    if length(ψ.dims) == 1
+        return mapreduce(n -> sqrt(n - 1) * ψ.data[n] * conj(ψ.data[n-1]), +, 2:ψ.dims[1])
+    else
+        R = CartesianIndices((ψ.dims...,))
+        off = circshift(ψ.dims, 1)
+        off[end] = 1
+
+        x = sum(R) do j
+            j_tuple = Tuple(j) .- 1
+            if 0 in j_tuple
+                return 0
+            end
+
+            J = dot(j_tuple, off) + 1
+            J2 = dot(j_tuple .- 1, off) + 1
+            return prod(sqrt.(j_tuple)) * ψ[J] * conj(ψ[J2])
+        end
 
-    return α, D' * ψ
+        return x
+    end
 end
 
 function get_coherence(ρ::QuantumObject{<:AbstractArray,OperatorQuantumObject})
-    a = destroy(prod(ρ.dims))
-    α = expect(a, ρ)
-    D = exp(α * a' - conj(α) * a)
+    if length(ρ.dims) == 1
+        return mapreduce(n -> sqrt(n - 1) * ρ.data[n, n-1], +, 2:ρ.dims[1])
+    else
+        R = CartesianIndices((ρ.dims...,))
+        off = circshift(ρ.dims, 1)
+        off[end] = 1
+
+        x = sum(R) do j
+            j_tuple = Tuple(j) .- 1
+            if 0 in j_tuple
+                return 0
+            end
+
+            J = dot(j_tuple, off) + 1
+            J2 = dot(j_tuple .- 1, off) + 1
+            return prod(sqrt.(j_tuple)) * ρ[J, J2]
+        end
+
+        return x
+    end
+end
+
+@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], α)
+
+    return D' * ψ
+end
+
+function remove_coherence(ρ::QuantumObject{<:AbstractArray,OperatorQuantumObject,1})
+    α = get_coherence(ρ)
+    D = displace(ρ.dims[1], α)
+
+    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}}``.
+
+It returns the expectation value of the number operator.
+"""
+function mean_occupation(ψ::QuantumObject{T,KetQuantumObject}) where {T}
+    if length(ψ.dims) == 1
+        return mapreduce(k -> abs2(ψ[k]) * (k - 1), +, 1:ρ.dims[1])
+    else
+        R = CartesianIndices((ψ.dims...,))
+        off = circshift(ψ.dims, 1)
+        off[end] = 1
+
+        x = sum(R) do j
+            j_tuple = Tuple(j) .- 1
+            J = dot(j_tuple, off) + 1
+            return abs2(ψ[J]) * prod(j_tuple)
+        end
+
+        return x
+    end
+end
+
+function mean_occupation(ρ::QuantumObject{T,OperatorQuantumObject}) where {T}
+    if length(ρ.dims) == 1
+        return real(mapreduce(k -> ρ[k, k] * (k - 1), +, 1:ρ.dims[1]))
+    else
+        R = CartesianIndices((ρ.dims...,))
+        off = circshift(ψ.dims, 1)
+        off[end] = 1
+
+        x = sum(R) do j
+            j_tuple = Tuple(j) .- 1
+            J = dot(j_tuple, off) + 1
+            return ρ[J, J] * prod(j_tuple)
+        end
+
+        return real(x)
+    end
+end
+
+function mean_occupation(ρ::QuantumObject{T,OperatorKetQuantumObject}) where {T}
+    if length(ρ.dims) > 1
+        throw(ArgumentError("Mean photon number not implemented for composite OPeratorKetQuantumObject"))
+    end
 
-    return α, D' * ρ * D
+    d = ρ.dims[1]
+    return real(mapreduce(k -> ρ[(k-1)*r+k] * (k - 1), +, 1:d))
 end
 
 @doc raw"""
diff --git a/test/core-test/quantum_objects.jl b/test/core-test/quantum_objects.jl
index 6d26ac4c2..3401e78d4 100644
--- a/test/core-test/quantum_objects.jl
+++ b/test/core-test/quantum_objects.jl
@@ -450,17 +450,54 @@
         end
     end
 
-    @testset "get coherence" begin
+    @testset "get and remove coherence" begin
         ψ = coherent(30, 3)
-        α, δψ = get_coherence(ψ)
-        @test isapprox(abs(α), 3, atol = 1e-5) && abs2(δψ[1]) > 0.999
+        α = get_coherence(ψ)
+        @test isapprox(abs(α), 3, atol = 1e-5)
+        δψ = remove_coherence(ψ)
+        @test abs2(δψ[1]) > 0.999
         ρ = ket2dm(ψ)
-        α, δρ = get_coherence(ρ)
-        @test isapprox(abs(α), 3, atol = 1e-5) && abs2(δρ[1, 1]) > 0.999
+        α = 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(ρ)
+            @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(ψ) ≈ N1
+        @test mean_occupation(ρ) ≈ N1
+        @test mean_occupation(v) ≈ N1
+
+        N2 = 4.0
+        Nc = N1 * N2
+        ψ2 = coherent(30, 2.0)
+        ψ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) ≈ Nc
+        @test mean_occupation(ρc) ≈ Nc
+
+        @testset "Type Inference (mean_occupation)" begin
+            @inferred mean_occupation(ψ)
+            @inferred mean_occupation(ρ)
+            @inferred mean_occupation(v)
         end
     end
 

From 9910a343d81b3604b038f07f665b7bf0fc4d94cc Mon Sep 17 00:00:00 2001
From: Adria Labay <adria.labay@gmail.com>
Date: Sun, 8 Dec 2024 18:55:29 +0100
Subject: [PATCH 2/5] export functions

---
 src/qobj/arithmetic_and_attributes.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/qobj/arithmetic_and_attributes.jl b/src/qobj/arithmetic_and_attributes.jl
index 1c161e98c..354722194 100644
--- a/src/qobj/arithmetic_and_attributes.jl
+++ b/src/qobj/arithmetic_and_attributes.jl
@@ -6,7 +6,7 @@ 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
 
 #    Broadcasting
 Base.broadcastable(x::QuantumObject) = x.data

From ae029d413b00aa160aadbff4507ca61826a08572 Mon Sep 17 00:00:00 2001
From: Adria Labay <adria.labay@gmail.com>
Date: Mon, 9 Dec 2024 09:46:03 +0100
Subject: [PATCH 3/5] skip use of cartesianindex to reduce allocations

---
 src/qobj/arithmetic_and_attributes.jl | 27 +++++++++++----------------
 1 file changed, 11 insertions(+), 16 deletions(-)

diff --git a/src/qobj/arithmetic_and_attributes.jl b/src/qobj/arithmetic_and_attributes.jl
index 354722194..59a42dd89 100644
--- a/src/qobj/arithmetic_and_attributes.jl
+++ b/src/qobj/arithmetic_and_attributes.jl
@@ -8,6 +8,8 @@ export proj, ptrace, purity, permute
 export tidyup, tidyup!
 export get_data, get_coherence, remove_coherence, mean_occupation
 
+import Base: _ind2sub
+
 #    Broadcasting
 Base.broadcastable(x::QuantumObject) = x.data
 for op in (:(+), :(-), :(*), :(/), :(^))
@@ -775,17 +777,13 @@ function mean_occupation(ψ::QuantumObject{T,KetQuantumObject}) where {T}
     if length(ψ.dims) == 1
         return mapreduce(k -> abs2(ψ[k]) * (k - 1), +, 1:ρ.dims[1])
     else
-        R = CartesianIndices((ψ.dims...,))
-        off = circshift(ψ.dims, 1)
-        off[end] = 1
+        t = Tuple(ψ.dims)
 
-        x = sum(R) do j
-            j_tuple = Tuple(j) .- 1
-            J = dot(j_tuple, off) + 1
-            return abs2(ψ[J]) * prod(j_tuple)
+        x = 0.0
+        for J in eachindex(ψ.data)
+            x += abs2(ψ[J]) * prod(Base._ind2sub(t, J) .- 1)
         end
-
-        return x
+        return real(x)
     end
 end
 
@@ -793,14 +791,11 @@ function mean_occupation(ρ::QuantumObject{T,OperatorQuantumObject}) where {T}
     if length(ρ.dims) == 1
         return real(mapreduce(k -> ρ[k, k] * (k - 1), +, 1:ρ.dims[1]))
     else
-        R = CartesianIndices((ρ.dims...,))
-        off = circshift(ψ.dims, 1)
-        off[end] = 1
+        t = Tuple(ρ.dims)
 
-        x = sum(R) do j
-            j_tuple = Tuple(j) .- 1
-            J = dot(j_tuple, off) + 1
-            return ρ[J, J] * prod(j_tuple)
+        x = 0.0im
+        for J in eachindex(ρ.data[:, 1])
+            x += ρ[J, J] * prod(Base._ind2sub(t, J) .- 1)
         end
 
         return real(x)

From 0d91c325fd19f2dc76c8bd1e5533aaff478fd71e Mon Sep 17 00:00:00 2001
From: Adria Labay <adria.labay@gmail.com>
Date: Sun, 15 Dec 2024 16:48:31 +0100
Subject: [PATCH 4/5] return list of mean occupation per dimension

---
 src/qobj/arithmetic_and_attributes.jl | 100 +++++++++++++-------------
 test/core-test/quantum_objects.jl     |  26 ++++---
 2 files changed, 66 insertions(+), 60 deletions(-)

diff --git a/src/qobj/arithmetic_and_attributes.jl b/src/qobj/arithmetic_and_attributes.jl
index 59a42dd89..fc4ebc68b 100644
--- a/src/qobj/arithmetic_and_attributes.jl
+++ b/src/qobj/arithmetic_and_attributes.jl
@@ -705,21 +705,13 @@ 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
-        R = CartesianIndices((ψ.dims...,))
-        off = circshift(ψ.dims, 1)
-        off[end] = 1
-
-        x = sum(R) do j
-            j_tuple = Tuple(j) .- 1
-            if 0 in j_tuple
-                return 0
-            end
-
-            J = dot(j_tuple, off) + 1
-            J2 = dot(j_tuple .- 1, off) + 1
-            return prod(sqrt.(j_tuple)) * ψ[J] * conj(ψ[J2])
-        end
+        off = sum(cumprod(reverse(ψ.dims[2:end]))) + 1
+        t = Tuple(reverse(ψ.dims))
 
+        x = 0.0im
+        for J in off+2:length(ψ.data)
+            x += ψ[J] * conj(ψ[J-off]) * prod(sqrt.(_ind2sub(t, J) .- 1))
+        end
         return x
     end
 end
@@ -728,25 +720,26 @@ function get_coherence(ρ::QuantumObject{<:AbstractArray,OperatorQuantumObject})
     if length(ρ.dims) == 1
         return mapreduce(n -> sqrt(n - 1) * ρ.data[n, n-1], +, 2:ρ.dims[1])
     else
-        R = CartesianIndices((ρ.dims...,))
-        off = circshift(ρ.dims, 1)
-        off[end] = 1
-
-        x = sum(R) do j
-            j_tuple = Tuple(j) .- 1
-            if 0 in j_tuple
-                return 0
-            end
-
-            J = dot(j_tuple, off) + 1
-            J2 = dot(j_tuple .- 1, off) + 1
-            return prod(sqrt.(j_tuple)) * ρ[J, J2]
-        end
+        off = sum(cumprod(reverse(ρ.dims[2:end]))) + 1
+        t = Tuple(reverse(ρ.dims))
 
+        x = 0.0im
+        for J in off+2:length(ρ.data[1, :])
+            x += ρ[J, J-off] * prod(sqrt.(_ind2sub(t, J) .- 1))
+        end
         return x
     end
 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"))
+    end
+
+    d = v.dims[1]
+    return mapreduce(n -> sqrt(n - 1) * v.data[(n-1)*d+n-1], +, 2:d)
+end
+
 @doc raw"""
     remove_coherence(ψ::QuantumObject)
 
@@ -773,42 +766,45 @@ Get the mean occupation number ``n`` by measuring the expectation value of the n
 
 It returns the expectation value of the number operator.
 """
-function mean_occupation(ψ::QuantumObject{T,KetQuantumObject}) where {T}
-    if length(ψ.dims) == 1
-        return mapreduce(k -> abs2(ψ[k]) * (k - 1), +, 1:ρ.dims[1])
-    else
-        t = Tuple(ψ.dims)
+function mean_occupation(ψ::QuantumObject{T,KetQuantumObject}; idx::Union{Int,Nothing} = nothing) where {T}
+    t = Tuple(reverse(ψ.dims))
+    mean_occ = zeros(length(ψ.dims))
 
-        x = 0.0
-        for J in eachindex(ψ.data)
-            x += abs2(ψ[J]) * prod(Base._ind2sub(t, J) .- 1)
-        end
-        return real(x)
+    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
 
-function mean_occupation(ρ::QuantumObject{T,OperatorQuantumObject}) where {T}
-    if length(ρ.dims) == 1
-        return real(mapreduce(k -> ρ[k, k] * (k - 1), +, 1:ρ.dims[1]))
-    else
-        t = Tuple(ρ.dims)
+mean_occupation(ψ::QuantumObject{T,KetQuantumObject,1}) where {T} = mapreduce(k -> abs2(ψ[k]) * (k - 1), +, 1:ψ.dims[1])
 
-        x = 0.0im
-        for J in eachindex(ρ.data[:, 1])
-            x += ρ[J, J] * prod(Base._ind2sub(t, J) .- 1)
-        end
+function mean_occupation(ρ::QuantumObject{T,OperatorQuantumObject}; idx::Union{Int,Nothing} = nothing) where {T}
+    t = Tuple(reverse(ρ.dims))
+    mean_occ = zeros(eltype(ρ.data), length(ρ.dims))
 
-        return real(x)
+    x = 0.0im
+    for J in eachindex(ρ.data[:, 1])
+        sub_indices = _ind2sub(t, J) .- 1
+        mean_occ .+= ρ[J, J] .* sub_indices
     end
+    reverse!(mean_occ)
+
+    return isnothing(idx) ? real.(mean_occ) : real(mean_occ[idx])
 end
 
-function mean_occupation(ρ::QuantumObject{T,OperatorKetQuantumObject}) where {T}
-    if length(ρ.dims) > 1
+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
 
-    d = ρ.dims[1]
-    return real(mapreduce(k -> ρ[(k-1)*r+k] * (k - 1), +, 1:d))
+    d = v.dims[1]
+    return real(mapreduce(k -> v[(k-1)*d+k] * (k - 1), +, 1:d))
 end
 
 @doc raw"""
diff --git a/test/core-test/quantum_objects.jl b/test/core-test/quantum_objects.jl
index 3401e78d4..57539c9b7 100644
--- a/test/core-test/quantum_objects.jl
+++ b/test/core-test/quantum_objects.jl
@@ -476,13 +476,17 @@
         ρ1 = ket2dm(ψ1)
         v1 = mat2vec(ρ1)
 
-        @test mean_occupation(ψ) ≈ N1
-        @test mean_occupation(ρ) ≈ N1
-        @test mean_occupation(v) ≈ N1
+        @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)
@@ -491,13 +495,19 @@
         @test mean_occupation(ρ2) ≈ N2
         @test mean_occupation(v2) ≈ N2
 
-        @test mean_occupation(ψc) ≈ Nc
-        @test mean_occupation(ρc) ≈ Nc
+        @test mean_occupation(ψc) ≈ Ns
+        @test mean_occupation(ρc) ≈ Ns
+        @test prod(mean_occupation(ψc)) ≈ Nc
+        @test prod(mean_occupation(ρc)) ≈ 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
 
         @testset "Type Inference (mean_occupation)" begin
-            @inferred mean_occupation(ψ)
-            @inferred mean_occupation(ρ)
-            @inferred mean_occupation(v)
+            @inferred mean_occupation(ψ1)
+            @inferred mean_occupation(ρ1)
+            @inferred mean_occupation(v1)
         end
     end
 

From 29cae761943e2d38ba00f368497e75210f9f2f47 Mon Sep 17 00:00:00 2001
From: Adria Labay <adria.labay@gmail.com>
Date: Sun, 22 Dec 2024 19:44:34 +0100
Subject: [PATCH 5/5] return list of a and n if multiple subsystems and tests

---
 src/qobj/arithmetic_and_attributes.jl | 124 +++++++++++++++++---------
 test/core-test/quantum_objects.jl     |  96 ++++++++++++++++++--
 2 files changed, 171 insertions(+), 49 deletions(-)

diff --git a/src/qobj/arithmetic_and_attributes.jl b/src/qobj/arithmetic_and_attributes.jl
index fc4ebc68b..1ce10f0ce 100644
--- a/src/qobj/arithmetic_and_attributes.jl
+++ b/src/qobj/arithmetic_and_attributes.jl
@@ -697,64 +697,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
@@ -764,7 +796,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))
@@ -785,7 +817,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
@@ -798,15 +829,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})
 
diff --git a/test/core-test/quantum_objects.jl b/test/core-test/quantum_objects.jl
index 57539c9b7..bd558a9ca 100644
--- a/test/core-test/quantum_objects.jl
+++ b/test/core-test/quantum_objects.jl
@@ -451,20 +451,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
@@ -497,17 +570,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