Support for new Dimensions structure in QuantumToolbox (#136)

Author: ytdHuang

Project.toml

@@ -1,7 +1,7 @@
 name = "HierarchicalEOM"
 uuid = "a62dbcb7-80f5-4d31-9a88-8b19fd92b128"
-authors = ["Yi-Te Huang <[email protected]>"]
-version = "2.3.3"
+authors = ["Yi-Te Huang"]
+version = "2.4.0"
 
 [deps]
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
@@ -37,7 +37,7 @@ LinearSolve = "2.4.2 - 2"
 OrdinaryDiffEqCore = "1"
 OrdinaryDiffEqLowOrderRK = "1"
 Pkg = "1"
-QuantumToolbox = "0.22 - 0.24"
+QuantumToolbox = "0.25"
 Reexport = "1"
 SciMLBase = "2"
 SciMLOperators = "0.3"

docs/src/ADOs.md

@@ -21,7 +21,7 @@ which is usually obtained after solving the [time evolution](@ref doc-Time-Evolu
 ## Fields
 The fields of the structure [`ADOs`](@ref) are as follows:
  - `data` : the vectorized auxiliary density operators
- - `dims` : the dimension list of the coupling operator (should be equal to the system dims).
+ - `dimensions` : the dimension list of the coupling operator (should be equal to the system dimensions).
  - `N` : the number of auxiliary density operators
 - `parity`: the [parity](@ref doc-Parity) label
 
@@ -31,6 +31,7 @@ One obtain the value of each fields as follows:
 ados::ADOs
 
 ados.data
+ados.dimensions
 ados.dims
 ados.N
 ados.parity

docs/src/heom_matrix/M_Boson.md

@@ -31,7 +31,7 @@ M_odd  = M_Boson(Hs, tier, Bath, ODD)
 The fields of the structure [`M_Boson`](@ref) are as follows:
  - `data` : the sparse matrix of HEOM Liouvillian superoperator
  - `tier` : the tier (cutoff level) for the bosonic hierarchy
- - `dims` : the dimension list of the coupling operator (should be equal to the system dims).
+ - `dimensions` : the dimension list of the coupling operator (should be equal to the system dimensions).
  - `N` : the number of total [ADOs](@ref doc-ADOs)
  - `sup_dim` : the dimension of system superoperator
  - `parity` : the [parity](@ref doc-Parity) label of the operator which HEOMLS is acting on.
@@ -44,6 +44,7 @@ M::M_Boson
 
 M.data
 M.tier
+M.dimensions
 M.dims
 M.N
 M.sup_dim

docs/src/heom_matrix/M_Boson_Fermion.md

@@ -36,7 +36,7 @@ The fields of the structure [`M_Boson_Fermion`](@ref) are as follows:
  - `data` : the sparse matrix of HEOM Liouvillian superoperator
  - `Btier` : the tier (cutoff level) for bosonic hierarchy
  - `Ftier` : the tier (cutoff level) for fermionic hierarchy
- - `dims` : the dimension list of the coupling operator (should be equal to the system dims).
+ - `dimensions` : the dimension list of the coupling operator (should be equal to the system dimensions).
  - `N` : the number of total [ADOs](@ref doc-ADOs)
  - `sup_dim` : the dimension of system superoperator
  - `parity` : the [parity](@ref doc-Parity) label of the operator which HEOMLS is acting on. 
@@ -51,6 +51,7 @@ M::M_Boson_Fermion
 M.data
 M.Btier
 M.Ftier
+M.dimensions
 M.dims
 M.N
 M.sup_dim

docs/src/heom_matrix/M_Fermion.md

@@ -31,7 +31,7 @@ M_odd  = M_Fermion(Hs, tier, Bath, ODD)
 The fields of the structure [`M_Fermion`](@ref) are as follows:
  - `data` : the sparse matrix of HEOM Liouvillian superoperator
  - `tier` : the tier (cutoff level) for the fermionic hierarchy
- - `dims` : the dimension list of the coupling operator (should be equal to the system dims).
+ - `dimensions` : the dimension list of the coupling operator (should be equal to the system dimensions).
  - `N` : the number of total [ADOs](@ref doc-ADOs)
  - `sup_dim` : the dimension of system superoperator
  - `parity` : the [parity](@ref doc-Parity) label of the operator which HEOMLS is acting on.
@@ -44,6 +44,7 @@ M::M_Fermion
 
 M.data
 M.tier
+M.dimensions
 M.dims
 M.N
 M.sup_dim

docs/src/heom_matrix/schrodinger_eq.md

@@ -31,7 +31,7 @@ M_odd  = M_S(Hs, ODD)
 The fields of the structure [`M_S`](@ref) are as follows:
  - `data` : the sparse matrix of HEOM Liouvillian superoperator
  - `tier` : the tier (cutoff level) for the hierarchy, which equals to `0` in this case
- - `dims` : the dimension list of the coupling operator (should be equal to the system dims).
+ - `dimensions` : the dimension list of the coupling operator (should be equal to the system dimensions).
  - `N` : the number of total [ADOs](@ref doc-ADOs), which equals to `1` (only the reduced density operator) in this case
  - `sup_dim` : the dimension of system superoperator
  - `parity::AbstractParity` : the [parity](@ref doc-Parity) label of the operator which HEOMLS is acting on.
@@ -42,6 +42,7 @@ M::M_S
 
 M.data
 M.tier
+M.dimensions
 M.dims
 M.N
 M.sup_dim

docs/src/libraryAPI.md

@@ -62,7 +62,7 @@ ODD
 HEOMSuperOp
 HEOMSuperOp(op, opParity::AbstractParity, refHEOMLS::AbstractHEOMLSMatrix)
 HEOMSuperOp(op, opParity::AbstractParity, refADOs::ADOs)
-HEOMSuperOp(op, opParity::AbstractParity, dims::SVector, N::Int)
+HEOMSuperOp(op, opParity::AbstractParity, dims, N::Int)
 AbstractHEOMLSMatrix
 M_S
 M_S(Hsys::QuantumObject, parity::AbstractParity=EVEN; verbose::Bool=true)

ext/HierarchicalEOM_CUDAExt.jl

@@ -22,17 +22,28 @@ Return a new HEOMLS-matrix-type object with `M.data` is in the type of `CuSparse
 function CuSparseMatrixCSC(M::T) where {T<:AbstractHEOMLSMatrix}
     A_gpu = _convert_to_gpu_matrix(M.data)
     if T <: M_S
-        return M_S(A_gpu, M.tier, M.dims, M.N, M.sup_dim, M.parity)
+        return M_S(A_gpu, M.tier, M.dimensions, M.N, M.sup_dim, M.parity)
     elseif T <: M_Boson
-        return M_Boson(A_gpu, M.tier, M.dims, M.N, M.sup_dim, M.parity, M.bath, M.hierarchy)
+        return M_Boson(A_gpu, M.tier, M.dimensions, M.N, M.sup_dim, M.parity, M.bath, M.hierarchy)
     elseif T <: M_Fermion
-        return M_Fermion(A_gpu, M.tier, M.dims, M.N, M.sup_dim, M.parity, M.bath, M.hierarchy)
+        return M_Fermion(A_gpu, M.tier, M.dimensions, M.N, M.sup_dim, M.parity, M.bath, M.hierarchy)
     else
-        return M_Boson_Fermion(A_gpu, M.Btier, M.Ftier, M.dims, M.N, M.sup_dim, M.parity, M.Bbath, M.Fbath, M.hierarchy)
+        return M_Boson_Fermion(
+            A_gpu,
+            M.Btier,
+            M.Ftier,
+            M.dimensions,
+            M.N,
+            M.sup_dim,
+            M.parity,
+            M.Bbath,
+            M.Fbath,
+            M.hierarchy,
+        )
     end
 end
 
-CuSparseMatrixCSC{ComplexF32}(M::HEOMSuperOp) = HEOMSuperOp(_convert_to_gpu_matrix(M.data), M.dims, M.N, M.parity)
+CuSparseMatrixCSC{ComplexF32}(M::HEOMSuperOp) = HEOMSuperOp(_convert_to_gpu_matrix(M.data), M.dimensions, M.N, M.parity)
 
 function _convert_to_gpu_matrix(A::AbstractSparseMatrix)
     if A isa CuSparseMatrixCSC{ComplexF32,Int32}
@@ -53,7 +64,7 @@ _convert_to_gpu_matrix(A::MatrixOperator) = MatrixOperator(_convert_to_gpu_matri
 _convert_to_gpu_matrix(A::ScaledOperator) = ScaledOperator(A.λ, _convert_to_gpu_matrix(A.L))
 _convert_to_gpu_matrix(A::AddedOperator) = AddedOperator(map(op -> _convert_to_gpu_matrix(op), A.ops))
 
-_Tr(M::Type{<:CuSparseMatrixCSC}, dims::SVector, N::Int) = CuSparseVector(_Tr(eltype(M), dims, N))
+_Tr(M::Type{<:CuSparseMatrixCSC}, dimensions::Dimensions, N::Int) = CuSparseVector(_Tr(eltype(M), dimensions, N))
 
 # change the type of `ADOs` to match the type of HEOMLS matrix
 _HandleVectorType(M::Type{<:CuSparseMatrixCSC}, V::SparseVector) = CuArray{_CType(eltype(M))}(V)

src/ADOs.jl

@@ -7,10 +7,13 @@ The Auxiliary Density Operators for HEOM model.
 
 # Fields
 - `data` : the vectorized auxiliary density operators
-- `dims` : the dimension list of the coupling operator (should be equal to the system dims).
+- `dimensions` : the dimension list of the coupling operator (should be equal to the system dimensions).
 - `N` : the number of auxiliary density operators
 - `parity`: the parity label (`EVEN` or `ODD`).
 
+!!! note "`dims` property"
+    For a given `ados::ADOs`, `ados.dims` or `getproperty(ados, :dims)` returns its `dimensions` in the type of integer-vector.
+
 # Methods
 One can obtain the density matrix for specific index (`idx`) by calling : `ados[idx]`.
 `HierarchicalEOM.jl` also supports the following calls (methods) :
@@ -26,15 +29,17 @@ end
 """
 struct ADOs
     data::SparseVector{ComplexF64,Int64}
-    dims::SVector
+    dimensions::Dimensions
     N::Int
     parity::AbstractParity
-end
 
-# these functions are for forward compatibility
-ADOs(data::SparseVector{ComplexF64,Int64}, dim::Int, N::Int, parity::AbstractParity) = ADOs(data, [dim], N, parity)
-ADOs(data::SparseVector{ComplexF64,Int64}, dims::AbstractVector, N::Int, parity::AbstractParity) =
-    ADOs(data, SVector{length(dims),Int}(dims), N, parity)
+    function ADOs(data::AbstractVector, dims, N::Int, parity::AbstractParity)
+        dimensions = _gen_dimensions(dims)
+        Vsize = size(data, 1)
+        ((Vsize / N) == prod(dimensions)^2) || error("The `dimensions` is not consistent with the ADOs number `N`.")
+        return new(sparsevec(data), dimensions, N, parity)
+    end
+end
 
 @doc raw"""
     ADOs(V, N, parity)
@@ -45,16 +50,7 @@ Generate the object of auxiliary density operators for HEOM model.
 - `N::Int` : the number of auxiliary density operators.
 - `parity::AbstractParity` : the parity label (`EVEN` or `ODD`). Default to `EVEN`.
 """
-function ADOs(V::AbstractVector, N::Int, parity::AbstractParity = EVEN)
-    # check the dimension of V
-    d = size(V, 1)
-    dim = √(d / N)
-    if isinteger(dim)
-        return ADOs(sparsevec(V), SVector{1,Int}(Int(dim)), N, parity)
-    else
-        error("The dimension of vector is not consistent with the ADOs number \"N\".")
-    end
-end
+ADOs(V::AbstractVector, N::Int, parity::AbstractParity = EVEN) = ADOs(V, isqrt(Int(size(V, 1) / N)), N, parity)
 
 @doc raw"""
     ADOs(ρ, N, parity)
@@ -67,11 +63,20 @@ Generate the object of auxiliary density operators for HEOM model.
 """
 function ADOs(ρ::QuantumObject, N::Int = 1, parity::AbstractParity = EVEN)
     _ρ = sparsevec(ket2dm(ρ).data)
-    return ADOs(sparsevec(_ρ.nzind, _ρ.nzval, N * length(_ρ)), ρ.dims, N, parity)
+    return ADOs(sparsevec(_ρ.nzind, _ρ.nzval, N * length(_ρ)), ρ.dimensions, N, parity)
 end
 ADOs(ρ, N::Int = 1, parity::AbstractParity = EVEN) =
     error("HierarchicalEOM doesn't support input `ρ` with type : $(typeof(ρ))")
 
+function Base.getproperty(ados::ADOs, key::Symbol)
+    # a comment here to avoid bad render by JuliaFormatter
+    if key === :dims
+        return dimensions_to_dims(getfield(ados, :dimensions))
+    else
+        return getfield(ados, key)
+    end
+end
+
 Base.checkbounds(A::ADOs, i::Int) =
     ((i > A.N) || (i < 1)) ? error("Attempt to access $(A.N)-element ADOs at index [$(i)]") : nothing
 
@@ -92,31 +97,33 @@ Base.lastindex(A::ADOs) = length(A)
 function Base.getindex(A::ADOs, i::Int)
     checkbounds(A, i)
 
-    D = prod(A.dims)
+    D = prod(A.dimensions)
     sup_dim = D^2
     back = sup_dim * i
-    return QuantumObject(reshape(A.data[(back-sup_dim+1):back], D, D), Operator, A.dims)
+    return QuantumObject(reshape(A.data[(back-sup_dim+1):back], D, D), Operator, A.dimensions)
 end
 
 function Base.getindex(A::ADOs, r::UnitRange{Int})
     checkbounds(A, r[1])
     checkbounds(A, r[end])
 
     result = []
-    D = prod(A.dims)
+    D = prod(A.dimensions)
     sup_dim = D^2
     for i in r
         back = sup_dim * i
-        push!(result, QuantumObject(reshape(A.data[(back-sup_dim+1):back], D, D), Operator, A.dims))
+        push!(result, QuantumObject(reshape(A.data[(back-sup_dim+1):back], D, D), Operator, A.dimensions))
     end
     return result
 end
 Base.getindex(A::ADOs, ::Colon) = getindex(A, 1:lastindex(A))
 
 Base.iterate(A::ADOs, state::Int = 1) = state > length(A) ? nothing : (A[state], state + 1)
 
-Base.show(io::IO, A::ADOs) =
-    print(io, "$(A.N) Auxiliary Density Operators with $(A.parity) and (system) dims = $(A.dims)\n")
+Base.show(io::IO, A::ADOs) = print(
+    io,
+    "$(A.N) Auxiliary Density Operators with $(A.parity) and (system) dims = $(_get_dims_string(A.dimensions))\n",
+)
 Base.show(io::IO, m::MIME"text/plain", A::ADOs) = show(io, A)
 
 @doc raw"""
@@ -130,8 +137,8 @@ Return the density matrix of the reduced state (system) from a given auxiliary d
 - `ρ::QuantumObject` : The density matrix of the reduced state
 """
 function getRho(ados::ADOs)
-    D = prod(ados.dims)
-    return QuantumObject(reshape(ados.data[1:(D^2)], D, D), Operator, ados.dims)
+    D = prod(ados.dimensions)
+    return QuantumObject(reshape(ados.data[1:(D^2)], D, D), Operator, ados.dimensions)
 end
 
 @doc raw"""
@@ -168,9 +175,9 @@ where ``O`` is the operator and ``\rho`` is the reduced density operator in the
 function QuantumToolbox.expect(op, ados::ADOs; take_real::Bool = true)
     if op isa HEOMSuperOp
         _check_sys_dim_and_ADOs_num(op, ados)
-        exp_val = dot(transpose(_Tr(eltype(ados), ados.dims, ados.N)), (SparseMatrixCSC(op) * ados).data)
+        exp_val = dot(transpose(_Tr(eltype(ados), ados.dimensions, ados.N)), (SparseMatrixCSC(op) * ados).data)
     else
-        _op = HandleMatrixType(op, ados.dims, "op (observable)"; type = Operator)
+        _op = HandleMatrixType(op, ados.dimensions, "op (observable)"; type = Operator)
         exp_val = tr(_op.data * getRho(ados).data)
     end
 
@@ -198,7 +205,7 @@ where ``O`` is the operator and ``\rho`` is the reduced density operator in one
 - `exp_val` : The expectation value
 """
 function QuantumToolbox.expect(op, ados_list::Vector{ADOs}; take_real::Bool = true)
-    dims = ados_list[1].dims
+    dimensions = ados_list[1].dimensions
     N = ados_list[1].N
     for i in 2:length(ados_list)
         _check_sys_dim_and_ADOs_num(ados_list[1], ados_list[i])
@@ -208,9 +215,9 @@ function QuantumToolbox.expect(op, ados_list::Vector{ADOs}; take_real::Bool = tr
         _check_sys_dim_and_ADOs_num(op, ados_list[1])
         _op = op
     else
-        _op = HEOMSuperOp(spre(op), EVEN, dims, N)
+        _op = HEOMSuperOp(spre(op), EVEN, dimensions, N)
     end
-    tr_op = transpose(_Tr(eltype(op), dims, N)) * SparseMatrixCSC(_op).data
+    tr_op = transpose(_Tr(eltype(op), dimensions, N)) * SparseMatrixCSC(_op).data
 
     exp_val = [dot(tr_op, ados.data) for ados in ados_list]
 

src/HeomBase.jl

@@ -2,6 +2,15 @@ export AbstractHEOMLSMatrix
 
 abstract type AbstractHEOMLSMatrix{T} end
 
+function Base.getproperty(M::AbstractHEOMLSMatrix, key::Symbol)
+    # a comment here to avoid bad render by JuliaFormatter
+    if key === :dims
+        return dimensions_to_dims(getfield(M, :dimensions))
+    else
+        return getfield(M, key)
+    end
+end
+
 @doc raw"""
     (M::AbstractHEOMLSMatrix)(p, t)
 
@@ -32,12 +41,12 @@ _get_SciML_matrix_wrapper(M::AddedOperator) = _get_SciML_matrix_wrapper(M.ops[1]
 _get_SciML_matrix_wrapper(M::AbstractHEOMLSMatrix) = _get_SciML_matrix_wrapper(M.data)
 
 # equal to : sparse(vec(system_identity_matrix))
-function _Tr(T::Type{<:Number}, dims::SVector, N::Int)
-    D = prod(dims)
+function _Tr(T::Type{<:Number}, dimensions::Dimensions, N::Int)
+    D = prod(dimensions)
     return SparseVector(N * D^2, [1 + n * (D + 1) for n in 0:(D-1)], ones(T, D))
 end
-_Tr(M::AbstractHEOMLSMatrix) = _Tr(_get_SciML_matrix_wrapper(M), M.dims, M.N)
-_Tr(M::Type{<:SparseMatrixCSC}, dims::SVector, N::Int) = _Tr(eltype(M), dims, N)
+_Tr(M::AbstractHEOMLSMatrix) = _Tr(_get_SciML_matrix_wrapper(M), M.dimensions, M.N)
+_Tr(M::Type{<:SparseMatrixCSC}, dimensions::Dimensions, N::Int) = _Tr(eltype(M), dimensions, N)
 
 function HandleMatrixType(
     M::AbstractQuantumObject,
@@ -55,19 +64,19 @@ function HandleMatrixType(
 end
 function HandleMatrixType(
     M::AbstractQuantumObject,
-    dims::SVector,
+    dimensions::Dimensions,
     MatrixName::String = "";
     type::T = nothing,
 ) where {T<:Union{Nothing,OperatorQuantumObject,SuperOperatorQuantumObject}}
-    if M.dims == dims
+    if M.dimensions == dimensions
         return HandleMatrixType(M, MatrixName; type = type)
     else
-        error("The dims of $(MatrixName) should be: $(dims)")
+        error("The dimensions of $(MatrixName) should be: $(_get_dims_string(dimensions))")
     end
 end
 HandleMatrixType(
     M,
-    dims::SVector,
+    dimensions::Dimensions,
     MatrixName::String = "";
     type::T = nothing,
 ) where {T<:Union{Nothing,OperatorQuantumObject,SuperOperatorQuantumObject}} =
@@ -86,7 +95,7 @@ _HandleVectorType(M::Type{<:SparseMatrixCSC}, V::SparseVector) = Vector{_CType(e
 function _HandleSteadyStateMatrix(M::AbstractHEOMLSMatrix{<:MatrixOperator})
     S = size(M, 1)
     ElType = eltype(M)
-    D = prod(M.dims)
+    D = prod(M.dimensions)
     A = copy(M.data.A)
     A[1, 1:S] .= 0
 
@@ -96,8 +105,8 @@ function _HandleSteadyStateMatrix(M::AbstractHEOMLSMatrix{<:MatrixOperator})
 end
 
 function _check_sys_dim_and_ADOs_num(A, B)
-    if (A.dims != B.dims)
-        error("Inconsistent system dimension (\"dims\").")
+    if (A.dimensions != B.dimensions)
+        error("Inconsistent system dimensions.")
     end
 
     if (A.N != B.N)