Refactor dimensions handling and update QuantumObject structure for improved clarity and consistency

Author: albertomercurio

src/qobj/arithmetic_and_attributes.jl

@@ -65,11 +65,11 @@ end
 
 function Base.:(*)(A::AbstractQuantumObject{Operator}, B::QuantumObject{Ket})
     check_mul_dimensions(get_dimensions_from(A), get_dimensions_to(B))
-    return QuantumObject(A.data * B.data, Ket(), Dimensions(get_dimensions_to(A)))
+    return QuantumObject(A.data * B.data, Ket(), ProductDimensions(get_dimensions_to(A)))
 end
 function Base.:(*)(A::QuantumObject{Bra}, B::AbstractQuantumObject{Operator})
     check_mul_dimensions(get_dimensions_from(A), get_dimensions_to(B))
-    return QuantumObject(A.data * B.data, Bra(), Dimensions(get_dimensions_from(B)))
+    return QuantumObject(A.data * B.data, Bra(), ProductDimensions(get_dimensions_from(B)))
 end
 function Base.:(*)(A::QuantumObject{Ket}, B::QuantumObject{Bra})
     check_dimensions(A, B)
@@ -653,15 +653,15 @@ Get the coherence value ``\alpha`` by measuring the expectation value of the des
 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|``.
 """
 function get_coherence(ψ::QuantumObject{Ket})
-    a = destroy(prod(ψ.dimensions))
+    a = destroy(hilbert_dimensions_to_size(ψ.dimensions)[1])
     α = expect(a, ψ)
     D = exp(α * a' - conj(α) * a)
 
     return α, D' * ψ
 end
 
 function get_coherence(ρ::QuantumObject{Operator})
-    a = destroy(prod(ρ.dimensions))
+    a = destroy(hilbert_dimensions_to_size(ρ.dimensions)[1])
     α = expect(a, ρ)
     D = exp(α * a' - conj(α) * a)
 

src/qobj/dimensions.jl

@@ -13,16 +13,19 @@ abstract type AbstractDimensions{M,N} end
     end
 
 A structure that describes the left-hand side (`to`) and right-hand side (`from`) dimensions of a quantum object.
+
+The `from` field can be different from `nothing` only for non-square [`Operator`](@ref) and [`SuperOperator`](@ref) quantum objects.
 """
 struct ProductDimensions{M,N,T1<:Tuple,T2<:Union{<:Tuple,Nothing}} <: AbstractDimensions{M,N}
     to::T1   # space acting on the left
     from::T2 # space acting on the right
 
-    function ProductDimensions(to::NTuple{M}, from::Union{NTuple,Nothing}) where {M}
+    function ProductDimensions(to::Union{AbstractVector,NTuple}, from::Union{NTuple,Nothing})
+        M = length(to)
         N = isnothing(from) ? M : length(from)
 
-        _non_static_array_warning("to", to)
-        isnothing(from) || _non_static_array_warning("from", from)
+        _non_static_array_warning("dims", to)
+        isnothing(from) || _non_static_array_warning("dims", from)
 
         to_space = _dims_tuple_of_space(to)
         from_space = _dims_tuple_of_space(from)
@@ -37,18 +40,25 @@ function ProductDimensions(dims::Union{AbstractVector,Tuple})
 end
 
 ProductDimensions(dims::Union{Int,AbstractSpace}) = ProductDimensions((dims,), nothing)
-ProductDimensions(dims::Union{AbstractVector{T},NTuple{N,T}}) where {T<:Integer,N} = ProductDimensions(dims, nothing)
+ProductDimensions(dims::Union{AbstractVector{<:Integer},NTuple{N,Integer}}) where {N} = ProductDimensions(dims, nothing)
+ProductDimensions(dims::Union{AbstractVector{<:AbstractSpace},NTuple{N,AbstractSpace}}) where {N} = ProductDimensions(dims, nothing)
 ProductDimensions(dims::ProductDimensions) = dims
 
 # obtain dims in the type of SVector with integers
