Resolve type instabilities when generating QuantumObjects

src/qobj/quantum_object.jl

@@ -141,6 +141,87 @@
     dims::Vector{Int}
 end
 
+function QuantumObject(
+    A::AbstractMatrix{T};
+    type::ObjType = nothing,
+    dims = nothing,
+) where {T,ObjType<:Union{Nothing,QuantumObjectType}}
+    if type isa Nothing
+        type = Operator # default type
+    end
+
+    if type != Operator && type != SuperOperator
+        throw(
+            ArgumentError(
+                "The argument type must be OperatorQuantumObject or SuperOperatorQuantumObject if the input array is a matrix.",
+            ),
+        )
+    end
+
+    _size = size(A)
+
+    if dims isa Nothing
+        if type isa OperatorQuantumObject
+            dims = [_size[1]]
+        elseif type isa SuperOperatorQuantumObject
+            dims = [isqrt(_size[1])]
+        end
+    else
+        _check_dims(dims)
+    end
+
+    _check_QuantumObject(type, dims, _size[1], _size[2])
+    return QuantumObject(A, type, dims)
+end
+
+function QuantumObject(
+    A::AbstractVector{T};
+    type::ObjType = nothing,
+    dims = nothing,
+) where {T,ObjType<:Union{Nothing,QuantumObjectType}}
+    if type isa Nothing
+        type = Ket # default type
+    end
+
+    if type != Ket && type != Bra && type != OperatorKet && type != OperatorBra
+        throw(
+            ArgumentError(
+                "The argument type must be KetQuantumObject, BraQuantumObject, OperatorKetQuantumObject or OperatorBraQuantumObject if the input array is a vector.",
+            ),
+        )
+    end
+
+    _size = (length(A), 1)
+
+    if dims isa Nothing
+        if (type isa KetQuantumObject) || (type isa BraQuantumObject)
+            dims = [_size[1]]
+        elseif (type isa OperatorKetQuantumObject) || (type isa OperatorBraQuantumObject)
+            dims = [isqrt(_size[1])]
+        end
+    else
+        _check_dims(dims)
+    end
+
+    if (type isa BraQuantumObject) || (type isa OperatorBraQuantumObject)
+        data = A'
+        _size = (1, length(A))
+    else
+        data = A
+    end
+
+    _check_QuantumObject(type, dims, _size[1], _size[2])
+    return QuantumObject(data, type, dims)
+end
+
+function QuantumObject(
+    A::AbstractArray{T,N};
+    type::ObjType = nothing,
+    dims = nothing,
+) where {T,N,ObjType<:Union{Nothing,QuantumObjectType}}
+    throw(ArgumentError("The input array must be a vector or a matrix."))
+end
+
 function Base.show(
     io::IO,
     QO::QuantumObject{<:AbstractArray{T},OpType},
@@ -175,51 +256,6 @@
     return show(io, MIME("text/plain"), op_data)
 end
 
-function QuantumObject(
-    A::AbstractArray{T,N};
-    type::ObjType = nothing,
-    dims = nothing,
-) where {T,N,ObjType<:Union{Nothing,QuantumObjectType}}
-
-    # only accept 1D- and 2D-array
-    if N == 1
-        Size = (length(A), 1)
-    else
-        Size = size(A)
-        (N > 2) && throw(DomainError(Size, "The dimension of the array is not compatible with Quantum Object"))
-    end
-
-    # decide QuantumObjectType from the size of A
-    if type isa Nothing
-        if Size[1] == Size[2]
-            type = Operator
-        elseif Size[2] == 1
-            type = Ket
-        elseif Size[1] == 1
-            type = Bra
-        else
-            throw(DomainError(Size, "The dimension of the array is not compatible with Quantum Object"))
-        end
-    end
-
-    # decide dims from the size of A and the given type
-    if dims isa Nothing
-        if (type isa KetQuantumObject) || (type isa OperatorQuantumObject)
-            dims = [Size[1]]
-        elseif (type isa SuperOperatorQuantumObject) || (type isa OperatorKetQuantumObject)
-            dims = [isqrt(Size[1])]
-        elseif type isa BraQuantumObject
-            dims = [Size[2]]
-        elseif type isa OperatorBraQuantumObject
-            dims = [isqrt(Size[2])]
-        end
-    else
-        _check_dims(dims)
-    end
-    _check_QuantumObject(type, dims, Size[1], Size[2])
-    return QuantumObject(A, type, dims)
-end
-
 _check_dims(dims::Vector{Int}) =
     all(>(0), dims) || throw(DomainError(dims, "The argument dims must be a vector with positive integers."))
 _check_dims(dims::Any) = throw(ArgumentError("The argument dims must be a vector with positive integers."))
@@ -269,8 +305,8 @@
     type::ObjType = A.type,
     dims = A.dims,
 ) where {T,N,ObjType<:QuantumObjectType}
-    Size = N == 1 ? (length(A), 1) : size(A)
-    _check_QuantumObject(type, dims, Size[1], Size[2])
+    _size = N == 1 ? (length(A), 1) : size(A)
+    _check_QuantumObject(type, dims, _size[1], _size[2])
     return QuantumObject(copy(A.data), type, dims)
 end
 

