Skip to content

Commit 34fa9b4

Browse files
Fe-r-ozFe-r-oz
committed
avoid scalar indexing
1 parent 8cb4e17 commit 34fa9b4

File tree

3 files changed

+63
-55
lines changed

3 files changed

+63
-55
lines changed

‎ext/QuantumToolboxCUDAExt.jl‎

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,8 @@ module QuantumToolboxCUDAExt
22

33
using QuantumToolbox
44
using QuantumToolbox: makeVal, getVal
5-
import QuantumToolbox: _sparse_similar, _convert_eltype_wordsize, allowed_setindex!
6-
import CUDA: cu, CuArray, allowscalar, @allowscalar, has_cuda
5+
import QuantumToolbox: _sparse_similar, _convert_eltype_wordsize
6+
import CUDA: cu, CuArray, allowscalar
77
import CUDA.CUSPARSE: CuSparseVector, CuSparseMatrixCSC, CuSparseMatrixCSR, AbstractCuSparseArray
88
import SparseArrays: SparseVector, SparseMatrixCSC, sparse, spzeros
99
import CUDA.Adapt: adapt
@@ -104,11 +104,22 @@ QuantumToolbox.to_dense(::Type{T}, A::AbstractCuSparseArray) where {T<:Number} =
104104

105105
QuantumToolbox._sparse_similar(A::CuSparseMatrixCSC, args...) = sparse(args..., fmt = :csc)
106106
QuantumToolbox._sparse_similar(A::CuSparseMatrixCSR, args...) = sparse(args..., fmt = :csr)
107-
QuantumToolbox._sparse_similar(A::CuSparseMatrixCSC, I::AbstractVector, J::AbstractVector, V::AbstractVector, m::Int, n::Int) =
108-
CuSparseMatrixCSC(sparse(I, J, V, m, n))
107+
QuantumToolbox._sparse_similar(
108+
A::CuSparseMatrixCSC,
109+
I::AbstractVector,
110+
J::AbstractVector,
111+
V::AbstractVector,
112+
m::Int,
113+
n::Int,
114+
) = CuSparseMatrixCSC(sparse(I, J, V, m, n))
109115
QuantumToolbox._sparse_similar(A::CuSparseMatrixCSC, m::Int, n::Int) = CuSparseMatrixCSC(spzeros(eltype(A), m, n))
110-
QuantumToolbox._sparse_similar(A::CuSparseMatrixCSR, I::AbstractVector, J::AbstractVector, V::AbstractVector, m::Int, n::Int) =
111-
CuSparseMatrixCSR(sparse(I, J, V, m, n))
116+
QuantumToolbox._sparse_similar(
117+
A::CuSparseMatrixCSR,
118+
I::AbstractVector,
119+
J::AbstractVector,
120+
V::AbstractVector,
121+
m::Int,
122+
n::Int,
123+
) = CuSparseMatrixCSR(sparse(I, J, V, m, n))
112124
QuantumToolbox._sparse_similar(A::CuSparseMatrixCSR, m::Int, n::Int) = CuSparseMatrixCSR(spzeros(eltype(A), m, n))
113-
QuantumToolbox.allowed_setindex!(A::AbstractCuSparseArray, v, I...) = @allowscalar A[I...] = v
114125
end

‎src/steadystate.jl‎

Lines changed: 45 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -355,76 +355,75 @@ function _steadystate_fourier(
355355
tol::R = 1e-8,
356356
kwargs...,
357357
) where {R<:Real}
358-
T1 = eltype(L_0)
359-
T2 = eltype(L_p)
360-
T3 = eltype(L_m)
361-
T = promote_type(T1, T2, T3)
358+
# Get matrix data and promote types
359+
L0_mat = get_data(L_0)
360+
Lp_mat = get_data(L_p)
361+
Lm_mat = get_data(L_m)
362+
T = promote_type(eltype(L0_mat), eltype(Lp_mat), eltype(Lm_mat))
362363

363-
L_0_mat = get_data(L_0)
364-
L_p_mat = get_data(L_p)
365-
L_m_mat = get_data(L_m)
366-
367-
N = size(L_0_mat, 1)
364+
N = size(L0_mat, 1)
368365
Ns = isqrt(N)
369366
n_fourier = 2 * n_max + 1
370367
n_list = (-n_max):n_max
371368

369+
# Create stabilization matrix Mn
372370
weight = one(T)
373-
rows = ones(Int, Ns)
374-
cols = [Ns * (j - 1) + j for j in 1:Ns]
375-
vals = fill(weight, Ns)
376-
Mn = _sparse_similar(L_0_mat, rows, cols, vals, N, N)
377-
L = L_0_mat + Mn
378-
379-
M = _sparse_similar(L_0_mat, n_fourier * N, n_fourier * N)
380-
for i in 1:(n_fourier-1)
381-
M_block = _sparse_similar(L_0_mat, N, N)
382-
copyto!(M_block, L_p_mat)
383-
M[(i*N+1):((i+1)*N), ((i-1)*N+1):(i*N)] = M_block
384-
M_block = _sparse_similar(L_0_mat, N, N)
385-
copyto!(M_block, L_m_mat)
386-
M[((i-1)*N+1):(i*N), (i*N+1):((i+1)*N)] = M_block
387-
end
371+
diag_indices = Ns * (0:(Ns-1)) .+ (1:Ns)
372+
Mn = _sparse_similar(L0_mat, ones(Int, Ns), diag_indices, fill(weight, Ns), N, N)
388373

