enhancements

Author: Fe-r-oz

ext/QuantumToolboxCUDAExt.jl

@@ -2,8 +2,8 @@ module QuantumToolboxCUDAExt
 
 using QuantumToolbox
 using QuantumToolbox: makeVal, getVal
-import QuantumToolbox: _sparse_similar, _convert_eltype_wordsize, allowed_setindex!
-import CUDA: cu, CuArray, allowscalar, @allowscalar, has_cuda
+import QuantumToolbox: _sparse_similar, _convert_eltype_wordsize
+import CUDA: cu, CuArray, allowscalar
 import CUDA.CUSPARSE: CuSparseVector, CuSparseMatrixCSC, CuSparseMatrixCSR, AbstractCuSparseArray
 import SparseArrays: SparseVector, SparseMatrixCSC, sparse, spzeros
 import CUDA.Adapt: adapt
@@ -104,11 +104,22 @@ QuantumToolbox.to_dense(::Type{T}, A::AbstractCuSparseArray) where {T<:Number} =
 
 QuantumToolbox._sparse_similar(A::CuSparseMatrixCSC, args...) = sparse(args..., fmt = :csc)
 QuantumToolbox._sparse_similar(A::CuSparseMatrixCSR, args...) = sparse(args..., fmt = :csr)
-QuantumToolbox._sparse_similar(A::CuSparseMatrixCSC, I::AbstractVector, J::AbstractVector, V::AbstractVector, m::Int, n::Int) =
-    CuSparseMatrixCSC(sparse(I, J, V, m, n))
+QuantumToolbox._sparse_similar(
+    A::CuSparseMatrixCSC,
+    I::AbstractVector,
+    J::AbstractVector,
+    V::AbstractVector,
+    m::Int,
+    n::Int,
+) = CuSparseMatrixCSC(sparse(I, J, V, m, n))
 QuantumToolbox._sparse_similar(A::CuSparseMatrixCSC, m::Int, n::Int) = CuSparseMatrixCSC(spzeros(eltype(A), m, n))
-QuantumToolbox._sparse_similar(A::CuSparseMatrixCSR, I::AbstractVector, J::AbstractVector, V::AbstractVector, m::Int, n::Int) =
-    CuSparseMatrixCSR(sparse(I, J, V, m, n))
+QuantumToolbox._sparse_similar(
+    A::CuSparseMatrixCSR,
+    I::AbstractVector,
+    J::AbstractVector,
+    V::AbstractVector,
+    m::Int,
+    n::Int,
+) = CuSparseMatrixCSR(sparse(I, J, V, m, n))
 QuantumToolbox._sparse_similar(A::CuSparseMatrixCSR, m::Int, n::Int) = CuSparseMatrixCSR(spzeros(eltype(A), m, n))
-QuantumToolbox.allowed_setindex!(A::AbstractCuSparseArray, v, I...) = @allowscalar A[I...] = v
 end

src/steadystate.jl