test/quantum_objects.jl

@@ -1,9 +1,15 @@
 @testset "Quantum Objects" begin
     # unsupported size of array
-    for a in [rand(ComplexF64, 3, 2), rand(ComplexF64, 2, 2, 2)]
-        for t in [nothing, Ket, Bra, Operator, SuperOperator, OperatorBra, OperatorKet]
-            @test_throws DomainError Qobj(a, type = t)
-        end
+    a = rand(ComplexF64, 3, 2)
+    for t in [nothing, Operator, SuperOperator]
+        @test_throws DomainError Qobj(a, type = t)
+    end
+    for t in [Ket, Bra, OperatorBra, OperatorKet]
+        @test_throws ArgumentError Qobj(a, type = t)
+    end
+    a = rand(ComplexF64, 2, 2, 2)
+    for t in [nothing, Ket, Bra, Operator, SuperOperator, OperatorBra, OperatorKet]
+        @test_throws ArgumentError Qobj(a, type = t)
     end
 
     # unsupported dims
@@ -16,19 +22,17 @@
     N = 10
     a = rand(ComplexF64, 10)
     # @test_logs (:warn, "The norm of the input data is not one.") QuantumObject(a)
-    @test_throws DomainError Qobj(a, type = Bra)
-    @test_throws DomainError Qobj(a, type = Operator)
-    @test_throws DomainError Qobj(a, type = SuperOperator)
-    @test_throws DomainError Qobj(a, type = OperatorBra)
-    @test_throws DomainError Qobj(a', type = Ket)
+    @test_throws ArgumentError Qobj(a, type = Operator)
+    @test_throws ArgumentError Qobj(a, type = SuperOperator)
+    @test_throws ArgumentError Qobj(a', type = Ket)
     @test_throws DomainError Qobj(a', type = Operator)
     @test_throws DomainError Qobj(a', type = SuperOperator)
-    @test_throws DomainError Qobj(a', type = OperatorKet)
+    @test_throws ArgumentError Qobj(a', type = OperatorKet)
     @test_throws DimensionMismatch Qobj(a, dims = [2])
-    @test_throws DimensionMismatch Qobj(a', dims = [2])
-    a2 = Qobj(a')
+    @test_throws DomainError Qobj(a', dims = [2])
+    a2 = Qobj(a, type = Bra)
     a3 = Qobj(a)
-    @test dag(a3) == a2
+    @test a3' == a2
     @test isket(a2) == false
     @test isbra(a2) == true
     @test isoper(a2) == false
@@ -43,7 +47,6 @@
     @test isoperbra(a3) == false
     @test Qobj(a3) == a3
     @test !(Qobj(a3) === a3)
-    @test isket(Qobj(Matrix([2 3])')) == true
 
     a = sprand(ComplexF64, 100, 100, 0.1)
     a2 = Qobj(a)
@@ -62,15 +65,14 @@
     @test isoperket(a3) == false
     @test isoperbra(a3) == false
     @test_throws DimensionMismatch Qobj(a, dims = [2])
-    @test_throws DimensionMismatch Qobj(a, dims = [2])
 
     # Operator-Ket, Operator-Bra tests
     H = 0.3 * sigmax() + 0.7 * sigmaz()
     L = liouvillian(H)
     ρ = Qobj(rand(ComplexF64, 2, 2))
     ρ_ket = mat2vec(ρ)
     ρ_bra = ρ_ket'
-    @test ρ_bra == Qobj(mat2vec(ρ.data)', type = OperatorBra)
+    @test ρ_bra == Qobj(mat2vec(ρ.data), type = OperatorBra)
     @test ρ == vec2mat(ρ_ket)
     @test isket(ρ_ket) == false
     @test isbra(ρ_ket) == false
@@ -94,7 +96,7 @@
     @test (ρ_bra * L')' == L * ρ_ket
     @test sum((conj(ρ) .* ρ).data) ≈ dot(ρ_ket, ρ_ket) ≈ ρ_bra * ρ_ket
     @test_throws DimensionMismatch Qobj(ρ_ket.data, type = OperatorKet, dims = [4])
-    @test_throws DimensionMismatch Qobj(ρ_bra.data, type = OperatorBra, dims = [4])
+    @test_throws ArgumentError Qobj(ρ_bra.data, type = OperatorBra, dims = [4])
 
     # matrix element
     s0 = Qobj(basis(4, 0).data; type = OperatorKet)
@@ -383,4 +385,19 @@
     @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" begin
+        for T in [Float32, Float64, ComplexF32, ComplexF64]
+            N = 25
+            a = rand(T, N)
+            for type in [Ket, Bra, OperatorKet, OperatorBra]
+                @inferred Qobj(a, type = type)
+            end
+
+            a = rand(T, N, N)
+            for type in [Operator, SuperOperator]
+                @inferred Qobj(a, type = type)
+            end
+        end
+    end
 end