-dimensions_to_dims(dimensions::NTuple{N,AbstractSpace}) where {N} = SVector{N}(map(dimensions_to_dims, dimensions))
-dimensions_to_dims(dimensions::ProductDimensions) =
-    SVector{2}(dimensions_to_dims(dimensions.to), dimensions_to_dims(dimensions.from))
+dimensions_to_dims(dimensions::NTuple{N,AbstractSpace}) where {N} = vcat(map(dimensions_to_dims, dimensions)...)
+function dimensions_to_dims(dimensions::ProductDimensions)
+    dims_to = dimensions_to_dims(dimensions.to)
+    isnothing(dimensions.from) && return dims_to
+    dims_from = dimensions_to_dims(dimensions.from)
+    return SVector{2}(dims_to, dims_from)
+end
 dimensions_to_dims(::Nothing) = nothing # for EigsolveResult.dimensions = nothing
 
 hilbert_dimensions_to_size(dimensions::ProductDimensions) =
     (hilbert_dimensions_to_size(dimensions.to), hilbert_dimensions_to_size(dimensions.from))
 hilbert_dimensions_to_size(dimensions::NTuple{N,AbstractSpace}) where {N} = prod(hilbert_dimensions_to_size, dimensions)
+hilbert_dimensions_to_size(dim::Int) = dim
+hilbert_dimensions_to_size(dimensions::Union{AbstractVector, NTuple{N,Integer}}) where {N} = prod(dimensions)
 hilbert_dimensions_to_size(::Nothing) = nothing
 
 liouvillian_dimensions_to_size(dimensions::ProductDimensions) =
@@ -66,9 +76,10 @@ Base.adjoint(dimensions::AbstractDimensions) = transpose(dimensions)
 Base.:(==)(dim1::ProductDimensions, dim2::ProductDimensions) = (dim1.to == dim2.to) && (dim1.from == dim2.from)
 
 _dims_tuple_of_space(dims::NTuple{N,AbstractSpace}) where {N} = dims
-function _dims_tuple_of_space(dims::NTuple{N, Integer}) where {N}
+function _dims_tuple_of_space(dims::Union{AbstractVector{<:Integer},SVector{M, <:Integer}, NTuple{M, Integer}}) where {M}
     _non_static_array_warning("dims", dims)
 
+    N = length(dims)
     N > 0 || throw(DomainError(N, "The length of `dims` must be larger or equal to 1."))
 
     return ntuple(dim -> HilbertSpace(dims[dim]), Val(N))
@@ -81,7 +92,10 @@ _gen_dimensions(dims) = ProductDimensions(dims)
 # this is used to show `dims` for Qobj and QobjEvo
 function _get_dims_string(dimensions::ProductDimensions)
     dims = dimensions_to_dims(dimensions)
-    isnothing(dims[2]) && return string(dims[1])
+    isnothing(dimensions.from) && return string(dims)
     return "[$(string(dims[1])), $(string(dims[2]))]"
 end
 _get_dims_string(::Nothing) = "nothing" # for EigsolveResult.dimensions = nothing
+
+space_one_list(dimensions::NTuple{N,AbstractSpace}) where {N} =
+    ntuple(i -> one(dimensions[i]), Val(sum(length, dimensions)))

src/qobj/functions.jl

