Change order of AbstractQuantumObject data type (#371)

qutip · Jan 20, 2025 · b69f2b2 · b69f2b2
2 parents 88bf11b + 5c7445f
commit b69f2b2
Show file tree

Hide file tree

Showing 32 changed files with 469 additions and 601 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - 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])
 - Improve lazy tensor warning for `SciMLOperators`. ([#370])
+- Change order of `AbstractQuantumObject` data type. For example, from `QuantumObject{DataType,ObjType,DimsType}` to `QuantumObject{ObjType,DimsType,DataType}`. ([#371])
 
 ## [v0.24.0]
 Release date: 2024-12-13
@@ -77,3 +78,4 @@ Release date: 2024-11-13
 [#353]: https://github.com/qutip/QuantumToolbox.jl/issues/353
 [#360]: https://github.com/qutip/QuantumToolbox.jl/issues/360
 [#370]: https://github.com/qutip/QuantumToolbox.jl/issues/370
+[#371]: https://github.com/qutip/QuantumToolbox.jl/issues/371
diff --git a/ext/QuantumToolboxCUDAExt.jl b/ext/QuantumToolboxCUDAExt.jl
@@ -1,6 +1,7 @@
 module QuantumToolboxCUDAExt
 
 using QuantumToolbox
+using QuantumToolbox: makeVal, getVal
 import CUDA: cu, CuArray
 import CUDA.CUSPARSE: CuSparseVector, CuSparseMatrixCSC, CuSparseMatrixCSR
 import SparseArrays: SparseVector, SparseMatrixCSC
@@ -10,60 +11,56 @@ import SparseArrays: SparseVector, SparseMatrixCSC
 
 If `A.data` is a dense array, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `CUDA.CuArray` for gpu calculations.
 """
-CuArray(A::QuantumObject{Tq}) where {Tq<:Union{Vector,Matrix}} = QuantumObject(CuArray(A.data), A.type, A.dimensions)
+CuArray(A::QuantumObject) = QuantumObject(CuArray(A.data), A.type, A.dimensions)
 
 @doc raw"""
     CuArray{T}(A::QuantumObject)
 
 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) where {T} = QuantumObject(CuArray{T}(A.data), A.type, A.dimensions)
 
 @doc raw"""
     CuSparseVector(A::QuantumObject)
 
 If `A.data` is a sparse vector, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `CUDA.CUSPARSE.CuSparseVector` for gpu calculations.
 """
-CuSparseVector(A::QuantumObject{<:SparseVector}) = QuantumObject(CuSparseVector(A.data), A.type, A.dimensions)
+CuSparseVector(A::QuantumObject) = QuantumObject(CuSparseVector(A.data), A.type, A.dimensions)
 
 @doc raw"""
     CuSparseVector{T}(A::QuantumObject)
 
 If `A.data` is a sparse vector, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `CUDA.CUSPARSE.CuSparseVector` with element type `T` for gpu calculations.
 """
-CuSparseVector{T}(A::QuantumObject{<:SparseVector}) where {T} =
-    QuantumObject(CuSparseVector{T}(A.data), A.type, A.dimensions)
+CuSparseVector{T}(A::QuantumObject) where {T} = QuantumObject(CuSparseVector{T}(A.data), A.type, A.dimensions)
 
 @doc raw"""
     CuSparseMatrixCSC(A::QuantumObject)
 
 If `A.data` is in the type of `SparseMatrixCSC`, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `CUDA.CUSPARSE.CuSparseMatrixCSC` for gpu calculations.
 """
-CuSparseMatrixCSC(A::QuantumObject{<:SparseMatrixCSC}) = QuantumObject(CuSparseMatrixCSC(A.data), A.type, A.dimensions)
+CuSparseMatrixCSC(A::QuantumObject) = QuantumObject(CuSparseMatrixCSC(A.data), A.type, A.dimensions)
 
 @doc raw"""
     CuSparseMatrixCSC{T}(A::QuantumObject)
 
 If `A.data` is in the type of `SparseMatrixCSC`, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `CUDA.CUSPARSE.CuSparseMatrixCSC` with element type `T` for gpu calculations.
 """
-CuSparseMatrixCSC{T}(A::QuantumObject{<:SparseMatrixCSC}) where {T} =
-    QuantumObject(CuSparseMatrixCSC{T}(A.data), A.type, A.dimensions)
+CuSparseMatrixCSC{T}(A::QuantumObject) where {T} = QuantumObject(CuSparseMatrixCSC{T}(A.data), A.type, A.dimensions)
 
 @doc raw"""
     CuSparseMatrixCSR(A::QuantumObject)
 
 If `A.data` is in the type of `SparseMatrixCSC`, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `CUDA.CUSPARSE.CuSparseMatrixCSR` for gpu calculations.
 """
-CuSparseMatrixCSR(A::QuantumObject{<:SparseMatrixCSC}) = QuantumObject(CuSparseMatrixCSR(A.data), A.type, A.dimensions)
+CuSparseMatrixCSR(A::QuantumObject) = QuantumObject(CuSparseMatrixCSR(A.data), A.type, A.dimensions)
 
 @doc raw"""
     CuSparseMatrixCSR(A::QuantumObject)
 
 If `A.data` is in the type of `SparseMatrixCSC`, return a new [`QuantumObject`](@ref) where `A.data` is in the type of `CUDA.CUSPARSE.CuSparseMatrixCSR` with element type `T` for gpu calculations.
 """
-CuSparseMatrixCSR{T}(A::QuantumObject{<:SparseMatrixCSC}) where {T} =
-    QuantumObject(CuSparseMatrixCSR{T}(A.data), A.type, A.dimensions)
+CuSparseMatrixCSR{T}(A::QuantumObject) where {T} = QuantumObject(CuSparseMatrixCSR{T}(A.data), A.type, A.dimensions)
 
 @doc raw"""
     cu(A::QuantumObject; word_size::Int=64)
@@ -74,15 +71,26 @@ Return a new [`QuantumObject`](@ref) where `A.data` is in the type of `CUDA` arr
 - `A::QuantumObject`: The [`QuantumObject`](@ref)
 - `word_size::Int`: The word size of the element type of `A`, can be either `32` or `64`. Default to `64`.
 """
-cu(A::QuantumObject; word_size::Int = 64) =
-    ((word_size == 64) || (word_size == 32)) ? cu(A, Val(word_size)) :
-    throw(DomainError(word_size, "The word size should be 32 or 64."))
-cu(A::QuantumObject{T}, word_size::TW) where {T<:Union{Vector,Matrix},TW<:Union{Val{32},Val{64}}} =
-    CuArray{_change_eltype(eltype(A), word_size)}(A)
-cu(A::QuantumObject{<:SparseVector}, word_size::TW) where {TW<:Union{Val{32},Val{64}}} =
-    CuSparseVector{_change_eltype(eltype(A), word_size)}(A)
-cu(A::QuantumObject{<:SparseMatrixCSC}, word_size::TW) where {TW<:Union{Val{32},Val{64}}} =
-    CuSparseMatrixCSC{_change_eltype(eltype(A), word_size)}(A)
+function cu(A::QuantumObject; word_size::Union{Val,Int} = Val(64))
+    _word_size = getVal(makeVal(word_size))
+
+    ((_word_size == 64) || (_word_size == 32)) || throw(DomainError(_word_size, "The word size should be 32 or 64."))
+
+    return cu(A, makeVal(word_size))
+end
+cu(A::QuantumObject, word_size::Union{Val{32},Val{64}}) = CuArray{_change_eltype(eltype(A), word_size)}(A)
+function cu(
+    A::QuantumObject{ObjType,DimsType,<:SparseVector},
+    word_size::Union{Val{32},Val{64}},
+) where {ObjType<:QuantumObjectType,DimsType<:AbstractDimensions}
+    return CuSparseVector{_change_eltype(eltype(A), word_size)}(A)
+end
+function cu(
+    A::QuantumObject{ObjType,DimsType,<:SparseMatrixCSC},
+    word_size::Union{Val{32},Val{64}},
+) where {ObjType<:QuantumObjectType,DimsType<:AbstractDimensions}
+    return CuSparseMatrixCSC{_change_eltype(eltype(A), word_size)}(A)
+end
 
 _change_eltype(::Type{T}, ::Val{64}) where {T<:Int} = Int64
 _change_eltype(::Type{T}, ::Val{32}) where {T<:Int} = Int32

diff --git a/ext/QuantumToolboxCairoMakieExt.jl b/ext/QuantumToolboxCairoMakieExt.jl
@@ -6,7 +6,7 @@ using CairoMakie: Axis, Axis3, Colorbar, Figure, GridLayout, heatmap!, surface!,
 @doc raw"""
     plot_wigner(
         library::Val{:CairoMakie},
-        state::QuantumObject{DT,OpType};
+        state::QuantumObject{OpType};
         xvec::Union{Nothing,AbstractVector} = nothing,
         yvec::Union{Nothing,AbstractVector} = nothing,
         g::Real = √2,
@@ -15,7 +15,7 @@ using CairoMakie: Axis, Axis3, Colorbar, Figure, GridLayout, heatmap!, surface!,
         location::Union{GridPosition,Nothing} = nothing,
         colorbar::Bool = false,
         kwargs...
-    ) where {DT,OpType}
+    ) where {OpType}
 
 Plot the [Wigner quasipropability distribution](https://en.wikipedia.org/wiki/Wigner_quasiprobability_distribution) of `state` using the [`CairoMakie`](https://github.com/MakieOrg/Makie.jl/tree/master/CairoMakie) plotting library.
 
@@ -44,7 +44,7 @@ Plot the [Wigner quasipropability distribution](https://en.wikipedia.org/wiki/Wi
 """
 function QuantumToolbox.plot_wigner(
     library::Val{:CairoMakie},
-    state::QuantumObject{DT,OpType};
+    state::QuantumObject{OpType};
     xvec::Union{Nothing,AbstractVector} = LinRange(-7.5, 7.5, 200),
     yvec::Union{Nothing,AbstractVector} = LinRange(-7.5, 7.5, 200),
     g::Real = √2,
@@ -53,7 +53,7 @@ function QuantumToolbox.plot_wigner(
     location::Union{GridPosition,Nothing} = nothing,
     colorbar::Bool = false,
     kwargs...,
-) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
+) where {OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
     QuantumToolbox.getVal(projection) == :two_dim ||
         QuantumToolbox.getVal(projection) == :three_dim ||
         throw(ArgumentError("Unsupported projection: $projection"))
@@ -74,7 +74,7 @@ end
 
 function _plot_wigner(
     ::Val{:CairoMakie},
-    state::QuantumObject{DT,OpType},
+    state::QuantumObject{OpType},
     xvec::AbstractVector,
     yvec::AbstractVector,
     projection::Val{:two_dim},
@@ -83,7 +83,7 @@ function _plot_wigner(
     location::Union{GridPosition,Nothing},
     colorbar::Bool;
     kwargs...,
-) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
+) where {OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
     fig, location = _getFigAndLocation(location)
 
     lyt = GridLayout(location)
@@ -107,7 +107,7 @@ end
 
 function _plot_wigner(
     ::Val{:CairoMakie},
-    state::QuantumObject{DT,OpType},
+    state::QuantumObject{OpType},
     xvec::AbstractVector,
     yvec::AbstractVector,
     projection::Val{:three_dim},
@@ -116,7 +116,7 @@ function _plot_wigner(
     location::Union{GridPosition,Nothing},
     colorbar::Bool;
     kwargs...,
-) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
+) where {OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}}
     fig, location = _getFigAndLocation(location)
 
     lyt = GridLayout(location)

diff --git a/src/correlations.jl b/src/correlations.jl
@@ -24,21 +24,16 @@ Returns the two-times correlation function of three operators ``\hat{A}``, ``\ha
 If the initial state `ψ0` is given as `nothing`, then the [`steadystate`](@ref) will be used as the initial state. Note that this is only implemented if `H` is constant ([`QuantumObject`](@ref)).
 """
 function correlation_3op_2t(
-    H::AbstractQuantumObject{DataType,HOpType},
-    ψ0::Union{Nothing,QuantumObject{<:AbstractArray{T1},StateOpType}},
+    H::AbstractQuantumObject{HOpType},
+    ψ0::Union{Nothing,QuantumObject{StateOpType}},
     tlist::AbstractVector,
     τlist::AbstractVector,
     c_ops::Union{Nothing,AbstractVector,Tuple},
-    A::QuantumObject{<:AbstractArray{T2},OperatorQuantumObject},
-    B::QuantumObject{<:AbstractArray{T3},OperatorQuantumObject},
-    C::QuantumObject{<:AbstractArray{T4},OperatorQuantumObject};
+    A::QuantumObject{OperatorQuantumObject},
+    B::QuantumObject{OperatorQuantumObject},
+    C::QuantumObject{OperatorQuantumObject};
     kwargs...,
 ) where {
-    DataType,
-    T1,
-    T2,
-    T3,
-    T4,
     HOpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject},
     StateOpType<:Union{KetQuantumObject,OperatorQuantumObject},
 }
@@ -79,20 +74,15 @@ Returns the one-time correlation function of three operators ``\hat{A}``, ``\hat
 If the initial state `ψ0` is given as `nothing`, then the [`steadystate`](@ref) will be used as the initial state. Note that this is only implemented if `H` is constant ([`QuantumObject`](@ref)).
 """
 function correlation_3op_1t(
-    H::AbstractQuantumObject{DataType,HOpType},
-    ψ0::Union{Nothing,QuantumObject{<:AbstractArray{T1},StateOpType}},
+    H::AbstractQuantumObject{HOpType},
+    ψ0::Union{Nothing,QuantumObject{StateOpType}},
     τlist::AbstractVector,
     c_ops::Union{Nothing,AbstractVector,Tuple},
-    A::QuantumObject{<:AbstractArray{T2},OperatorQuantumObject},
-    B::QuantumObject{<:AbstractArray{T3},OperatorQuantumObject},
-    C::QuantumObject{<:AbstractArray{T4},OperatorQuantumObject};
+    A::QuantumObject{OperatorQuantumObject},
+    B::QuantumObject{OperatorQuantumObject},
+    C::QuantumObject{OperatorQuantumObject};
     kwargs...,
 ) where {
-    DataType,
-    T1,
-    T2,
-    T3,
-    T4,
     HOpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject},
     StateOpType<:Union{KetQuantumObject,OperatorQuantumObject},
 }
@@ -119,20 +109,16 @@ If the initial state `ψ0` is given as `nothing`, then the [`steadystate`](@ref)
 When `reverse=true`, the correlation function is calculated as ``\left\langle \hat{A}(t) \hat{B}(t + \tau) \right\rangle``.
 """
 function correlation_2op_2t(
-    H::AbstractQuantumObject{DataType,HOpType},
-    ψ0::Union{Nothing,QuantumObject{<:AbstractArray{T1},StateOpType}},
+    H::AbstractQuantumObject{HOpType},
+    ψ0::Union{Nothing,QuantumObject{StateOpType}},
     tlist::AbstractVector,
     τlist::AbstractVector,
     c_ops::Union{Nothing,AbstractVector,Tuple},
-    A::QuantumObject{<:AbstractArray{T2},OperatorQuantumObject},
-    B::QuantumObject{<:AbstractArray{T3},OperatorQuantumObject};
+    A::QuantumObject{OperatorQuantumObject},
+    B::QuantumObject{OperatorQuantumObject};
     reverse::Bool = false,
     kwargs...,
 ) where {
-    DataType,
-    T1,
-    T2,
-    T3,
     HOpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject},
     StateOpType<:Union{KetQuantumObject,OperatorQuantumObject},
 }
@@ -163,19 +149,15 @@ If the initial state `ψ0` is given as `nothing`, then the [`steadystate`](@ref)
 When `reverse=true`, the correlation function is calculated as ``\left\langle \hat{A}(0) \hat{B}(\tau) \right\rangle``.
 """
 function correlation_2op_1t(
-    H::AbstractQuantumObject{DataType,HOpType},
-    ψ0::Union{Nothing,QuantumObject{<:AbstractArray{T1},StateOpType}},
+    H::AbstractQuantumObject{HOpType},
+    ψ0::Union{Nothing,QuantumObject{StateOpType}},
     τlist::AbstractVector,
     c_ops::Union{Nothing,AbstractVector,Tuple},
-    A::QuantumObject{<:AbstractArray{T2},OperatorQuantumObject},
-    B::QuantumObject{<:AbstractArray{T3},OperatorQuantumObject};
+    A::QuantumObject{OperatorQuantumObject},
+    B::QuantumObject{OperatorQuantumObject};
     reverse::Bool = false,
     kwargs...,
 ) where {
-    DataType,
-    T1,
-    T2,
-    T3,
     HOpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject},
     StateOpType<:Union{KetQuantumObject,OperatorQuantumObject},
 }