extended methods for expect and variance

Author: ytdHuang

src/qobj/functions.jl

@@ -17,7 +17,7 @@ ket2dm(ψ::QuantumObject{<:AbstractArray{T},KetQuantumObject}) where {T} = ψ *
 ket2dm(ρ::QuantumObject{<:AbstractArray{T},OperatorQuantumObject}) where {T} = ρ
 
 @doc raw"""
-    expect(O::QuantumObject, ψ::QuantumObject)
+    expect(O::QuantumObject, ψ::Union{QuantumObject,Vector{QuantumObject}})
 
 Expectation value of the [`Operator`](@ref) `O` with the state `ψ`. The state can be a [`Ket`](@ref), [`Bra`](@ref) or [`Operator`](@ref).
 
@@ -27,6 +27,8 @@ If `ψ` is a density matrix ([`Operator`](@ref)), the function calculates ``\tex
 
 The function returns a real number if `O` is of `Hermitian` type or `Symmetric` type, and returns a complex number otherwise. You can make an operator `O` hermitian by using `Hermitian(O)`.
 
+Note that `ψ` can also be given as a list of [`QuantumObject`](@ref), it returns a list of expectation values.
+
 # Examples
 
 ```
@@ -77,20 +79,33 @@ function expect(
 ) where {TF<:Number,TR<:Real,T2}
     return real(tr(O * ρ))
 end
+function expect(
+    O::QuantumObject{<:AbstractArray{T1},OperatorQuantumObject},
+    ρ::Vector{<:QuantumObject},
+) where {T1}
+    _expect = _ρ -> expect(O, _ρ)
+    return _expect.(ρ)
+end
 
 @doc raw"""
-    variance(O::QuantumObject, ψ::QuantumObject)
+    variance(O::QuantumObject, ψ::Union{QuantumObject,Vector{QuantumObject}})
 
 Variance of the [`Operator`](@ref) `O`: ``\langle\hat{O}^2\rangle - \langle\hat{O}\rangle^2``,
 
 where ``\langle\hat{O}\rangle`` is the expectation value of `O` with the state `ψ` (see also [`expect`](@ref)), and the state `ψ` can be a [`Ket`](@ref), [`Bra`](@ref) or [`Operator`](@ref).
 
 The function returns a real number if `O` is hermitian, and returns a complex number otherwise.
+
+Note that `ψ` can also be given as a list of [`QuantumObject`](@ref), it returns a list of expectation values.
 """
 variance(
     O::QuantumObject{<:AbstractArray{T1},OperatorQuantumObject},
     ψ::QuantumObject{<:AbstractArray{T2}},
 ) where {T1,T2} = expect(O^2, ψ) - expect(O, ψ)^2
+variance(
+    O::QuantumObject{<:AbstractArray{T1},OperatorQuantumObject},
+    ψ::Vector{<:QuantumObject},
+) where {T1} = expect(O^2, ψ) .- expect(O, ψ).^2
 
 @doc raw"""
     sparse_to_dense(A::QuantumObject)

src/qobj/superoperators.jl

@@ -25,6 +25,7 @@ Since the density matrix is vectorized in [`OperatorKet`](@ref) form: ``|\hat{\r
 ```math
 \mathcal{O} \left(\hat{A}\right) \left[ \hat{\rho} \right] = \hat{\mathbb{1}} \otimes \hat{A} ~ |\hat{\rho}\rangle\rangle
 ```
+(see the section in documentation: [Superoperators and Vectorized Operators](@ref doc:Superoperators-and-Vectorized-Operators) for more details)
 
 The optional argument `Id_cache` can be used to pass a precomputed identity matrix. This can be useful when
 the same function is applied multiple times with a known Hilbert space dimension.
@@ -42,6 +43,7 @@ Since the density matrix is vectorized in [`OperatorKet`](@ref) form: ``|\hat{\r
 ```math
 \mathcal{O} \left(\hat{B}\right) \left[ \hat{\rho} \right] = \hat{B}^T \otimes \hat{\mathbb{1}} ~ |\hat{\rho}\rangle\rangle
 ```
+(see the section in documentation: [Superoperators and Vectorized Operators](@ref doc:Superoperators-and-Vectorized-Operators) for more details)
 
 The optional argument `Id_cache` can be used to pass a precomputed identity matrix. This can be useful when
 the same function is applied multiple times with a known Hilbert space dimension.
@@ -59,6 +61,7 @@ Since the density matrix is vectorized in [`OperatorKet`](@ref) form: ``|\hat{\r
 ```math
 \mathcal{O} \left(\hat{A}, \hat{B}\right) \left[ \hat{\rho} \right] = \hat{B}^T \otimes \hat{A} ~ |\hat{\rho}\rangle\rangle = \textrm{spre}(A) * \textrm{spost}(B) ~ |\hat{\rho}\rangle\rangle
 ```
+(see the section in documentation: [Superoperators and Vectorized Operators](@ref doc:Superoperators-and-Vectorized-Operators) for more details)
 
 See also [`spre`](@ref) and [`spost`](@ref).
 """

test/quantum_objects.jl

@@ -384,13 +384,20 @@
         @test expect(a, ρ) ≈ tr(a * ρ)
         @test variance(a, ρ) ≈ tr(a^2 * ρ) - tr(a * ρ)^2
 
+        # when input is a vector of states
+        xlist = [1.0, 1.0im, -1.0, -1.0im]
+        ψlist = [normalize!(basis(N, 4) + x * basis(N, 3)) for x in xlist]
+        @test all(expect(a', ψlist) .≈ xlist)
+
         @testset "Type Inference (expect)" begin
             @inferred expect(a, ψ)
             @inferred expect(a, ψ')
             @inferred variance(a, ψ)
             @inferred variance(a, ψ')
             @inferred expect(a, ρ)
             @inferred variance(a, ρ)
+            @inferred expect(a, ψlist)
+            @inferred variance(a, ψlist)
         end
     end