Change the type of dims for QuantumObject to SVector for type stabilities (qutip#217)

* Improve c_ops handling

* Format code

* Change dims to SVector

* Fix JET errors

* Fix dfd_mesolve issues

* Fix entanglement issues

* Fix low rank dynamics

* Fix all tests

* Format code

* Change to StaticArraysCore.jl and minor changes

* Make a few comments on the documentation

* Update some docstring

* Minor changes

* FIx error on Documentation

Author: albertomercurio

Project.toml

@@ -1,7 +1,7 @@
 name = "QuantumToolbox"
 uuid = "6c2fb7c5-b903-41d2-bc5e-5a7c320b9fab"
 authors = ["Alberto Mercurio", "Luca Gravina", "Yi-Te Huang"]
-version = "0.12.1"
+version = "0.13.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -18,6 +18,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -41,6 +42,7 @@ Random = "<0.0.1, 1"
 Reexport = "1"
 SparseArrays = "<0.0.1, 1"
 SpecialFunctions = "2"
+StaticArraysCore = "1"
 Test = "<0.0.1, 1"
 julia = "1.7"
 

docs/make.jl

@@ -28,6 +28,7 @@ const PAGES = [
             "users_guide/QuantumObject/QuantumObject.md",
             "users_guide/QuantumObject/QuantumObject_functions.md",
         ],
+        "The Importance of Type-Stability" => "users_guide/type_stability.md",
         "Manipulating States and Operators" => "users_guide/states_and_operators.md",
         "Tensor Products and Partial Traces" => "users_guide/tensor.md",
         "Time Evolution and Dynamics" => [

docs/src/users_guide/QuantumObject/QuantumObject.md

@@ -39,9 +39,15 @@ Qobj(rand(4, 4))
 ```
 
 ```@example Qobj
-Qobj(rand(4, 4), dims = [2, 2])
+M = rand(ComplexF64, 4, 4)
+Qobj(M, dims = [2, 2])  # dims as Vector
+Qobj(M, dims = (2, 2))  # dims as Tuple (recommended)
+Qobj(M, dims = SVector(2, 2)) # dims as StaticArrays.SVector (recommended)
 ```
 
+> [!IMPORTANT]
+> Please note that here we put the `dims` as a tuple `(2, 2)`. Although it supports also `Vector` type (`dims = [2, 2]`), it is recommended to use `Tuple` or `SVector` from [`StaticArrays.jl`](https://github.com/JuliaArrays/StaticArrays.jl) to improve performance. For a brief explanation on the impact of the type of `dims`, see the Section [The Importance of Type-Stability](@ref doc:Type-Stability).
+
 ```@example Qobj
 Qobj(rand(4, 4), type = SuperOperator)
 ```

docs/src/users_guide/type_stability.md

@@ -0,0 +1,3 @@
+# [The Importance of Type-Stability](@id doc:Type-Stability)
+
+This page is still under construction.

src/QuantumToolbox.jl

@@ -32,6 +32,8 @@ import OrdinaryDiffEq: OrdinaryDiffEqAlgorithm
 import Pkg
 import Random
 import SpecialFunctions: loggamma
+@reexport import StaticArraysCore: SVector
+import StaticArraysCore: MVector
 
 # Setting the number of threads to 1 allows
 # to achieve better performances for more massive parallelizations

src/metrics.jl

@@ -62,21 +62,21 @@ function entropy_vn(
 end
 
 """
-    entanglement(QO::QuantumObject, sel::Vector)
+    entanglement(QO::QuantumObject, sel::Union{Int,AbstractVector{Int},Tuple})
 
 Calculates the entanglement by doing the partial trace of `QO`, selecting only the dimensions
 with the indices contained in the `sel` vector, and then using the Von Neumann entropy [`entropy_vn`](@ref).
 """
 function entanglement(
     QO::QuantumObject{<:AbstractArray{T},OpType},
-    sel::Vector{Int},
+    sel::Union{AbstractVector{Int},Tuple},
 ) where {T,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
     ψ = normalize(QO)
     ρ_tr = ptrace(ψ, sel)
     entropy = entropy_vn(ρ_tr)
     return (entropy > 0) * entropy
 end
-entanglement(QO::QuantumObject, sel::Int) = entanglement(QO, [sel])
+entanglement(QO::QuantumObject, sel::Int) = entanglement(QO, (sel,))
 
 @doc raw"""
     tracedist(ρ::QuantumObject, σ::QuantumObject)

src/qobj/arithmetic_and_attributes.jl

@@ -471,7 +471,7 @@ proj(ψ::QuantumObject{<:AbstractArray{T},KetQuantumObject}) where {T} = ψ * ψ
 proj(ψ::QuantumObject{<:AbstractArray{T},BraQuantumObject}) where {T} = ψ' * ψ
 
 @doc raw"""
-    ptrace(QO::QuantumObject, sel::Vector{Int})
+    ptrace(QO::QuantumObject, sel)
 
 [Partial trace](https://en.wikipedia.org/wiki/Partial_trace) of a quantum state `QO` leaving only the dimensions
 with the indices present in the `sel` vector.
@@ -487,7 +487,7 @@ Quantum Object:   type=Ket   dims=[2, 2]   size=(4,)
  0.0 + 0.0im
  0.0 + 0.0im
 
-julia> ptrace(ψ, [2])
+julia> ptrace(ψ, 2)
 Quantum Object:   type=Operator   dims=[2]   size=(2, 2)   ishermitian=true
 2×2 Matrix{ComplexF64}:
  0.0+0.0im  0.0+0.0im
@@ -504,63 +504,79 @@ Quantum Object:   type=Ket   dims=[2, 2]   size=(4,)
                 0.0 + 0.0im
  0.7071067811865475 + 0.0im
 
-julia> ptrace(ψ, [1])
+julia> ptrace(ψ, 1)
 Quantum Object:   type=Operator   dims=[2]   size=(2, 2)   ishermitian=true
 2×2 Matrix{ComplexF64}:
  0.5+0.0im  0.0+0.0im
  0.0+0.0im  0.5+0.0im
 ```
 """
-function ptrace(QO::QuantumObject{<:AbstractArray{T1},KetQuantumObject}, sel::Vector{T2}) where {T1,T2<:Integer}
+function ptrace(QO::QuantumObject{<:AbstractArray,KetQuantumObject}, sel::Union{AbstractVector{Int},Tuple})
     length(QO.dims) == 1 && return QO
 
-    ρtr, dkeep = _ptrace_ket(QO.data, QO.dims, sel)
+    ρtr, dkeep = _ptrace_ket(QO.data, QO.dims, SVector(sel))
     return QuantumObject(ρtr, dims = dkeep)
 end
 
-ptrace(QO::QuantumObject{<:AbstractArray{T1},BraQuantumObject}, sel::Vector{T2}) where {T1,T2<:Integer} =
-    ptrace(QO', sel)
+ptrace(QO::QuantumObject{<:AbstractArray,BraQuantumObject}, sel::Union{AbstractVector{Int},Tuple}) = ptrace(QO', sel)
 
-function ptrace(QO::QuantumObject{<:AbstractArray{T1},OperatorQuantumObject}, sel::Vector{T2}) where {T1,T2<:Integer}
+function ptrace(QO::QuantumObject{<:AbstractArray,OperatorQuantumObject}, sel::Union{AbstractVector{Int},Tuple})
     length(QO.dims) == 1 && return QO
 
-    ρtr, dkeep = _ptrace_oper(QO.data, QO.dims, sel)
+    ρtr, dkeep = _ptrace_oper(QO.data, QO.dims, SVector(sel))
     return QuantumObject(ρtr, dims = dkeep)
 end
-ptrace(QO::QuantumObject, sel::Int) = ptrace(QO, [sel])
-
-function _ptrace_ket(QO::AbstractArray{T1}, dims::Vector{<:Integer}, sel::Vector{T2}) where {T1,T2<:Integer}
-    rd = dims
-    nd = length(rd)
-
-    nd == 1 && return QO, rd
-
-    dkeep = rd[sel]
-    qtrace = filter!(e -> e ∉ sel, Vector(1:nd))
-    dtrace = @view(rd[qtrace])
+ptrace(QO::QuantumObject, sel::Int) = ptrace(QO, SVector(sel))
+
+function _ptrace_ket(QO::AbstractArray, dims::Union{SVector,MVector}, sel)
+    nd = length(dims)
+
+    nd == 1 && return QO, dims
+
+    qtrace = filter(i -> i ∉ sel, 1:nd)
+    dkeep = dims[sel]
+    dtrace = dims[qtrace]
+    # Concatenate sel and qtrace without loosing the length information
+    sel_qtrace = ntuple(Val(nd)) do i
+        if i <= length(sel)
+            @inbounds sel[i]
+        else
+            @inbounds qtrace[i-length(sel)]
+        end
+    end
 
-    vmat = reshape(QO, reverse(rd)...)
-    topermute = nd + 1 .- vcat(sel, qtrace)
+    vmat = reshape(QO, reverse(dims)...)
+    topermute = nd + 1 .- sel_qtrace
     vmat = PermutedDimsArray(vmat, topermute)
     vmat = reshape(vmat, prod(dkeep), prod(dtrace))
 
     return vmat * vmat', dkeep
 end
 
-function _ptrace_oper(QO::AbstractArray{T1}, dims::Vector{<:Integer}, sel::Vector{T2}) where {T1,T2<:Integer}
-    rd = dims
-    nd = length(rd)
-
-    nd == 1 && return QO, rd
-
-    dkeep = rd[sel]
-    qtrace = filter!(e -> e ∉ sel, Vector(1:nd))
-    dtrace = @view(rd[qtrace])
+function _ptrace_oper(QO::AbstractArray, dims::Union{SVector,MVector}, sel)
+    nd = length(dims)
+
+    nd == 1 && return QO, dims
+
+    qtrace = filter(i -> i ∉ sel, 1:nd)
+    dkeep = dims[sel]
+    dtrace = dims[qtrace]
+    # Concatenate sel and qtrace without loosing the length information
+    qtrace_sel = ntuple(Val(2 * nd)) do i
+        if i <= length(qtrace)
+            @inbounds qtrace[i]
+        elseif i <= 2 * length(qtrace)
+            @inbounds qtrace[i-length(qtrace)] + nd
+        elseif i <= 2 * length(qtrace) + length(sel)
+            @inbounds sel[i-length(qtrace)-length(sel)]
+        else
+            @inbounds sel[i-length(qtrace)-2*length(sel)] + nd
+        end
+    end
 
-    ρmat = reshape(QO, reverse!(repeat(rd, 2))...)
-    topermute = 2 * nd + 1 .- vcat(qtrace, qtrace .+ nd, sel, sel .+ nd)
-    reverse!(topermute)
-    ρmat = PermutedDimsArray(ρmat, topermute)
+    ρmat = reshape(QO, reverse(vcat(dims, dims))...)
+    topermute = 2 * nd + 1 .- qtrace_sel
+    ρmat = PermutedDimsArray(ρmat, reverse(topermute))
 
     ## TODO: Check if it works always
 
@@ -639,7 +655,7 @@ function get_coherence(ρ::QuantumObject{<:AbstractArray,OperatorQuantumObject})
 end
 
 @doc raw"""
-    permute(A::QuantumObject, order::Vector{Int})
+    permute(A::QuantumObject, order::Union{AbstractVector{Int},Tuple})
 
 Permute the tensor structure of a [`QuantumObject`](@ref) `A` according to the specified `order` list
 
@@ -660,28 +676,36 @@ true
 """
 function permute(
     A::QuantumObject{<:AbstractArray{T},ObjType},
-    order::AbstractVector{Int},
+    order::Union{AbstractVector{Int},Tuple},
 ) where {T,ObjType<:Union{KetQuantumObject,BraQuantumObject,OperatorQuantumObject}}
     (length(order) != length(A.dims)) &&
         throw(ArgumentError("The order list must have the same length as the number of subsystems (A.dims)"))
 
     !isperm(order) && throw(ArgumentError("$(order) is not a valid permutation of the subsystems (A.dims)"))
 
+    _non_static_array_warning("order", order)
+
+    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, length(order))
+    dims, perm = _dims_and_perm(A.type, A.dims, order_svector, length(order_svector))
 
-    return QuantumObject(reshape(PermutedDimsArray(reshape(A.data, dims...), perm), size(A)), A.type, A.dims[order])
+    return QuantumObject(
+        reshape(PermutedDimsArray(reshape(A.data, dims...), Tuple(perm)), size(A)),
+        A.type,
+        A.dims[order_svector],
+    )
 end
 
 function _dims_and_perm(
     ::ObjType,
-    dims::AbstractVector{Int},
+    dims::SVector{N,Int},
     order::AbstractVector{Int},
     L::Int,
-) where {ObjType<:Union{KetQuantumObject,BraQuantumObject}}
-    return reverse(dims), reverse!((L + 1) .- order)
+) where {ObjType<:Union{KetQuantumObject,BraQuantumObject},N}
+    return reverse(dims), reverse((L + 1) .- order)
 end
 
-function _dims_and_perm(::OperatorQuantumObject, dims::AbstractVector{Int}, order::AbstractVector{Int}, L::Int)
-    return reverse!([dims; dims]), reverse!((2 * L + 1) .- [order; order .+ L])
+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

src/qobj/eigsolve.jl

@@ -6,11 +6,11 @@ export EigsolveResult
 export eigenenergies, eigenstates, eigsolve, eigsolve_al
 
 @doc raw"""
-    struct EigsolveResult{T1<:Vector{<:Number}, T2<:AbstractMatrix{<:Number}, ObjType<:Union{Nothing,OperatorQuantumObject,SuperOperatorQuantumObject}}
+    struct EigsolveResult{T1<:Vector{<:Number}, T2<:AbstractMatrix{<:Number}, ObjType<:Union{Nothing,OperatorQuantumObject,SuperOperatorQuantumObject},N}
         values::T1
         vectors::T2
         type::ObjType
-        dims::Vector{Int}
+        dims::SVector{N,Int}
         iter::Int
         numops::Int
         converged::Bool
@@ -22,7 +22,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::Vector{Int}`: the `dims` of [`QuantumObject`](@ref)
+- `dims::SVector`: 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
@@ -54,11 +54,12 @@ struct EigsolveResult{
     T1<:Vector{<:Number},
     T2<:AbstractMatrix{<:Number},
     ObjType<:Union{Nothing,OperatorQuantumObject,SuperOperatorQuantumObject},
+    N,
 }
     values::T1
     vectors::T2
     type::ObjType
-    dims::Vector{Int}
+    dims::SVector{N,Int}
     iter::Int
     numops::Int
     converged::Bool
@@ -271,7 +272,7 @@ function _eigsolve(
     A,
     b::AbstractVector{T},
     type::ObjType,
-    dims::Vector{Int},
+    dims::SVector,
     k::Int = 1,
     m::Int = max(20, 2 * k + 1);
     tol::Real = 1e-8,
@@ -398,7 +399,7 @@ function eigsolve(
     A;
     v0::Union{Nothing,AbstractVector} = nothing,
     type::Union{Nothing,OperatorQuantumObject,SuperOperatorQuantumObject} = nothing,
-    dims::Vector{Int} = Int[],
+    dims = SVector{0,Int}(),
     sigma::Union{Nothing,Real} = nothing,
     k::Int = 1,
     krylovdim::Int = max(20, 2 * k + 1),
@@ -411,6 +412,8 @@ 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