support real and imag for Qobj (#202)

qutip · Aug 13, 2024 · 346bdad · 346bdad
2 parents 1dc43b9 + 15531bc
commit 346bdad
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/src/qobj/quantum_object.jl b/src/qobj/quantum_object.jl
@@ -330,6 +330,9 @@ Base.isapprox(A::QuantumObject{<:AbstractArray{T}}, B::QuantumObject{<:AbstractA
 Base.:(==)(A::QuantumObject{<:AbstractArray{T}}, B::QuantumObject{<:AbstractArray{T}}) where {T} =
     (A.type == B.type) && (A.dims == B.dims) && (A.data == B.data)
 
+Base.real(x::QuantumObject) = QuantumObject(real(x.data), x.type, x.dims)
+Base.imag(x::QuantumObject) = QuantumObject(imag(x.data), x.type, x.dims)
+
 SparseArrays.sparse(A::QuantumObject{<:AbstractArray{T}}) where {T} = QuantumObject(sparse(A.data), A.type, A.dims)
 SparseArrays.nnz(A::QuantumObject{<:AbstractSparseArray}) = nnz(A.data)
 SparseArrays.nonzeros(A::QuantumObject{<:AbstractSparseArray}) = nonzeros(A.data)

diff --git a/test/quantum_objects.jl b/test/quantum_objects.jl
@@ -145,6 +145,8 @@
         @test isequal(a4, a3) == false
         @test a4 ≈ a2
 
+        @test real(a2).data == real(a)
+        @test imag(a2).data == imag(a)
         @test +a2 == a2
         @test -(-a2) == a2
         @test a2^3 ≈ a2 * a2 * a2