@@ -186,41 +186,35 @@ julia> a.dims, O.dims
 ([20], [20, 20])
 ```
 """
-Base.kron(A::AbstractQuantumObject, B::AbstractQuantumObject)
-
-for type in (:Operator, :SuperOperator)
-    @eval function Base.kron(
-        A::AbstractQuantumObject{$type,ProductDimensions},
-        B::AbstractQuantumObject{$type,ProductDimensions},
-    )
-        QType = promote_op_type(A, B)
-        _lazy_tensor_warning(A.data, B.data)
-
-        A_dims_from = get_dimensions_from(A)
-        B_dims_from = get_dimensions_from(B)
-        AB_dims_to = (get_dimensions_to(A)..., get_dimensions_to(B)...)
-        AB_dims_from = (isnothing(A_dims_from) && isnothing(B_dims_from)) ? nothing : (A_dims_from..., B_dims_from...)
-
-        return QType(kron(A.data, B.data), $type(), ProductDimensions(AB_dims_to, AB_dims_from))
-    end
+function Base.kron(
+    A::AbstractQuantumObject{ObjType,<:ProductDimensions},
+    B::AbstractQuantumObject{ObjType,<:ProductDimensions},
+) where {ObjType<:Union{Ket, Bra, Operator}}
+    QType = promote_op_type(A, B)
+    _lazy_tensor_warning(A.data, B.data)
+
+    A_dims_from = get_dimensions_from(A)
+    B_dims_from = get_dimensions_from(B)
+    AB_dims_to = (get_dimensions_to(A)..., get_dimensions_to(B)...)
+    AB_dims_from = (A.dimensions.from === B.dimensions.from === nothing) ? nothing : (A_dims_from..., B_dims_from...)
+
+    return QType(kron(A.data, B.data), A.type, ProductDimensions(AB_dims_to, AB_dims_from))
 end
 
-# if A and B are different type (must return Operator with GeneralDimensions)
 for AOpType in (:Ket, :Bra, :Operator)
     for BOpType in (:Ket, :Bra, :Operator)
         if (AOpType != BOpType)
             @eval begin
                 function Base.kron(A::AbstractQuantumObject{$AOpType}, B::AbstractQuantumObject{$BOpType})
                     QType = promote_op_type(A, B)
                     _lazy_tensor_warning(A.data, B.data)
-                    return QType(
-                        kron(A.data, B.data),
-                        Operator(),
-                        GeneralDimensions(
-                            (get_dimensions_to(A)..., get_dimensions_to(B)...),
-                            (get_dimensions_from(A)..., get_dimensions_from(B)...),
-                        ),
-                    )
+
+                    A_dims_from = get_dimensions_from(A)
+                    B_dims_from = get_dimensions_from(B)
+                    AB_dims_to = (get_dimensions_to(A)..., get_dimensions_to(B)...)
+                    AB_dims_from = (A.dimensions.from === B.dimensions.from === nothing) ? nothing : (A_dims_from..., B_dims_from...)
+
+                    return QType(kron(A.data, B.data), Operator(), ProductDimensions(AB_dims_to, AB_dims_from))
                 end
             end
         end

src/qobj/quantum_object.jl

@@ -55,8 +55,7 @@ struct QuantumObject{ObjType<:QuantumObjectType,DimType<:AbstractDimensions,Data
 
         ObjType = _check_type(type)
 
-        _size = _get_size(data)
-        _check_QuantumObject(type, dimensions, _size[1], _size[2])
+        _check_QuantumObject(type, dimensions, size(data, 1), size(data, 2))
 
         return new{ObjType,typeof(dimensions),DT}(data, type, dimensions)
     end
@@ -72,12 +71,10 @@ Generate [`QuantumObject`](@ref) with a given `A::AbstractArray` and specified `
     `Qobj` is a synonym of `QuantumObject`.
 """
 function QuantumObject(A::AbstractMatrix{T}; type = nothing, dims = nothing) where {T}
-    _size = _get_size(A)
-
     _check_type(type)
 
-    if type isa Nothing
-        type = (_size[1] == 1 && _size[2] > 1) ? Bra() : Operator() # default type
+    if isnothing(type)
+        type = (size(A, 1) == 1 && size(A, 2) > 1) ? Bra() : Operator() # default type
     elseif !(type isa Operator) && !(type isa SuperOperator) && !(type isa Bra) && !(type isa OperatorBra)
         throw(
             ArgumentError(
@@ -86,15 +83,15 @@ function QuantumObject(A::AbstractMatrix{T}; type = nothing, dims = nothing) whe
         )
     end
 
-    if dims isa Nothing
+    if isnothing(dims)
         if type isa Bra
-            dims = Dimensions(_size[2])
+            dims = ProductDimensions(size(A, 2))
         elseif type isa Operator
-            dims =
-                (_size[1] == _size[2]) ? Dimensions(_size[1]) :
-                GeneralDimensions(SVector{2}(SVector{1}(_size[1]), SVector{1}(_size[2])))
+            dims_to = (size(A, 1), )
+            dims_from = (size(A, 1) == size(A, 2)) ? nothing : (size(A, 2), )
+            dims = ProductDimensions(dims_to, dims_from)
         elseif type isa SuperOperator || type isa OperatorBra
-            dims = Dimensions(isqrt(_size[2]))
+            dims = ProductDimensions(isqrt(size(A, 2)))
         end
     end
 
@@ -103,18 +100,18 @@ end
 
 function QuantumObject(A::AbstractVector{T}; type = nothing, dims = nothing) where {T}
     _check_type(type)
-    if type isa Nothing
+
+    if isnothing(type)
         type = Ket() # default type
     elseif !(type isa Ket) && !(type isa OperatorKet)
         throw(ArgumentError("The argument type must be Ket() or OperatorKet() if the input array is a vector."))
     end
 
-    if dims isa Nothing
-        _size = _get_size(A)
+    if isnothing(dims)
         if type isa Ket
-            dims = Dimensions(_size[1])
+            dims = ProductDimensions(size(A, 1))
         elseif type isa OperatorKet
-            dims = Dimensions(isqrt(_size[1]))
+            dims = ProductDimensions(isqrt(size(A, 1)))
         end
     end
 
@@ -126,10 +123,9 @@ function QuantumObject(A::AbstractArray{T,N}; type = nothing, dims = nothing) wh
 end
 
 function QuantumObject(A::QuantumObject; type = A.type, dims = A.dimensions)
-    _size = _get_size(A.data)
     dimensions = _gen_dimensions(dims)
     _check_type(type)
-    _check_QuantumObject(type, dimensions, _size[1], _size[2])
+    _check_QuantumObject(type, dimensions, size(A, 1), size(A, 2))
     return QuantumObject(copy(A.data), type, dimensions)
 end
 

src/qobj/quantum_object_base.jl

@@ -210,8 +210,8 @@ end
 
 get_dimensions_to(A::AbstractQuantumObject) = A.dimensions.to
 get_dimensions_from(A::AbstractQuantumObject) = isnothing(A.dimensions.from) ? A.dimensions.to : A.dimensions.from
-get_dimensions_from(::AbstractQuantumObject{Ket}) = nothing
-get_dimensions_to(::AbstractQuantumObject{Bra}) = nothing
+get_dimensions_from(A::AbstractQuantumObject{Ket}) = space_one_list(A.dimensions.to)
+get_dimensions_to(A::AbstractQuantumObject{Bra}) = space_one_list(A.dimensions.to)
 
 # functions for getting Float or Complex element type
 _float_type(A::AbstractQuantumObject) = _float_type(eltype(A))

src/qobj/quantum_object_evo.jl

@@ -124,8 +124,7 @@ struct QuantumObjectEvolution{
 
         dimensions = _gen_dimensions(dims)
 
-        _size = _get_size(data)
-        _check_QuantumObject(type, dimensions, _size[1], _size[2])
+        _check_QuantumObject(type, dimensions, size(data, 1), size(data, 2))
 
         return new{ObjType,typeof(dimensions),DT}(data, type, dimensions)
     end
@@ -158,16 +157,15 @@ Generate a [`QuantumObjectEvolution`](@ref) object from a [`SciMLOperator`](http
 Note that `QobjEvo` is a synonym of `QuantumObjectEvolution`
 """
 function QuantumObjectEvolution(data::AbstractSciMLOperator; type = Operator(), dims = nothing)