@@ -345,86 +345,113 @@ function steadystate_fourier(
     return _steadystate_fourier(L_0, L_p, L_m, ωd, solver; n_max = n_max, tol = tol, kwargs...)
 end
 
+import SparseArrays: sparse, spdiagm, I
+
 function _steadystate_fourier(
-    L_0::QuantumObject{SuperOperator},
-    L_p::QuantumObject{SuperOperator},
-    L_m::QuantumObject{SuperOperator},
+    L0::QuantumObject{SuperOperator},
+    Lp::QuantumObject{SuperOperator},
+    Lm::QuantumObject{SuperOperator},
     ωd::Number,
     solver::SteadyStateLinearSolver;
     n_max::Integer = 1,
     tol::R = 1e-8,
     kwargs...,
 ) where {R<:Real}
-    T1 = eltype(L_0)
-    T2 = eltype(L_p)
-    T3 = eltype(L_m)
-    T = promote_type(T1, T2, T3)
+    A0 = get_data(L0)
+    Ap = get_data(Lp)
+    Am = get_data(Lm)
 
-    L_0_mat = get_data(L_0)
-    L_p_mat = get_data(L_p)
-    L_m_mat = get_data(L_m)
-
-    N = size(L_0_mat, 1)
+    N = size(A0, 1)
     Ns = isqrt(N)
-    n_fourier = 2 * n_max + 1
+    n_fourier = 2*n_max + 1
     n_list = (-n_max):n_max
 
+    T = promote_type(eltype(A0), eltype(Ap), eltype(Am))
     weight = one(T)
-    rows = ones(Int, Ns)
-    cols = [Ns * (j - 1) + j for j in 1:Ns]
-    vals = fill(weight, Ns)
-    Mn = _sparse_similar(L_0_mat, rows, cols, vals, N, N)
-    L = L_0_mat + Mn
-
-    M = _sparse_similar(L_0_mat, n_fourier * N, n_fourier * N)
-    for i in 1:(n_fourier-1)
-        M_block = _sparse_similar(L_0_mat, N, N)
-        copyto!(M_block, L_p_mat)
-        M[(i*N+1):((i+1)*N), ((i-1)*N+1):(i*N)] = M_block
-        M_block = _sparse_similar(L_0_mat, N, N)
-        copyto!(M_block, L_m_mat)
-        M[((i-1)*N+1):(i*N), (i*N+1):((i+1)*N)] = M_block
-    end
 
-    for i in 1:n_fourier
-        n = n_list[i]
-        M_block = _sparse_similar(L_0_mat, N, N)
-        copyto!(M_block, L)
-        M_block -= 1im * ωd * n * sparse(I, N, N)
-        M[((i-1)*N+1):(i*N), ((i-1)*N+1):(i*N)] = M_block
+    # Stabilization block Mn = weight * I (size N)
+    rows = _dense_similar(A0, Ns)
+    cols = _dense_similar(A0, Ns)
+    vals = _dense_similar(A0, Ns)
+    fill!(rows, 1)
+    for j in 1:Ns
+        cols[j] = Ns*(j-1) + j
+        vals[j] = weight
+    end
+    Mn = _sparse_similar(A0, rows, cols, vals, N, N)
+
+    # Base Liouvillian L = A0 + Mn
+    L = A0 + Mn
+
+    # Build off-diagonal blocks via Kron
+    # Kp shifts down, Km shifts up in Fourier index
+    Kp = _sparse_similar(A0, spdiagm(-1 => ones(T, n_fourier-1)))
+    Km = _sparse_similar(A0, spdiagm(1 => ones(T, n_fourier-1)))
+    M = kron(Kp, Ap) + kron(Km, Am)
+
+    # Precompute identity block on N×N
+    # identity in same sparse format as A0
+    Id_vec = _dense_similar(A0, N)
+    fill!(Id_vec, one(T))
+    Id = _sparse_similar(A0, spdiagm(0 => Id_vec))
+
+    # Build diagonal blocks via Kron
+    for (i, n) in enumerate(n_list)
+        # projector onto Fourier index i
+        Pi = sparse([i], [i], one(T), n_fourier, n_fourier)
+        Pi_blk = _sparse_similar(A0, Pi)
+        # block = L - i*ωd*n*I_N
+        D = L - (1im * ωd * n) * Id
+        D_blk = _sparse_similar(A0, D)
+        M += kron(Pi_blk, D_blk)
     end
 
-    v0 = similar(L_0_mat, n_fourier * N)
+    # Initialize RHS vector v0 v0
+    v0 = _dense_similar(A0, n_fourier * N)
     fill!(v0, zero(T))
-    target_idx = n_max*N + 1
-    QuantumToolbox.allowed_setindex!(v0, weight, target_idx)
+    tidx = n_max*N + 1
+    allowed_setindex!(v0, weight, tidx)
 
+    # Preconditioners
     (haskey(kwargs, :Pl) || haskey(kwargs, :Pr)) && error("The use of preconditioners must be defined in the solver.")
     if !isnothing(solver.Pl)
-        kwargs = merge((; kwargs...), (Pl = solver.Pl(M),))
+        kwargs = (; kwargs..., Pl = solver.Pl(M))
     elseif isa(M, SparseMatrixCSC)
-        kwargs = merge((; kwargs...), (Pl = ilu(M, τ = 0.01),))
+        kwargs = (; kwargs..., Pl = ilu(M, τ = 0.01))
+    end
+    if !isnothing(solver.Pr)
+        kwargs = (; kwargs..., Pr = solver.Pr(M))
+    end
+    if !haskey(kwargs, :abstol)
+        kwargs = (; kwargs..., abstol = tol)
+    end
+    if !haskey(kwargs, :reltol)
+        kwargs = (; kwargs..., reltol = tol)
     end
-    !isnothing(solver.Pr) && (kwargs = merge((; kwargs...), (Pr = solver.Pr(M),)))
-    !haskey(kwargs, :abstol) && (kwargs = merge((; kwargs...), (abstol = tol,)))
-    !haskey(kwargs, :reltol) && (kwargs = merge((; kwargs...), (reltol = tol,)))
 
+    # Solve linear system
+    prob = LinearProblem(M, v0)
+    ρtot = solve(prob, solver.alg; kwargs...).u
+    prob = LinearProblem(M, v0)
+    ρtot = solve(prob, solver.alg; kwargs...).u
     prob = LinearProblem(M, v0)
     ρtot = solve(prob, solver.alg; kwargs...).u
 
+    # Extract ρ0 and normalize
     offset1 = n_max * N
     offset2 = (n_max + 1) * N
-    ρ0 = reshape(ρtot[(offset1+1):offset2], Ns, Ns)
-    ρ0_tr = tr(ρ0)
-    ρ0 = ρ0 / ρ0_tr
-    ρ0 = QuantumObject((ρ0 + ρ0') / 2, type = Operator(), dims = L_0.dimensions)
-    ρtot = ρtot / ρ0_tr
+    blk0 = reshape(ρtot[(offset1+1):offset2], Ns, Ns)
+    tr0 = tr(blk0)
+    blk0 ./= tr0
+    ρ0 = QuantumObject((blk0 + blk0')/2, type = Operator(), dims = L0.dimensions)
 
+    # Collect higher-order Fourier components
     ρ_list = [ρ0]
     for i in 0:(n_max-1)
-        ρi_m = reshape(ρtot[(offset1-(i+1)*N+1):(offset1-i*N)], Ns, Ns)
-        ρi_m = QuantumObject(ρi_m, type = Operator(), dims = L_0.dimensions)
-        push!(ρ_list, ρi_m)
+        idx1 = offset1 - (i+1)*N + 1
+        idx2 = offset1 - i*N
+        blk = reshape(ρtot[idx1:idx2], Ns, Ns)
+        push!(ρ_list, QuantumObject(blk, type = Operator(), dims = L0.dimensions))
     end
 
     return ρ_list

src/utilities.jl

@@ -200,5 +200,3 @@ _convert_eltype_wordsize(::Type{T}, ::Val{64}) where {T<:AbstractFloat} = Float6
 _convert_eltype_wordsize(::Type{T}, ::Val{32}) where {T<:AbstractFloat} = Float32
 _convert_eltype_wordsize(::Type{Complex{T}}, ::Val{64}) where {T<:Union{Int,AbstractFloat}} = ComplexF64
 _convert_eltype_wordsize(::Type{Complex{T}}, ::Val{32}) where {T<:Union{Int,AbstractFloat}} = ComplexF32
-
-function allowed_setindex! end