389-
for i in 1:n_fourier
390-
n = n_list[i]
391-
M_block = _sparse_similar(L_0_mat, N, N)
392-
copyto!(M_block, L)
393-
M_block -= 1im * ωd * n * sparse(I, N, N)
394-
M[((i-1)*N+1):(i*N), ((i-1)*N+1):(i*N)] = M_block
395-
end
374+
# Base Liouvillian
375+
L = L0_mat + Mn
376+
377+
# Build off-diagonal blocks
378+
Kp = _sparse_similar(L0_mat, spdiagm(-1 => ones(T, n_fourier-1)))
379+
Km = _sparse_similar(L0_mat, spdiagm(1 => ones(T, n_fourier-1)))
380+
M = kron(Kp, Lm_mat) + kron(Km, Lp_mat)
396381

397-
v0 = similar(L_0_mat, n_fourier * N)
398-
fill!(v0, zero(T))
399-
target_idx = n_max*N + 1
400-
QuantumToolbox.allowed_setindex!(v0, weight, target_idx)
382+
# Build diagonal blocks
383+
n_vals = -1im * ωd * T.(n_list)
384+
diag_blocks = [L + n * sparse(I, N, N) for n in n_vals]
385+
block_diag = blockdiag(diag_blocks...)
386+
M += block_diag
401387

388+
# Initialize RHS vector
389+
v0 = zeros(T, n_fourier * N)
390+
v0[n_max*N+1] = weight
391+
392+
# Preconditioners setup
402393
(haskey(kwargs, :Pl) || haskey(kwargs, :Pr)) && error("The use of preconditioners must be defined in the solver.")
403394
if !isnothing(solver.Pl)
404-
kwargs = merge((; kwargs...), (Pl = solver.Pl(M),))
395+
kwargs = (; kwargs..., Pl = solver.Pl(M))
405396
elseif isa(M, SparseMatrixCSC)
406-
kwargs = merge((; kwargs...), (Pl = ilu(M, τ = 0.01),))
397+
kwargs = (; kwargs..., Pl = ilu(M, τ = 0.01))
398+
end
399+
if !isnothing(solver.Pr)
400+
kwargs = (; kwargs..., Pr = solver.Pr(M))
401+
end
402+
if !haskey(kwargs, :abstol)
403+
kwargs = (; kwargs..., abstol = tol)
404+
end
405+
if !haskey(kwargs, :reltol)
406+
kwargs = (; kwargs..., reltol = tol)
407407
end
408-
!isnothing(solver.Pr) && (kwargs = merge((; kwargs...), (Pr = solver.Pr(M),)))
409-
!haskey(kwargs, :abstol) && (kwargs = merge((; kwargs...), (abstol = tol,)))
410-
!haskey(kwargs, :reltol) && (kwargs = merge((; kwargs...), (reltol = tol,)))
411408

409+
# Solve linear system
412410
prob = LinearProblem(M, v0)
413411
ρtot = solve(prob, solver.alg; kwargs...).u
414412

413+
# Extract ρ0 and normalize
415414
offset1 = n_max * N
416415
offset2 = (n_max + 1) * N
417-
ρ0 = reshape(ρtot[(offset1+1):offset2], Ns, Ns)
416+
ρ0 = Matrix(reshape(ρtot[(offset1+1):offset2], Ns, Ns)) # Explicitly convert to Matrix
418417
ρ0_tr = tr(ρ0)
419-
ρ0 = ρ0 / ρ0_tr
418+
ρ0 ./= ρ0_tr
420419
ρ0 = QuantumObject((ρ0 + ρ0') / 2, type = Operator(), dims = L_0.dimensions)
421-
ρtot = ρtot / ρ0_tr
422420

421+
# Collect higher-order Fourier components
423422
ρ_list = [ρ0]
424423
for i in 0:(n_max-1)
425-
ρi_m = reshape(ρtot[(offset1-(i+1)*N+1):(offset1-i*N)], Ns, Ns)
426-
ρi_m = QuantumObject(ρi_m, type = Operator(), dims = L_0.dimensions)
427-
push!(ρ_list, ρi_m)
424+
idx_range = (offset1-(i+1)*N+1):(offset1-i*N)
425+
ρi_m = Matrix(reshape(ρtot[idx_range], Ns, Ns)) # Explicitly convert to Matrix
426+
push!(ρ_list, QuantumObject(ρi_m, type = Operator(), dims = L_0.dimensions))
428427
end
429428

430429
return ρ_list

‎src/utilities.jl‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -200,5 +200,3 @@ _convert_eltype_wordsize(::Type{T}, ::Val{64}) where {T<:AbstractFloat} = Float6
200200
_convert_eltype_wordsize(::Type{T}, ::Val{32}) where {T<:AbstractFloat} = Float32
201201
_convert_eltype_wordsize(::Type{Complex{T}}, ::Val{64}) where {T<:Union{Int,AbstractFloat}} = ComplexF64
202202
_convert_eltype_wordsize(::Type{Complex{T}}, ::Val{32}) where {T<:Union{Int,AbstractFloat}} = ComplexF32
203-
204-
function allowed_setindex! end

0 commit comments

Comments
 (0)