Several updates for Dimensions structure (#367)

ytdHuang · web-flow · commit 309a6704bba4 · 2025-01-07T20:02:36.000+09:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Change the structure of block diagonalization functions, using `BlockDiagonalForm` struct and changing the function name from `bdf` to `block_diagonal_form`. ([#349])
 - Add **GPUArrays** compatibility for `ptrace` function, by using **KernelAbstractions.jl**. ([#350])
+- Introduce `Space`, `Dimensions`, `GeneralDimensions` structures to support wider definitions and operations of `Qobj/QobjEvo`, and potential functionalities in the future. ([#271], [#353], [#360])
 
 ## [v0.24.0]
 Release date: 2024-12-13
@@ -55,6 +56,7 @@ Release date: 2024-11-13
 [v0.24.0]: https://github.com/qutip/QuantumToolbox.jl/releases/tag/v0.24.0
 [#86]: https://github.com/qutip/QuantumToolbox.jl/issues/86
 [#139]: https://github.com/qutip/QuantumToolbox.jl/issues/139
+[#271]: https://github.com/qutip/QuantumToolbox.jl/issues/271
 [#292]: https://github.com/qutip/QuantumToolbox.jl/issues/292
 [#306]: https://github.com/qutip/QuantumToolbox.jl/issues/306
 [#309]: https://github.com/qutip/QuantumToolbox.jl/issues/309
@@ -71,3 +73,5 @@ Release date: 2024-11-13
 [#347]: https://github.com/qutip/QuantumToolbox.jl/issues/347
 [#349]: https://github.com/qutip/QuantumToolbox.jl/issues/349
 [#350]: https://github.com/qutip/QuantumToolbox.jl/issues/350
+[#353]: https://github.com/qutip/QuantumToolbox.jl/issues/353
+[#360]: https://github.com/qutip/QuantumToolbox.jl/issues/360
diff --git a/ext/QuantumToolboxCUDAExt.jl b/ext/QuantumToolboxCUDAExt.jl
@@ -17,7 +17,8 @@ CuArray(A::QuantumObject{Tq}) where {Tq<:Union{Vector,Matrix}} = QuantumObject(C
 
 If `A.data` is a dense array, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `CUDA.CuArray` with element type `T` for gpu calculations.
 """
-CuArray{T}(A::QuantumObject{Tq}) where {T,Tq<:Union{Vector,Matrix}} = QuantumObject(CuArray{T}(A.data), A.type, A.dimensions)
+CuArray{T}(A::QuantumObject{Tq}) where {T,Tq<:Union{Vector,Matrix}} =
+    QuantumObject(CuArray{T}(A.data), A.type, A.dimensions)
 
 @doc raw"""
     CuSparseVector(A::QuantumObject)
diff --git a/src/negativity.jl b/src/negativity.jl
@@ -81,7 +81,7 @@ end
 function _partial_transpose(ρ::QuantumObject{DT,OperatorQuantumObject}, mask::Vector{Bool}) where {DT<:AbstractArray}
     isa(ρ.dimensions, GeneralDimensions) &&
         (get_dimensions_to(ρ) != get_dimensions_from(ρ)) &&
-        throw(ArgumentError("Invalid partial transpose for dims = $(ρ.dims)"))
+        throw(ArgumentError("Invalid partial transpose for dims = $(_get_dims_string(ρ.dimensions))"))
 
     mask2 = [1 + Int(i) for i in mask]
     # mask2 has elements with values equal to 1 or 2
@@ -106,7 +106,7 @@ end
 function _partial_transpose(ρ::QuantumObject{<:AbstractSparseArray,OperatorQuantumObject}, mask::Vector{Bool})
     isa(ρ.dimensions, GeneralDimensions) &&
         (get_dimensions_to(ρ) != get_dimensions_from(ρ)) &&
-        throw(ArgumentError("Invalid partial transpose for dims = $(ρ.dims)"))
+        throw(ArgumentError("Invalid partial transpose for dims = $(_get_dims_string(ρ.dimensions))"))
 
     M, N = size(ρ)
     dimsTuple = Tuple(dimensions_to_dims(get_dimensions_to(ρ)))
diff --git a/src/qobj/arithmetic_and_attributes.jl b/src/qobj/arithmetic_and_attributes.jl
@@ -612,9 +612,10 @@ 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})
+    # TODO: support for special cases when some of the subsystems have same `to` and `from` space
     isa(QO.dimensions, GeneralDimensions) &&
         (get_dimensions_to(QO) != get_dimensions_from(QO)) &&
-        throw(ArgumentError("Invalid partial trace for dims = $(QO.dims)"))
+        throw(ArgumentError("Invalid partial trace for dims = $(_get_dims_string(QO.dimensions))"))
 
     _non_static_array_warning("sel", sel)
 
@@ -800,6 +801,11 @@ function permute(
     A::QuantumObject{<:AbstractArray{T},ObjType},
     order::Union{AbstractVector{Int},Tuple},
 ) where {T,ObjType<:Union{KetQuantumObject,BraQuantumObject,OperatorQuantumObject}}
+    # TODO: fix this (should be able to permute arbitrary GeneralDimensions)
+    isa(A.dimensions, GeneralDimensions) &&
+        (get_dimensions_to(A) != get_dimensions_from(A)) &&
+        throw(ArgumentError("Invalid permute for dims = $(_get_dims_string(A.dimensions))"))
+
     (length(order) != length(A.dimensions)) &&
         throw(ArgumentError("The order list must have the same length as the number of subsystems (A.dims)"))
 
@@ -829,8 +835,9 @@ _dims_and_perm(::OperatorQuantumObject, dims::SVector{N,Int}, order::AbstractVec
     reverse(vcat(dims, dims)), reverse((2 * L + 1) .- vcat(order, order .+ L))
 
 # if dims originates from GeneralDimensions
-_dims_and_perm(::OperatorQuantumObject, dims::SVector{2,SVector{N,Int}}, order::AbstractVector{Int}, L::Int) where {N} =
-    reverse(vcat(dims[1], dims[2])), reverse((2 * L + 1) .- vcat(order, order .+ L))
+# TODO: fix this
+#= _dims_and_perm(::OperatorQuantumObject, dims::SVector{2,SVector{N,Int}}, order::AbstractVector{Int}, L::Int) where {N} =
+    reverse(vcat(dims[1], dims[2])), reverse((2 * L + 1) .- vcat(order, order .+ L)) =#
 
 _order_dimensions(dimensions::Dimensions{N}, order::AbstractVector{Int}) where {N} = Dimensions{N}(dimensions.to[order])
 _order_dimensions(dimensions::GeneralDimensions{N}, order::AbstractVector{Int}) where {N} =
diff --git a/src/qobj/dimensions.jl b/src/qobj/dimensions.jl
@@ -6,7 +6,7 @@ export AbstractDimensions, Dimensions, GeneralDimensions
 
 abstract type AbstractDimensions{N} end
 
-@doc raw""""
+@doc raw"""
     struct Dimensions{N} <: AbstractDimensions{N}
         to::NTuple{N,AbstractSpace}
     end
@@ -31,7 +31,7 @@ Dimensions(dims::Any) = throw(
     ),
 )
 
-@doc raw""""
+@doc raw"""
     struct GeneralDimensions{N} <: AbstractDimensions{N}
         to::NTuple{N,AbstractSpace}
         from::NTuple{N,AbstractSpace}
@@ -71,7 +71,8 @@ _gen_dimensions(dims::Any) = Dimensions(dims)
 # obtain dims in the type of SVector with integers
 dimensions_to_dims(dimensions::NTuple{N,AbstractSpace}) where {N} = vcat(map(dimensions_to_dims, dimensions)...)
 dimensions_to_dims(dimensions::Dimensions) = dimensions_to_dims(dimensions.to)
-dimensions_to_dims(dimensions::GeneralDimensions) = SVector{2}(dimensions_to_dims(dimensions.to), dimensions_to_dims(dimensions.from))
+dimensions_to_dims(dimensions::GeneralDimensions) =
+    SVector{2}(dimensions_to_dims(dimensions.to), dimensions_to_dims(dimensions.from))
 
 Base.length(::AbstractDimensions{N}) where {N} = N
 
@@ -89,4 +90,4 @@ _get_dims_string(dimensions::Dimensions) = string(dimensions_to_dims(dimensions)
 function _get_dims_string(dimensions::GeneralDimensions)
     dims = dimensions_to_dims(dimensions)
     return "[$(string(dims[1])), $(string(dims[2]))]"
-end
+end
diff --git a/src/qobj/eigsolve.jl b/src/qobj/eigsolve.jl
@@ -351,7 +351,7 @@ end
         c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
         alg::OrdinaryDiffEqAlgorithm = Tsit5(),
         params::NamedTuple = NamedTuple(),
-        ρ0::AbstractMatrix = rand_dm(prod(H.dims)).data,
+        ρ0::AbstractMatrix = rand_dm(prod(H.dimensions)).data,
         k::Int = 1,
         krylovdim::Int = min(10, size(H, 1)),
         maxiter::Int = 200,
@@ -390,7 +390,7 @@ function eigsolve_al(
     c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     params::NamedTuple = NamedTuple(),
-    ρ0::AbstractMatrix = rand_dm(prod(H.dims)).data,
+    ρ0::AbstractMatrix = rand_dm(prod(H.dimensions)).data,
     k::Int = 1,
     krylovdim::Int = min(10, size(H, 1)),
     maxiter::Int = 200,
diff --git a/src/qobj/quantum_object.jl b/src/qobj/quantum_object.jl
@@ -18,7 +18,7 @@ export QuantumObject
 !!! note "`dims` property"
     For a given `H::QuantumObject`, `H.dims` or `getproperty(H, :dims)` returns its `dimensions` in the type of integer-vector.
 
-Julia struct representing any quantum objects.
+Julia structure representing any time-independent quantum objects. For time-dependent cases, see [`QuantumObjectEvolution`](@ref).
 
 # Examples
 
@@ -150,7 +150,15 @@ function Base.show(
     },
 }
     op_data = QO.data
-    println(io, "\nQuantum Object:   type=", QO.type, "   dims=", _get_dims_string(QO.dimensions), "   size=", size(op_data))
+    println(
+        io,
+        "\nQuantum Object:   type=",
+        QO.type,
+        "   dims=",
+        _get_dims_string(QO.dimensions),
+        "   size=",
+        size(op_data),
+    )
     return show(io, MIME("text/plain"), op_data)
 end
 
diff --git a/src/qobj/quantum_object_base.jl b/src/qobj/quantum_object_base.jl
@@ -219,6 +219,7 @@ Base.getproperty(A::AbstractQuantumObject, key::Symbol) = getproperty(A, makeVal
 Base.getproperty(A::AbstractQuantumObject, ::Val{:dims}) = dimensions_to_dims(getfield(A, :dimensions))
 Base.getproperty(A::AbstractQuantumObject, ::Val{K}) where {K} = getfield(A, K)
 
+# this returns `to` in GeneralDimensions representation
 get_dimensions_to(A::AbstractQuantumObject{DT,KetQuantumObject,Dimensions{N}}) where {DT,N} = A.dimensions.to
 get_dimensions_to(A::AbstractQuantumObject{DT,BraQuantumObject,Dimensions{N}}) where {DT,N} = space_one_list(N)
 get_dimensions_to(A::AbstractQuantumObject{DT,OperatorQuantumObject,Dimensions{N}}) where {DT,N} = A.dimensions.to
@@ -229,6 +230,7 @@ get_dimensions_to(
 ) where {DT,ObjType<:Union{SuperOperatorQuantumObject,OperatorBraQuantumObject,OperatorKetQuantumObject},N} =
     A.dimensions.to
 
+# this returns `from` in GeneralDimensions representation
 get_dimensions_from(A::AbstractQuantumObject{DT,KetQuantumObject,Dimensions{N}}) where {DT,N} = space_one_list(N)
 get_dimensions_from(A::AbstractQuantumObject{DT,BraQuantumObject,Dimensions{N}}) where {DT,N} = A.dimensions.to
 get_dimensions_from(A::AbstractQuantumObject{DT,OperatorQuantumObject,Dimensions{N}}) where {DT,N} = A.dimensions.to
diff --git a/src/qobj/quantum_object_evo.jl b/src/qobj/quantum_object_evo.jl
@@ -22,7 +22,7 @@ Julia struct representing any time-dependent quantum object. The `data` field is
 
 where ``c_i(p, t)`` is a function that depends on the parameters `p` and time `t`, and ``\hat{O}_i`` are the operators that form the quantum object. 
 
-For more information about `AbstractSciMLOperator`, see the [SciML](https://docs.sciml.ai/SciMLOperators/stable/) documentation.
+For time-independent cases, see [`QuantumObject`](@ref), and for more information about `AbstractSciMLOperator`, see the [SciML](https://docs.sciml.ai/SciMLOperators/stable/) documentation.
 
 # Examples
 This operator can be initialized in the same way as the QuTiP `QobjEvo` object. For example
diff --git a/src/qobj/space.jl b/src/qobj/space.jl
@@ -11,7 +11,7 @@ abstract type AbstractSpace end
         size::Int
     end
 
-A structure that describes a single Hilbert space.
+A structure that describes a single Hilbert space with size = `size`.
 """
 struct Space <: AbstractSpace
     size::Int
diff --git a/src/qobj/superoperators.jl b/src/qobj/superoperators.jl
@@ -141,7 +141,7 @@ lindblad_dissipator(O::AbstractQuantumObject{DT,OperatorQuantumObject}, Id_cache
 lindblad_dissipator(O::AbstractQuantumObject{DT,SuperOperatorQuantumObject}, Id_cache = nothing) where {DT} = O
 
 @doc raw"""
-    liouvillian(H::AbstractQuantumObject, c_ops::Union{Nothing,AbstractVector,Tuple}=nothing, Id_cache=I(prod(H.dims)))
+    liouvillian(H::AbstractQuantumObject, c_ops::Union{Nothing,AbstractVector,Tuple}=nothing, Id_cache=I(prod(H.dimensions)))
 
 Construct the Liouvillian [`SuperOperator`](@ref) for a system Hamiltonian ``\hat{H}`` and a set of collapse operators ``\{\hat{C}_n\}_n``:
 
diff --git a/test/core-test/negativity_and_partial_transpose.jl b/test/core-test/negativity_and_partial_transpose.jl
@@ -35,6 +35,7 @@
             end
         end
         @test_throws ArgumentError partial_transpose(A_dense, [true])
+        @test_throws ArgumentError partial_transpose(Qobj(zeros(ComplexF64, 3, 2)), [true]) # invalid GeneralDimensions
 
         @testset "Type Inference (partial_transpose)" begin
             @inferred partial_transpose(A_dense, [true, false, true])
diff --git a/test/core-test/quantum_objects.jl b/test/core-test/quantum_objects.jl
@@ -253,6 +253,16 @@
         @test opstring ==
               "\nQuantum Object:   type=Operator   dims=$a_dims   size=$a_size   ishermitian=$a_isherm\n$datastring"
 
+        # GeneralDimensions
+        Gop = tensor(a, ψ)
+        opstring = sprint((t, s) -> show(t, "text/plain", s), Gop)
+        datastring = sprint((t, s) -> show(t, "text/plain", s), Gop.data)
+        Gop_dims = [[N, N], [N, 1]]
+        Gop_size = size(Gop)
+        Gop_isherm = isherm(Gop)
+        @test opstring ==
+              "Quantum Object:   type=Operator   dims=$Gop_dims   size=$Gop_size   ishermitian=$Gop_isherm\n$datastring"
+
         a = spre(a)
         opstring = sprint((t, s) -> show(t, "text/plain", s), a)
         datastring = sprint((t, s) -> show(t, "text/plain", s), a.data)
@@ -655,14 +665,14 @@
 
         # use GeneralDimensions to do partial trace
         ρ1_compound = Qobj(zeros(ComplexF64, 2, 2), dims = ((2, 1), (2, 1)))
-        basis2 = [tensor(eye(2), basis(2, i)) for i in 0:1]
-        for b in basis2
-            ρ1_compound += b' * ρ * b
+        II = qeye(2)
+        basis_list = [basis(2, i) for i in 0:1]
+        for b in basis_list
+            ρ1_compound += tensor(II, b') * ρ * tensor(II, b)
         end
         ρ2_compound = Qobj(zeros(ComplexF64, 2, 2), dims = ((1, 2), (1, 2)))
-        basis1 = [tensor(basis(2, i), eye(2)) for i in 0:1]
-        for b in basis1
-            ρ2_compound += b' * ρ * b
+        for b in basis_list
+            ρ2_compound += tensor(b', II) * ρ * tensor(b, II)
         end
         @test ρ1.data ≈ ρ1_ptr.data ≈ ρ1_compound.data
         @test ρ2.data ≈ ρ2_ptr.data ≈ ρ2_compound.data
@@ -724,6 +734,7 @@
         @test_throws ArgumentError ptrace(ρtotal, (0, 2))
         @test_throws ArgumentError ptrace(ρtotal, (2, 5))
         @test_throws ArgumentError ptrace(ρtotal, (2, 2, 3))
+        @test_throws ArgumentError ptrace(Qobj(zeros(ComplexF64, 3, 2)), 1) # invalid GeneralDimensions
 
         @testset "Type Inference (ptrace)" begin
             @inferred ptrace(ρ, 1)
@@ -736,6 +747,7 @@
     end
 
     @testset "permute" begin
+        # standard Dimensions
         ket_a = Qobj(rand(ComplexF64, 2))
         ket_b = Qobj(rand(ComplexF64, 3))
         ket_c = Qobj(rand(ComplexF64, 4))
@@ -770,6 +782,21 @@
         @test_throws ArgumentError permute(op_bdca, wrong_order1)
         @test_throws ArgumentError permute(op_bdca, wrong_order2)
 
+        # GeneralDimensions
+        # TODO: support for GeneralDimensions
+        #= 
+        Gop_d = Qobj(rand(ComplexF64, 5, 6))
+        compound_bdca = permute(tensor(ket_a, op_b, bra_c, Gop_d), (2, 4, 3, 1))
+        compound_dacb = permute(tensor(ket_a, op_b, bra_c, Gop_d), (4, 1, 3, 2))
+        @test compound_bdca ≈ tensor(op_b, Gop_d, bra_c, ket_a)
+        @test compound_dacb ≈ tensor(Gop_d, ket_a, bra_c, op_b)
+        @test compound_bdca.dims == [[3, 5, 1, 2], [3, 6, 4, 1]]
+        @test compound_dacb.dims == [[5, 2, 1, 3], [6, 1, 4, 3]]
+        @test isoper(compound_bdca)
+        @test isoper(compound_dacb)
+        =#
+        @test_throws ArgumentError permute(Qobj(zeros(ComplexF64, 3, 2)), (1,))
+
         @testset "Type Inference (permute)" begin
             @inferred permute(ket_bdca, (2, 4, 3, 1))
             @inferred permute(bra_bdca, (2, 4, 3, 1))
