first try of Space and Dimensions (#359)

Author: ytdHuang

src/QuantumToolbox.jl

@@ -79,6 +79,8 @@ include("progress_bar.jl")
 include("linear_maps.jl")
 
 # Quantum Object
+include("qobj/space.jl")
+include("qobj/dimensions.jl")
 include("qobj/quantum_object_base.jl")
 include("qobj/quantum_object.jl")
 include("qobj/quantum_object_evo.jl")

src/negativity.jl

@@ -77,28 +77,37 @@ end
 
 # for dense matrices
 function _partial_transpose(ρ::QuantumObject{<:AbstractArray,OperatorQuantumObject}, mask::Vector{Bool})
+    isa(ρ.dims, CompoundDimensions) &&
+        (ρ.to != ρ.from) &&
+        throw(ArgumentError("Invalid partial transpose for dims = $(ρ.dims)"))
+
     mask2 = [1 + Int(i) for i in mask]
     # mask2 has elements with values equal to 1 or 2
     #   1 - the subsystem don't need to be transposed
     #   2 - the subsystem need be transposed
 
     nsys = length(mask2)
+    dimslist = dims_to_list(ρ.to)
     pt_dims = reshape(Vector(1:(2*nsys)), (nsys, 2))
     pt_idx = [
         [pt_dims[n, mask2[n]] for n in 1:nsys] # origin   value in mask2
         [pt_dims[n, 3-mask2[n]] for n in 1:nsys]  # opposite value in mask2 (1 -> 2, and 2 -> 1)
     ]
     return QuantumObject(
-        reshape(permutedims(reshape(ρ.data, (ρ.dims..., ρ.dims...)), pt_idx), size(ρ)),
+        reshape(permutedims(reshape(ρ.data, (dimslist..., dimslist...)), pt_idx), size(ρ)),
         Operator,
-        ρ.dims,
+        Dimensions(ρ.dims.to),
     )
 end
 
 # for sparse matrices
 function _partial_transpose(ρ::QuantumObject{<:AbstractSparseArray,OperatorQuantumObject}, mask::Vector{Bool})
+    isa(ρ.dims, CompoundDimensions) &&
+        (ρ.to != ρ.from) &&
+        throw(ArgumentError("Invalid partial transpose for dims = $(ρ.dims)"))
+
     M, N = size(ρ)
-    dimsTuple = Tuple(ρ.dims)
+    dimsTuple = Tuple(dims_to_list(ρ.to))
     colptr = ρ.data.colptr
     rowval = ρ.data.rowval
     nzval = ρ.data.nzval

src/qobj/arithmetic_and_attributes.jl

@@ -50,26 +50,63 @@ for op in (:(+), :(-), :(*))
     end
 end
 
+function check_mul_dims(from::SVector{NA,AbstractSpace}, to::SVector{NB,AbstractSpace}) where {NA,NB}
+    (from != to) && throw(
+        DimensionMismatch(
+            "The quantum object with dims = $from can not multiply a quantum object with dims = $to on the right-hand side.",
+        ),
+    )
+    return nothing
+end
+
+for ADimType in (:Dimensions, :CompoundDimensions)
+    for BDimType in (:Dimensions, :CompoundDimensions)
+        if ADimType == BDimType == :Dimensions
+            @eval begin
+                function LinearAlgebra.:(*)(
+                    A::AbstractQuantumObject{DT1,OperatorQuantumObject,$ADimType{NA}},
+                    B::AbstractQuantumObject{DT2,OperatorQuantumObject,$BDimType{NB}},
+                ) where {DT1,DT2,NA,NB}
+                    check_dims(A, B)
+                    QType = promote_op_type(A, B)
+                    return QType(A.data * B.data, Operator, A.dims)
+                end
+            end
+        else
+            @eval begin
+                function LinearAlgebra.:(*)(
+                    A::AbstractQuantumObject{DT1,OperatorQuantumObject,$ADimType{NA}},
+                    B::AbstractQuantumObject{DT2,OperatorQuantumObject,$BDimType{NB}},
+                ) where {DT1,DT2,NA,NB}
+                    check_mul_dims(A.from, B.to)
+                    QType = promote_op_type(A, B)
+                    return QType(A.data * B.data, Operator, CompoundDimensions(A.to, B.from))
+                end
+            end
+        end
+    end
+end
+
 function LinearAlgebra.:(*)(
     A::AbstractQuantumObject{DT1,OperatorQuantumObject},
     B::QuantumObject{DT2,KetQuantumObject},
 ) where {DT1,DT2}
-    check_dims(A, B)
-    return QuantumObject(A.data * B.data, Ket, A.dims)
+    check_mul_dims(A.from, B.to)
+    return QuantumObject(A.data * B.data, Ket, Dimensions(A.to))
 end
 function LinearAlgebra.:(*)(
     A::QuantumObject{DT1,BraQuantumObject},
     B::AbstractQuantumObject{DT2,OperatorQuantumObject},
 ) where {DT1,DT2}
-    check_dims(A, B)
-    return QuantumObject(A.data * B.data, Bra, A.dims)
+    check_mul_dims(A.from, B.to)
+    return QuantumObject(A.data * B.data, Bra, Dimensions(B.from))
 end
 function LinearAlgebra.:(*)(
     A::QuantumObject{DT1,KetQuantumObject},
     B::QuantumObject{DT2,BraQuantumObject},
 ) where {DT1,DT2}
     check_dims(A, B)
-    return QuantumObject(A.data * B.data, Operator, A.dims)
+    return QuantumObject(A.data * B.data, Operator, A.dims) # to align with QuTiP, don't use kron(A, B) to do it.
 end
 function LinearAlgebra.:(*)(
     A::QuantumObject{DT1,BraQuantumObject},
@@ -196,7 +233,7 @@ Lazy matrix transpose of the [`AbstractQuantumObject`](@ref).
 LinearAlgebra.transpose(
     A::AbstractQuantumObject{DT,OpType},
 ) where {DT,OpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject}} =
-    get_typename_wrapper(A)(transpose(A.data), A.type, A.dims)
+    get_typename_wrapper(A)(transpose(A.data), A.type, transpose(A.dims))
 
 @doc raw"""
     A'
@@ -211,13 +248,15 @@ Lazy adjoint (conjugate transposition) of the [`AbstractQuantumObject`](@ref)
 LinearAlgebra.adjoint(
     A::AbstractQuantumObject{DT,OpType},
 ) where {DT,OpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject}} =
-    get_typename_wrapper(A)(adjoint(A.data), A.type, A.dims)
-LinearAlgebra.adjoint(A::QuantumObject{DT,KetQuantumObject}) where {DT} = QuantumObject(adjoint(A.data), Bra, A.dims)
-LinearAlgebra.adjoint(A::QuantumObject{DT,BraQuantumObject}) where {DT} = QuantumObject(adjoint(A.data), Ket, A.dims)
+    get_typename_wrapper(A)(adjoint(A.data), A.type, transpose(A.dims))
+LinearAlgebra.adjoint(A::QuantumObject{DT,KetQuantumObject}) where {DT} =
+    QuantumObject(adjoint(A.data), Bra, transpose(A.dims))
+LinearAlgebra.adjoint(A::QuantumObject{DT,BraQuantumObject}) where {DT} =
+    QuantumObject(adjoint(A.data), Ket, transpose(A.dims))
 LinearAlgebra.adjoint(A::QuantumObject{DT,OperatorKetQuantumObject}) where {DT} =
-    QuantumObject(adjoint(A.data), OperatorBra, A.dims)
+    QuantumObject(adjoint(A.data), OperatorBra, transpose(A.dims))
 LinearAlgebra.adjoint(A::QuantumObject{DT,OperatorBraQuantumObject}) where {DT} =
-    QuantumObject(adjoint(A.data), OperatorKet, A.dims)
+    QuantumObject(adjoint(A.data), OperatorKet, transpose(A.dims))
 
 @doc raw"""
     inv(A::AbstractQuantumObject)
@@ -557,14 +596,19 @@ function ptrace(QO::QuantumObject{<:AbstractArray,KetQuantumObject}, sel::Union{
         (n_d == 1) && return ket2dm(QO) # ptrace should always return Operator
     end
 
+    dimslist = dims_to_list(QO.dims)
     _sort_sel = sort(SVector{length(sel),Int}(sel))
-    ρtr, dkeep = _ptrace_ket(QO.data, QO.dims, _sort_sel)
-    return QuantumObject(ρtr, type = Operator, dims = dkeep)
+    ρtr, dkeep = _ptrace_ket(QO.data, dimslist, _sort_sel)
+    return QuantumObject(ρtr, type = Operator, dims = Dimensions(dkeep))
 end
 
 ptrace(QO::QuantumObject{<:AbstractArray,BraQuantumObject}, sel::Union{AbstractVector{Int},Tuple}) = ptrace(QO', sel)
 
 function ptrace(QO::QuantumObject{<:AbstractArray,OperatorQuantumObject}, sel::Union{AbstractVector{Int},Tuple})
+    isa(QO.dims, CompoundDimensions) &&
+        (QO.to != QO.from) &&
+        throw(ArgumentError("Invalid partial trace for dims = $(QO.dims)"))
+
     _non_static_array_warning("sel", sel)
 
     n_s = length(sel)
@@ -580,9 +624,10 @@ function ptrace(QO::QuantumObject{<:AbstractArray,OperatorQuantumObject}, sel::U
         (n_d == 1) && return QO
     end
 
+    dimslist = dims_to_list(QO.to)
     _sort_sel = sort(SVector{length(sel),Int}(sel))
-    ρtr, dkeep = _ptrace_oper(QO.data, QO.dims, _sort_sel)
-    return QuantumObject(ρtr, type = Operator, dims = dkeep)
+    ρtr, dkeep = _ptrace_oper(QO.data, dimslist, _sort_sel)
+    return QuantumObject(ρtr, type = Operator, dims = Dimensions(dkeep))
 end
 ptrace(QO::QuantumObject, sel::Int) = ptrace(QO, SVector(sel))
 
@@ -758,24 +803,31 @@ function permute(
     order_svector = SVector{length(order),Int}(order) # convert it to SVector for performance
 
     # obtain the arguments: dims for reshape; perm for PermutedDimsArray
-    dims, perm = _dims_and_perm(A.type, A.dims, order_svector, length(order_svector))
+    dimslist = dims_to_list(A.dims)
+    dims, perm = _dims_and_perm(A.type, dimslist, order_svector, length(order_svector))
 
     return QuantumObject(
         reshape(permutedims(reshape(A.data, dims...), Tuple(perm)), size(A)),
         A.type,
-        A.dims[order_svector],
+        Dimensions(A.dims.to[order_svector]),
     )
 end
 
-function _dims_and_perm(
+_dims_and_perm(
     ::ObjType,
-    dims::SVector{N,Int},
+    dimslist::SVector{N,Int},
     order::AbstractVector{Int},
     L::Int,
-) where {ObjType<:Union{KetQuantumObject,BraQuantumObject},N}
-    return reverse(dims), reverse((L + 1) .- order)
-end
+) where {ObjType<:Union{KetQuantumObject,BraQuantumObject},N} = reverse(dimslist), reverse((L + 1) .- order)
 
-function _dims_and_perm(::OperatorQuantumObject, dims::SVector{N,Int}, order::AbstractVector{Int}, L::Int) where {N}
-    return reverse(vcat(dims, dims)), reverse((2 * L + 1) .- vcat(order, order .+ L))
-end
+# if dims originates from Dimensions
+_dims_and_perm(::OperatorQuantumObject, dimslist::SVector{N,Int}, order::AbstractVector{Int}, L::Int) where {N} =
+    reverse(vcat(dimslist, dimslist)), reverse((2 * L + 1) .- vcat(order, order .+ L))
+
+# if dims originates from CompoundDimensions
+_dims_and_perm(
+    ::OperatorQuantumObject,
+    dimslist::SVector{2,SVector{N,Int}},
+    order::AbstractVector{Int},
+    L::Int,
+) where {N} = reverse(vcat(dimslist[1], dimslist[2])), reverse((2 * L + 1) .- vcat(order, order .+ L))

src/qobj/dimensions.jl

@@ -0,0 +1,78 @@
+export AbstractDimensions, Dimensions, CompoundDimensions
+
+abstract type AbstractDimensions{N} end
+
+# this show function is for printing AbstractDimensions
+function Base.show(io::IO, svec::SVector{N,AbstractSpace}) where {N}
+    print(io, "[")
+    join(io, string.(svec), ", ")
+    return print(io, "]")
+end
+
+struct Dimensions{N} <: AbstractDimensions{N}
+    to::SVector{N,AbstractSpace}
+end
+function Dimensions(dims::Union{AbstractVector{T},NTuple{N,T}}) where {T<:Integer,N}
+    _non_static_array_warning("dims", dims)
+    L = length(dims)
+    (L > 0) || throw(DomainError(dims, "The argument dims must be of non-zero length"))
+
+    return Dimensions{L}(SVector{L,AbstractSpace}(Space.(dims)))
+end
+Dimensions(dims::Int) = Dimensions(SVector{1,Int}(dims))
+Dimensions(dims::Any) = throw(
+    ArgumentError(
+        "The argument dims must be a Tuple or a StaticVector of non-zero length and contain only positive integers.",
+    ),
+)
+
+Base.show(io::IO, D::Dimensions) = print(io, D.to)
+
+struct CompoundDimensions{N} <: AbstractDimensions{N}
+    # note that the number `N` should be the same for both `to` and `from`
+    to::SVector{N,AbstractSpace}   # space acting on the left
+    from::SVector{N,AbstractSpace} # space acting on the right
+end
+function CompoundDimensions(dims::Union{AbstractVector{T},NTuple{N,T}}) where {T<:Union{AbstractVector,NTuple},N}
+    (length(dims) != 2) && throw(ArgumentError("Invalid dims = $dims"))
+
+    _non_static_array_warning("dims[1]", dims[1])
+    _non_static_array_warning("dims[2]", dims[2])
+
+    L1 = length(dims[1])
+    L2 = length(dims[2])
+    ((L1 > 0) && (L1 == L2)) || throw(
+        DomainError(
+            (L1, L2),
+            "The length of the arguments `dims[1]` and `dims[2]` must be in the same length and have at least one element.",
+        ),
+    )
+
+    return CompoundDimensions{L1}(
+        SVector{L1,AbstractSpace}(Space.(dims[1])),
+        SVector{L1,AbstractSpace}(Space.(dims[2])),
+    )
+end
+
+Base.show(io::IO, D::CompoundDimensions) = print(io, "[", D.to, ", ", D.from, "]")
+
+_gen_dims(dims::AbstractDimensions) = dims
+_gen_dims(dims::Union{AbstractVector{T},NTuple{N,T}}) where {T<:Integer,N} = Dimensions(dims)
+_gen_dims(dims::Union{AbstractVector{T},NTuple{N,T}}) where {T<:Union{AbstractVector,NTuple},N} =
+    CompoundDimensions(dims)
+_gen_dims(dims::Any) = Dimensions(dims)
+
+# obtain dims in the type of SVector with integers
+dims_to_list(dimsvec::SVector{N,AbstractSpace}) where {N} = SVector{N,Int}(ntuple(i -> dimsvec[i].size, Val(N)))
+dims_to_list(dims::Dimensions) = dims_to_list(dims.to)
+dims_to_list(dims::CompoundDimensions) = SVector{2}(dims_to_list(dims.to), dims_to_list(dims.from))
+
+Base.:(==)(vect::AbstractVector{T}, dims::AbstractDimensions) where {T} = vect == dims_to_list(dims)
+Base.:(==)(dims::AbstractDimensions, vect::AbstractVector{T}) where {T} = vect == dims
+
+Base.length(::AbstractDimensions{N}) where {N} = N
+
+Base.prod(dims::Dimensions) = prod(dims.to)
+
+LinearAlgebra.transpose(dims::Dimensions) = dims
+LinearAlgebra.transpose(dims::CompoundDimensions) = CompoundDimensions(dims.from, dims.to) # switch `to` and `from`

src/qobj/eigsolve.jl

@@ -7,11 +7,11 @@ export eigenenergies, eigenstates, eigsolve
 export eigsolve_al
 
 @doc raw"""
-    struct EigsolveResult{T1<:Vector{<:Number}, T2<:AbstractMatrix{<:Number}, ObjType<:Union{Nothing,OperatorQuantumObject,SuperOperatorQuantumObject},N}
+    struct EigsolveResult{T1<:Vector{<:Number}, T2<:AbstractMatrix{<:Number}, ObjType<:Union{Nothing,OperatorQuantumObject,SuperOperatorQuantumObject},DimType<:AbstractDimensions}
         values::T1
         vectors::T2
         type::ObjType
-        dims::SVector{N,Int}
+        dims::DimType
         iter::Int
         numops::Int
         converged::Bool
@@ -23,7 +23,7 @@ A struct containing the eigenvalues, the eigenvectors, and some information from
 - `values::AbstractVector`: the eigenvalues
 - `vectors::AbstractMatrix`: the transformation matrix (eigenvectors)
 - `type::Union{Nothing,QuantumObjectType}`: the type of [`QuantumObject`](@ref), or `nothing` means solving eigen equation for general matrix
-- `dims::SVector`: the `dims` of [`QuantumObject`](@ref)
+- `dims::AbstractDimensions`: the `dims` of [`QuantumObject`](@ref)
 - `iter::Int`: the number of iteration during the solving process
 - `numops::Int` : number of times the linear map was applied in krylov methods
 - `converged::Bool`: Whether the result is converged
@@ -70,12 +70,12 @@ struct EigsolveResult{
     T1<:Vector{<:Number},
     T2<:AbstractMatrix{<:Number},
     ObjType<:Union{Nothing,OperatorQuantumObject,SuperOperatorQuantumObject},
-    N,
+    DimType<:AbstractDimensions,
 }
     values::T1
     vectors::T2
     type::ObjType
-    dims::SVector{N,Int}
+    dims::DimType
     iter::Int
     numops::Int
     converged::Bool
@@ -159,7 +159,7 @@ function _eigsolve(
     A,
     b::AbstractVector{T},
     type::ObjType,
-    dims::SVector,
+    dims::AbstractDimensions,
     k::Int = 1,
     m::Int = max(20, 2 * k + 1);
     tol::Real = 1e-8,
@@ -296,7 +296,7 @@ function eigsolve(
     A;
     v0::Union{Nothing,AbstractVector} = nothing,
     type::Union{Nothing,OperatorQuantumObject,SuperOperatorQuantumObject} = nothing,
-    dims = SVector{0,Int}(),
+    dims = Dimensions{0}(SVector{0,AbstractSpace}()),
     sigma::Union{Nothing,Real} = nothing,
     k::Int = 1,
     krylovdim::Int = max(20, 2 * k + 1),
@@ -309,8 +309,6 @@ function eigsolve(
     isH = ishermitian(A)
     v0 === nothing && (v0 = normalize!(rand(T, size(A, 1))))
 
-    dims = SVector(dims)
-
     if sigma === nothing
         res = _eigsolve(A, v0, type, dims, k, krylovdim, tol = tol, maxiter = maxiter)
         vals = res.values