-    _size = _get_size(data)
     _check_type(type)
 
-    if dims isa Nothing
+    if isnothing(dims)
         if type isa Operator
-            dims =
-                (_size[1] == _size[2]) ? Dimensions(_size[1]) :
-                GeneralDimensions(SVector{2}(SVector{1}(_size[1]), SVector{1}(_size[2])))
+            dims_to = (size(data, 1), )
+            dims_from = (size(data, 1) == size(data, 2)) ? nothing : (size(data, 2), )
+            dims = ProductDimensions(dims_to, dims_from)
         elseif type isa SuperOperator
-            dims = Dimensions(isqrt(_size[2]))
+            dims = ProductDimensions(isqrt(size(data, 2)))
         end
     end
 

src/qobj/space.jl

@@ -22,6 +22,9 @@ struct HilbertSpace <: AbstractSpace
     end
 end
 
+Base.one(::HilbertSpace) = HilbertSpace(1)
+Base.length(::HilbertSpace) = 1
+
 dimensions_to_dims(s::HilbertSpace) = SVector{1,Int}(s.size)
 hilbert_dimensions_to_size(s::HilbertSpace) = s.size
 liouvillian_dimensions_to_size(s::HilbertSpace) = s.size^2

src/utilities.jl

@@ -143,10 +143,6 @@ makeVal(x) = Val(x)
 getVal(x::Val{T}) where {T} = T
 getVal(x) = x # getVal for any other type
 
