Fix type instabilities for qobj methods

Author: albertomercurio

src/qobj/arithmetic_and_attributes.jl

@@ -239,8 +239,8 @@ julia> tr(a' * a)
 """
 LinearAlgebra.tr(
     A::QuantumObject{<:AbstractArray{T},OpType},
-) where {T,OpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject}} =
-    ishermitian(A) ? real(tr(A.data)) : tr(A.data)
+) where {T,OpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject}} = tr(A.data)
+LinearAlgebra.tr(A::QuantumObject{<:Union{<:Hermitian{TF}, Symmetric{TR}},OpType}) where {TF<:BlasFloat,TR<:Real,OpType<:OperatorQuantumObject} = real(tr(A.data))
 
 @doc raw"""
     svdvals(A::QuantumObject)
@@ -673,6 +673,9 @@ julia> ψ_123 = tensor(ψ1, ψ2, ψ3)
 julia> permute(ψ_123, [2, 1, 3]) ≈ tensor(ψ2, ψ1, ψ3)
 true
 ```
+
+!!! warning "Beware of type-stability!"
+    It is highly recommended to use `permute(A, order)` with `order` as `Tuple` or `SVector` to keep type stability. See the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 function permute(
     A::QuantumObject{<:AbstractArray{T},ObjType},
@@ -691,7 +694,7 @@ function permute(
     dims, perm = _dims_and_perm(A.type, A.dims, order_svector, length(order_svector))
 
     return QuantumObject(
-        reshape(PermutedDimsArray(reshape(A.data, dims...), Tuple(perm)), size(A)),
+        reshape(permutedims(reshape(A.data, dims...), Tuple(perm)), size(A)),
         A.type,
         A.dims[order_svector],
     )

src/qobj/functions.jl

@@ -25,7 +25,7 @@ If `ψ` is a [`Ket`](@ref) or [`Bra`](@ref), the function calculates ``\langle\p
 
 If `ψ` is a density matrix ([`Operator`](@ref)), the function calculates ``\textrm{Tr} \left[ \hat{O} \hat{\psi} \right]``
 
-The function returns a real number if `O` is hermitian, and returns a complex number otherwise.
+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)`.
 
 # Examples
 
@@ -34,31 +34,48 @@ julia> ψ = 1 / √2 * (fock(10,2) + fock(10,4));
 
 julia> a = destroy(10);
 
-julia> expect(a' * a, ψ) ≈ 3
-true
+julia> expect(a' * a, ψ) |> round
+3.0 + 0.0im
+
+julia> expect(Hermitian(a' * a), ψ) |> round
+3.0
 ```
 """
 function expect(
     O::QuantumObject{<:AbstractArray{T1},OperatorQuantumObject},
     ψ::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
 ) where {T1,T2}
-    ψd = ψ.data
-    Od = O.data
-    return ishermitian(O) ? real(dot(ψd, Od, ψd)) : dot(ψd, Od, ψd)
+    return dot(ψ.data, O.data, ψ.data)
 end
 function expect(
     O::QuantumObject{<:AbstractArray{T1},OperatorQuantumObject},
     ψ::QuantumObject{<:AbstractArray{T2},BraQuantumObject},
 ) where {T1,T2}
-    ψd = ψ.data'
-    Od = O.data
-    return ishermitian(O) ? real(dot(ψd, Od, ψd)) : dot(ψd, Od, ψd)
+    return expect(O, ψ')
 end
 function expect(
     O::QuantumObject{<:AbstractArray{T1},OperatorQuantumObject},
     ρ::QuantumObject{<:AbstractArray{T2},OperatorQuantumObject},
 ) where {T1,T2}
-    return ishermitian(O) ? real(tr(O * ρ)) : tr(O * ρ)
+    return tr(O * ρ)
+end
+function expect(
+    O::QuantumObject{<:Union{<:Hermitian{TF},<:Symmetric{TR}},OperatorQuantumObject},
+    ψ::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
+) where {TF<:Number,TR<:Real,T2}
+    return real(dot(ψ.data, O.data, ψ.data))
+end
+function expect(
+    O::QuantumObject{<:Union{<:Hermitian{TF},<:Symmetric{TR}},OperatorQuantumObject},
+    ψ::QuantumObject{<:AbstractArray{T2},BraQuantumObject},
+) where {TF<:Number,TR<:Real,T2}
+    return real(expect(O, ψ'))
+end
+function expect(
+    O::QuantumObject{<:Union{<:Hermitian{TF},<:Symmetric{TR}},OperatorQuantumObject},
+    ρ::QuantumObject{<:AbstractArray{T2},OperatorQuantumObject},
+) where {TF<:Number,TR<:Real,T2}
+    return real(tr(O * ρ))
 end
 
 @doc raw"""

test/quantum_objects.jl

@@ -317,38 +317,56 @@
 
     @testset "projection" begin
         N = 10
-        a = fock(N, 3)
-        @test proj(a) ≈ proj(a') ≈ sparse(ket2dm(a)) ≈ projection(N, 3, 3)
-        @test isket(a') == false
-        @test isbra(a') == true
-        @test shape(a) == (N,)
-        @test shape(a') == (1, N)
-        @test norm(a) ≈ 1
-        @test norm(a') ≈ 1
+        ψ = fock(N, 3)
+        @test proj(ψ) ≈ proj(ψ') ≈ sparse(ket2dm(ψ)) ≈ projection(N, 3, 3)
+        @test isket(ψ') == false
+        @test isbra(ψ') == true
+        @test shape(ψ) == (N,)
+        @test shape(ψ') == (1, N)
+        @test norm(ψ) ≈ 1
+        @test norm(ψ') ≈ 1
+
+        @testset "Type Inference (proj)" begin
+            @inferred proj(ψ)
+            @inferred proj(ψ')
+        end
     end
 
     @testset "dot product" begin
         ψ = rand_ket(10)
         @test dot(ψ, ψ) ≈ ψ' * ψ ≈ norm(ψ) ≈ 1.0
+
+        @testset "Type Inference (dot)" begin
+            @inferred dot(ψ, ψ)
+            @inferred ψ' * ψ
+            @inferred norm(ψ)
+        end
     end
 
     @testset "normalization" begin
         # normalize, normalize!, unit
         N = 10
-        a = Qobj(rand(ComplexF64, N))
-        M = a * a'
-        @test (norm(a) ≈ 1) == false
+        ψ = Qobj(rand(ComplexF64, N))
+        M = ψ * ψ'
+        @test (norm(ψ) ≈ 1) == false
         @test (norm(M) ≈ 1) == false
-        @test (norm(unit(a)) ≈ 1) == true
+        @test (norm(unit(ψ)) ≈ 1) == true
         @test (norm(unit(M)) ≈ 1) == true
-        @test (norm(a) ≈ 1) == false # Again, to be sure that it is still non-normalized
+        @test (norm(ψ) ≈ 1) == false # Again, to be sure that it is still non-normalized
         @test (norm(M) ≈ 1) == false # Again, to be sure that it is still non-normalized
-        normalize!(a)
+        normalize!(ψ)
         normalize!(M)
-        @test (norm(a) ≈ 1) == true
+        @test (norm(ψ) ≈ 1) == true
         @test (norm(M) ≈ 1) == true
-        @test M ≈ a * a'
+        @test M ≈ ψ * ψ'
         @test (unit(qeye(N)) ≈ (qeye(N) / N)) == true
+
+        @testset "Type Inference (normalize)" begin
+            ψ = Qobj(rand(ComplexF64, N))
+            M = ket2dm(ψ)
+            @inferred normalize(ψ)
+            @inferred normalize(M)
+        end
     end
 
     @testset "expectation value" begin
@@ -362,6 +380,19 @@
         ψ = fock(N, 3)
         @test norm(ψ' * a) ≈ 2
         @test expect(a' * a, ψ' * a) ≈ 16
+
+        ρ = rand_dm(N)
+        @test expect(a, ρ) ≈ tr(a * ρ)
+        @test variance(a, ρ) ≈ tr(a^2 * ρ) - tr(a * ρ)^2
+
+        @testset "Type Inference (expect)" begin
+            @inferred expect(a, ψ)
+            @inferred expect(a, ψ')
+            @inferred variance(a, ψ)
+            @inferred variance(a, ψ')
+            @inferred expect(a, ρ)
+            @inferred variance(a, ρ)
+        end
     end
 
     @testset "get coherence" begin
@@ -371,6 +402,11 @@
         ρ = ket2dm(ψ)
         α, δρ = get_coherence(ρ)
         @test isapprox(abs(α), 3, atol = 1e-5) && abs2(δρ[1, 1]) > 0.999
+
+        @testset "Type Inference (get_coherence)" begin
+            @inferred get_coherence(ψ)
+            @inferred get_coherence(ρ)
+        end
     end
 
     @testset "SVD and Schatten p-norm" begin
@@ -382,6 +418,13 @@
         @test svdvals(vs)[1] ≈ √(vs' * vs)
         @test norm(Md, 1) ≈ sum(sqrt, abs.(eigenenergies(Md' * Md))) atol = 1e-6
         @test norm(Ms, 1) ≈ sum(sqrt, abs.(eigenenergies(Ms' * Ms))) atol = 1e-6
+
+        @testset "Type Inference (SVD and Schatten p-norm)" begin
+            @inferred svdvals(vd)
+            @inferred svdvals(vs)
+            @inferred norm(Md, 1)
+            @inferred norm(Ms, 1)
+        end
     end
 
     @testset "purity" begin
@@ -395,6 +438,13 @@
         @test purity(ψd) ≈ norm(ψd)^2 ≈ 1.0
         @test purity(ρ1) ≈ 1.0
         @test (1.0 / N) <= purity(ρ2) <= 1.0
+
+        @testset "Type Inference (purity)" begin
+            @inferred purity(ψ)
+            @inferred purity(ψd)
+            @inferred purity(ρ1)
+            @inferred purity(ρ2)
+        end
     end
 
     @testset "trace distance" begin
@@ -407,6 +457,13 @@
         @test tracedist(ρz0, ψz1) ≈ 1.0
         @test tracedist(ψz1, ρz0) ≈ 1.0
         @test tracedist(ρz0, ρz1) ≈ 1.0
+
+        @testset "Type Inference (trace distance)" begin
+            @inferred tracedist(ψz0, ψx0)
+            @inferred tracedist(ρz0, ψz1)
+            @inferred tracedist(ψz1, ρz0)
+            @inferred tracedist(ρz0, ρz1)
+        end
     end
 
     @testset "sqrt and fidelity" begin
@@ -418,6 +475,13 @@
         @test sqrtm(M0) ≈ sqrtm(sparse_to_dense(M0))
         @test isapprox(fidelity(M0, M1), fidelity(ψ1, M0); atol = 1e-6)
         @test isapprox(fidelity(ψ1, ψ2), fidelity(ket2dm(ψ1), ket2dm(ψ2)); atol = 1e-6)
+
+        @testset "Type Inference (sqrt and fidelity)" begin
+            @inferred sqrtm(M0)
+            @inferred fidelity(M0, M1)
+            @inferred fidelity(ψ1, M0)
+            @inferred fidelity(ψ1, ψ2)
+        end
     end
 
     @testset "log, exp, sinm, cosm" begin
@@ -469,6 +533,13 @@
         ρ2 = dense_to_sparse(ρ1)
         @test tidyup(ρ2, tol) != ρ2
         @test dense_to_sparse(tidyup(ρ1, tol)) == tidyup(ρ2, tol)
+
+        @testset "Type Inference (tidyup)" begin
+            @inferred tidyup(ψ1, tol)
+            @inferred tidyup(ρ1, tol)
+            @inferred tidyup(ψ2, tol)
+            @inferred tidyup(ρ2, tol)
+        end
     end
 
     @testset "ptrace" begin
@@ -550,5 +621,11 @@
         @test_throws ArgumentError permute(bra_bdca, wrong_order2)
         @test_throws ArgumentError permute(op_bdca, wrong_order1)
         @test_throws ArgumentError permute(op_bdca, wrong_order2)
+
+        @testset "Type Inference (permute)" begin
+            @inferred permute(ket_bdca, (2, 4, 3, 1))
+            @inferred permute(bra_bdca, (2, 4, 3, 1))
+            @inferred permute(op_bdca, (2, 4, 3, 1))
+        end
     end
 end