introduce inner function _convert_u0

Author: ytdHuang

ext/QuantumToolboxCUDAExt.jl

@@ -1,6 +1,7 @@
 module QuantumToolboxCUDAExt
 
 using QuantumToolbox
+import QuantumToolbox: _convert_u0
 import CUDA: cu, CuArray
 import CUDA.CUSPARSE: CuSparseVector, CuSparseMatrixCSC, CuSparseMatrixCSR
 import SparseArrays: SparseVector, SparseMatrixCSC
@@ -89,4 +90,7 @@ _change_eltype(::Type{T}, ::Val{32}) where {T<:AbstractFloat} = Float32
 _change_eltype(::Type{Complex{T}}, ::Val{64}) where {T<:Union{Int,AbstractFloat}} = ComplexF64
 _change_eltype(::Type{Complex{T}}, ::Val{32}) where {T<:Union{Int,AbstractFloat}} = ComplexF32
 
+# make sure u0 in time evolution is dense vector and has complex element type
+_convert_u0(u0::Union{CuArray{T},CuSparseVector{T}}) where {T<:Number} = convert(CuArray{complex(T)}, u0)
+
 end

src/time_evolution/mcsolve.jl

@@ -190,7 +190,7 @@ function mcsolveProblem(
     c_ops isa Nothing &&
         throw(ArgumentError("The list of collapse operators must be provided. Use sesolveProblem instead."))
 
-    t_l = _convert_tlist(eltype(H), tlist) # Convert it to support GPUs and avoid type instabilities for OrdinaryDiffEq.jl
+    t_l = convert(Vector{real(eltype(ψ0))}, tlist) # Convert it to support GPUs and avoid type instabilities for OrdinaryDiffEq.jl
 
     H_eff = H - 1im * mapreduce(op -> op' * op, +, c_ops) / 2
 

src/time_evolution/mesolve.jl

@@ -122,9 +122,9 @@ function mesolveProblem(
     is_time_dependent = !(H_t isa Nothing)
     progress_bar_val = makeVal(progress_bar)
 
-    t_l = _convert_tlist(eltype(H), tlist) # Convert it to support GPUs and avoid type instabilities for OrdinaryDiffEq.jl
+    ρ0 = _convert_u0(mat2vec(ket2dm(ψ0).data))
 
-    ρ0 = mat2vec(ket2dm(ψ0).data)
+    t_l = convert(Vector{real(eltype(ρ0))}, tlist) # Convert it to support GPUs and avoid type instabilities for OrdinaryDiffEq.jl
 
     L = liouvillian(H, c_ops).data
     progr = ProgressBar(length(t_l), enable = getVal(progress_bar_val))

src/time_evolution/sesolve.jl

@@ -103,9 +103,9 @@ function sesolveProblem(
     is_time_dependent = !(H_t isa Nothing)
     progress_bar_val = makeVal(progress_bar)
 
-    t_l = _convert_tlist(eltype(H), tlist) # Convert it to support GPUs and avoid type instabilities for OrdinaryDiffEq.jl
+    ϕ0 = _convert_u0(get_data(ψ0))
 
-    ϕ0 = get_data(ψ0)
+    t_l = convert(Vector{real(eltype(ϕ0))}, tlist) # Convert it to support GPUs and avoid type instabilities for OrdinaryDiffEq.jl
 
     U = -1im * get_data(H)
     progr = ProgressBar(length(t_l), enable = getVal(progress_bar_val))

src/utilities.jl

@@ -66,12 +66,5 @@ _non_static_array_warning(argname, arg::AbstractVector{T}) where {T} =
           join(arg, ", ") *
           ")` instead of `$argname = $arg`." maxlog = 1
 
-# convert tlist in time evolution
-_convert_tlist(::Type{Int32}, tlist::AbstractVector) = _convert_tlist(Val(32), tlist)
-_convert_tlist(::Type{Float32}, tlist::AbstractVector) = _convert_tlist(Val(32), tlist)
-_convert_tlist(::Type{ComplexF32}, tlist::AbstractVector) = _convert_tlist(Val(32), tlist)
-_convert_tlist(::Type{Int64}, tlist::AbstractVector) = _convert_tlist(Val(64), tlist)
-_convert_tlist(::Type{Float64}, tlist::AbstractVector) = _convert_tlist(Val(64), tlist)
-_convert_tlist(::Type{ComplexF64}, tlist::AbstractVector) = _convert_tlist(Val(64), tlist)
-_convert_tlist(::Val{32}, tlist::AbstractVector) = convert(Vector{Float32}, tlist)
-_convert_tlist(::Val{64}, tlist::AbstractVector) = convert(Vector{Float64}, tlist)
+# make sure u0 in time evolution is dense vector and has complex element type
+_convert_u0(u0::AbstractVector{T}) where {T<:Number} = convert(Vector{complex(T)}, u0)