Merge pull request #119 from albertomercurio/dev/patch-1

Make steadystate function almost compatible with GPU arrays
qutip · May 22, 2024 · e73ee0a · e73ee0a
2 parents f0a6e1c + 9757448
commit e73ee0a
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 5 deletions.
diff --git a/Project.toml b/Project.toml
@@ -4,6 +4,7 @@ authors = ["Alberto Mercurio"]
 version = "0.8.2"
 
 [deps]
+ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
@@ -25,6 +26,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 QuantumToolboxCUDAExt = "CUDA"
 
 [compat]
+ArrayInterface = "6, 7"
 CUDA = "5"
 DiffEqCallbacks = "2, 3"
 FFTW = "1.5"

diff --git a/src/QuantumToolbox.jl b/src/QuantumToolbox.jl
@@ -22,6 +22,7 @@ else
 end
 
 # other dependencies (in alphabetical order)
+import ArrayInterface: allowed_getindex, allowed_setindex!
 import DiffEqCallbacks: DiscreteCallback, PeriodicCallback, PresetTimeCallback
 import FFTW: fft, fftshift
 import Graphs: connected_components, DiGraph

diff --git a/src/general_functions.jl b/src/general_functions.jl
@@ -442,3 +442,6 @@ Convert a quantum object from matrix ([`OperatorQuantumObject`](@ref)-type) to v
 """
 mat2vec(A::QuantumObject{<:AbstractArray{T},OperatorQuantumObject}) where {T} =
     QuantumObject(mat2vec(A.data), OperatorKet, A.dims)
+
+_get_dense_similar(A::AbstractArray, args...) = similar(A, args...)
+_get_dense_similar(A::AbstractSparseMatrix, args...) = similar(nonzeros(A), args...)
diff --git a/src/steadystate.jl b/src/steadystate.jl
@@ -26,15 +26,25 @@ function _steadystate(
     L::QuantumObject{<:AbstractArray{T},SuperOperatorQuantumObject},
     solver::SteadyStateDirectSolver,
 ) where {T}
-    L_tmp = copy(L.data)
+    L_tmp = L.data
     N = prod(L.dims)
     weight = norm(L_tmp, 1) / length(L_tmp)
-    v0 = zeros(ComplexF64, N^2) # This is not scalable for GPU arrays
-    v0[1] = weight
 
-    L_tmp[1, [N * (i - 1) + i for i in 1:N]] .+= weight
+    v0 = _get_dense_similar(L_tmp, N^2)
+    fill!(v0, 0)
+    allowed_setindex!(v0, weight, 1) # Because scalar indexing is not allowed on GPU arrays
 
-    ρss_vec = L_tmp \ v0
+    idx_range = collect(1:N)
+    rows = _get_dense_similar(L_tmp, N)
+    cols = _get_dense_similar(L_tmp, N)
+    datas = _get_dense_similar(L_tmp, N)
+    fill!(rows, 1)
+    copyto!(cols, N .* (idx_range .- 1) .+ idx_range)
+    fill!(datas, weight)
+    Tn = sparse(rows, cols, datas, N^2, N^2)
+    L_tmp = L_tmp + Tn
+
+    ρss_vec = L_tmp \ v0 # This is still not supported on GPU, yet
     ρss = reshape(ρss_vec, N, N)
     ρss = (ρss + ρss') / 2 # Hermitianize
     return QuantumObject(ρss, Operator, L.dims)