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] 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 52e6859e..1c161e98 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 6d26ac4c..3401e78d 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