-_get_size(A::AbstractMatrix) = size(A)
-_get_size(A::AbstractVector) = (length(A), 1)
-_get_size(A::AbstractSciMLOperator) = size(A)
-
 _non_static_array_warning(argname, arg::Tuple{}) =
     throw(ArgumentError("The argument $argname must be a Tuple or a StaticVector of non-zero length."))
 _non_static_array_warning(argname, arg::Union{SVector{N,T},MVector{N,T},NTuple{N,T}}) where {N,T} = nothing

test/core-test/quantum_objects.jl

@@ -17,7 +17,7 @@
     # DomainError: incompatible between size of array and type
     @testset "DomainError" begin
         a = rand(ComplexF64, 3, 2)
-        for t in [SuperOperator(), Bra(), OperatorBra()]
+        for t in (Bra(), OperatorBra())
             @test_throws DomainError Qobj(a, type = t)
         end
 
@@ -33,17 +33,14 @@
         @test_throws DomainError Qobj(rand(ComplexF64, 2, 1), type = Operator()) # should be type = Bra
 
         # check that Ket, Bra, SuperOperator, OperatorKet, and OperatorBra don't support GeneralDimensions
-        @test_throws DomainError Qobj(rand(ComplexF64, 2), type = Ket(), dims = ((2,), (1,)))
-        @test_throws DomainError Qobj(rand(ComplexF64, 1, 2), type = Bra(), dims = ((1,), (2,)))
-        @test_throws DomainError Qobj(rand(ComplexF64, 4, 4), type = SuperOperator(), dims = ((2,), (2,)))
-        @test_throws DomainError Qobj(rand(ComplexF64, 4), type = OperatorKet(), dims = ((2,), (1,)))
-        @test_throws DomainError Qobj(rand(ComplexF64, 1, 4), type = OperatorBra(), dims = ((1,), (2,)))
+        @test_throws DimensionMismatch Qobj(rand(ComplexF64, 2), type = Ket(), dims = ((2,), (1,)))
+        @test_throws DimensionMismatch Qobj(rand(ComplexF64, 1, 2), type = Bra(), dims = ((1,), (2,)))
     end
 
     # unsupported type of dims
     @testset "unsupported dims" begin
-        @test_throws ArgumentError Qobj(rand(2, 2), dims = 2.0)
-        @test_throws ArgumentError Qobj(rand(2, 2), dims = 2.0 + 0.0im)
+        @test_throws MethodError Qobj(rand(2, 2), dims = 2.0)
+        @test_throws MethodError Qobj(rand(2, 2), dims = 2.0 + 0.0im)
         @test_throws DomainError Qobj(rand(2, 2), dims = 0)
         @test_throws DomainError Qobj(rand(2, 2), dims = (2, -2))
         @test_logs (
@@ -367,8 +364,8 @@
             end
 
             UnionType = Union{
-                QuantumObject{Bra,Dimensions{1,Tuple{Space}},Matrix{T}},
-                QuantumObject{Operator,Dimensions{1,Tuple{Space}},Matrix{T}},
+                QuantumObject{Bra,ProductDimensions{1,1,Tuple{HilbertSpace},Nothing},Matrix{T}},
+                QuantumObject{Operator,ProductDimensions{1,1,Tuple{HilbertSpace},Nothing},Matrix{T}},
             }
             a = rand(T, 1, N)
             @inferred UnionType Qobj(a)
@@ -377,8 +374,8 @@
             end
 
             UnionType2 = Union{
-                QuantumObject{Operator,GeneralDimensions{1,1,Tuple{Space},Tuple{Space}},Matrix{T}},
-                QuantumObject{Operator,Dimensions{1,Tuple{Space}},Matrix{T}},
+                QuantumObject{Operator,ProductDimensions{1,1,Tuple{HilbertSpace},Tuple{HilbertSpace}},Matrix{T}},
+                QuantumObject{Operator,ProductDimensions{1,1,Tuple{HilbertSpace},Nothing},Matrix{T}},
             }
             a = rand(T, N, N)
             @inferred UnionType Qobj(a)