Skip to content

Commit 08bb40d

Browse files
committed
Scalar indexing is disallowed is Cuda.jl so use allowscalar
1 parent 11525d5 commit 08bb40d

File tree

1 file changed

+6
-12
lines changed

1 file changed

+6
-12
lines changed

src/steadystate.jl

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -369,8 +369,7 @@ function _steadystate_fourier(
369369
n_fourier = 2 * n_max + 1
370370
n_list = (-n_max):n_max
371371

372-
# Create arrays with proper GPU/CPU backend
373-
weight = one(T)
372+
weight = 1
374373
rows = _dense_similar(L_0_mat, Ns)
375374
cols = _dense_similar(L_0_mat, Ns)
376375
vals = _dense_similar(L_0_mat, Ns)
@@ -382,7 +381,6 @@ function _steadystate_fourier(
382381
Mn = _sparse_similar(L_0_mat, rows, cols, vals, N, N)
383382
L = L_0_mat + Mn
384383

385-
# Initialize the big matrix M with proper backend
386384
M = _sparse_similar(L_0_mat, n_fourier * N, n_fourier * N)
387385

388386
# Add superdiagonal blocks (L_m)
@@ -391,7 +389,7 @@ function _steadystate_fourier(
391389
cols_block = _dense_similar(L_0_mat, N)
392390
fill!(rows_block, i)
393391
fill!(cols_block, i+1)
394-
block = _sparse_similar(L_0_mat, rows_block, cols_block, ones(T, N), n_fourier, n_fourier)
392+
block = _sparse_similar(L_0_mat, rows_block, cols_block, ones(N), n_fourier, n_fourier)
395393
M += kron(block, L_m_mat)
396394
end
397395

@@ -401,7 +399,7 @@ function _steadystate_fourier(
401399
cols_block = _dense_similar(L_0_mat, N)
402400
fill!(rows_block, i+1)
403401
fill!(cols_block, i)
404-
block = _sparse_similar(L_0_mat, rows_block, cols_block, ones(T, N), n_fourier, n_fourier)
402+
block = _sparse_similar(L_0_mat, rows_block, cols_block, ones(N), n_fourier, n_fourier)
405403
M += kron(block, L_p_mat)
406404
end
407405

@@ -413,16 +411,14 @@ function _steadystate_fourier(
413411
cols_block = _dense_similar(L_0_mat, N)
414412
fill!(rows_block, i)
415413
fill!(cols_block, i)
416-
block = _sparse_similar(L_0_mat, rows_block, cols_block, ones(T, N), n_fourier, n_fourier)
414+
block = _sparse_similar(L_0_mat, rows_block, cols_block, ones(N), n_fourier, n_fourier)
417415
M += kron(block, block_diag)
418416
end
419417

420-
# Initialize solution vector with proper backend
421418
v0 = _dense_similar(L_0_mat, n_fourier * N)
422-
fill!(v0, zero(T))
423-
allowed_setindex!(v0, weight, n_max * N + 1)
419+
fill!(v0, 0)
420+
@allowscalar v0[n_max * N + 1] = weight
424421

425-
# Prepare preconditioners if needed
426422
if !isnothing(solver.Pl)
427423
kwargs = merge((; kwargs...), (Pl = solver.Pl(M),))
428424
elseif isa(M, SparseMatrixCSC)
@@ -431,11 +427,9 @@ function _steadystate_fourier(
431427
!isnothing(solver.Pr) && (kwargs = merge((; kwargs...), (Pr = solver.Pr(M),)))
432428
kwargs = merge((abstol = tol, reltol = tol), kwargs)
433429

434-
# Solve the linear system
435430
prob = LinearProblem(M, v0)
436431
ρtot = solve(prob, solver.alg; kwargs...).u
437432

438-
# Extract results
439433
offset1 = n_max * N
440434
offset2 = (n_max + 1) * N
441435
ρ0 = reshape(ρtot[(offset1+1):offset2], Ns, Ns)

0 commit comments

Comments
 (0)