Add support to dense eigen on GPU

benchmarks/eigenvalues.jl

@@ -17,7 +17,7 @@ function benchmark_eigenvalues!(SUITE)
 
     SUITE["Eigenvalues"]["eigenstates"]["dense"] = @benchmarkable eigenstates($L)
     SUITE["Eigenvalues"]["eigenstates"]["sparse"] =
-        @benchmarkable eigenstates($L, sparse = true, sigma = 0.01, eigvals = 5)
+        @benchmarkable eigenstates($L, sparse = Val(true), sigma = 0.01, eigvals = 5)
 
     return nothing
 end

src/qobj/eigsolve.jl

@@ -29,9 +29,9 @@ One can obtain the eigenvalues and the corresponding [`QuantumObject`](@ref)-typ
 julia> result = eigenstates(sigmax())
 EigsolveResult:   type=Operator()   dims=[2]
 values:
-2-element Vector{ComplexF64}:
- -1.0 + 0.0im
-  1.0 + 0.0im
+2-element Vector{Float64}:
+ -1.0
+  1.0
 vectors:
 2×2 Matrix{ComplexF64}:
  -0.707107+0.0im  0.707107+0.0im
@@ -40,9 +40,9 @@ vectors:
 julia> λ, ψ, U = result;
 
 julia> λ
-2-element Vector{ComplexF64}:
- -1.0 + 0.0im
-  1.0 + 0.0im
+2-element Vector{Float64}:
+ -1.0
+  1.0
 
 julia> ψ
 2-element Vector{QuantumObject{Ket, Dimensions{1, Tuple{Space}}, Vector{ComplexF64}}}:
@@ -64,7 +64,7 @@ julia> U
 ```
 """
 struct EigsolveResult{
-    T1<:Vector{<:Number},
+    T1<:AbstractVector{<:Number},
     T2<:AbstractMatrix{<:Number},
     ObjType<:Union{Nothing,Operator,SuperOperator},
     DimType<:Union{Nothing,AbstractDimensions},
@@ -532,12 +532,12 @@ julia> using LinearAlgebra;
 julia> E, ψ, U = eigen(H)
 EigsolveResult:   type=Operator()   dims=[5]
 values:
-5-element Vector{ComplexF64}:
-       -2.8569700138728 + 0.0im
-    -1.3556261799742608 + 0.0im
- 1.3322676295501878e-15 + 0.0im
-     1.3556261799742677 + 0.0im
-     2.8569700138728056 + 0.0im
+5-element Vector{Float64}:
+ -2.8569700138728
+ -1.3556261799742608
+  1.3322676295501878e-15
+  1.3556261799742677
+  2.8569700138728056
 vectors:
 5×5 Matrix{ComplexF64}:
   0.106101+0.0im  -0.471249-0.0im  …   0.471249+0.0im  0.106101+0.0im
@@ -551,11 +551,11 @@ true
 ```
 """
 function LinearAlgebra.eigen(A::QuantumObject{OpType}; kwargs...) where {OpType<:Union{Operator,SuperOperator}}
-    MT = typeof(A.data)
+    # This creates a weak Union type on CPU. See https://github.com/JuliaLang/LinearAlgebra.jl/issues/1498
     F = eigen(to_dense(A.data); kwargs...)
-    # This fixes a type inference issue. But doesn't work for GPU arrays
-    E::mat2vec(to_dense(MT)) = F.values
-    U::to_dense(MT) = F.vectors
+
+    E = F.values
+    U = F.vectors
     settings.auto_tidyup && tidyup!(U)
 
     return EigsolveResult(E, U, A.type, A.dimensions, 0, 0, true)
@@ -570,51 +570,57 @@ LinearAlgebra.eigvals(A::QuantumObject{OpType}; kwargs...) where {OpType<:Union{
     eigvals(to_dense(A.data); kwargs...)
 
 @doc raw"""
-    eigenenergies(A::QuantumObject; sparse::Bool=false, kwargs...)
+    eigenenergies(A::QuantumObject; sparse::Union{Bool,Val}=Val(false), kwargs...)
 
 Calculate the eigenenergies
 
 # Arguments
 - `A::QuantumObject`: the [`QuantumObject`](@ref) to solve eigenvalues
-- `sparse::Bool`: if `false` call [`eigvals(A::QuantumObject; kwargs...)`](@ref), otherwise call [`eigsolve`](@ref). Default to `false`.
+- `sparse::Union{Bool,Val}`: if `false` call [`eigvals(A::QuantumObject; kwargs...)`](@ref), otherwise call [`eigsolve`](@ref). Default to `Val(false)`.
 - `kwargs`: Additional keyword arguments passed to the solver. If `sparse=true`, the keyword arguments are passed to [`eigsolve`](@ref), otherwise to [`eigen`](@ref).
 
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `eigenenergies(A; sparse=Val(sparse))` instead of `eigenenergies(A; sparse=sparse)`. See the [related Section](@ref doc:Type-Stability) about type stability for more details.
+
 # Returns
 - `::Vector{<:Number}`: a list of eigenvalues
 """
 function eigenenergies(
     A::QuantumObject{OpType};
-    sparse::Bool = false,
+    sparse::Union{Bool,Val} = Val(false),
     kwargs...,
 ) where {OpType<:Union{Operator,SuperOperator}}
-    if !sparse
-        return eigvals(A; kwargs...)
-    else
+    if getVal(sparse)
         return eigsolve(A; kwargs...).values
+    else
+        return eigvals(A; kwargs...)
     end
 end
 
 @doc raw"""
-    eigenstates(A::QuantumObject; sparse::Bool=false, kwargs...)
+    eigenstates(A::QuantumObject; sparse::Union{Bool,Val}=Val(false), kwargs...)
 
 Calculate the eigenvalues and corresponding eigenvectors
 
 # Arguments
 - `A::QuantumObject`: the [`QuantumObject`](@ref) to solve eigenvalues and eigenvectors
-- `sparse::Bool`: if `false` call [`eigen(A::QuantumObject; kwargs...)`](@ref), otherwise call [`eigsolve`](@ref). Default to `false`.
+- `sparse::Union{Bool,Val}`: if `false` call [`eigen(A::QuantumObject; kwargs...)`](@ref), otherwise call [`eigsolve`](@ref). Default to `Val(false)`.
 - `kwargs`: Additional keyword arguments passed to the solver. If `sparse=true`, the keyword arguments are passed to [`eigsolve`](@ref), otherwise to [`eigen`](@ref).
 
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `eigenstates(A; sparse=Val(sparse))` instead of `eigenstates(A; sparse=sparse)`. See the [related Section](@ref doc:Type-Stability) about type stability for more details.
+
 # Returns
 - `::EigsolveResult`: containing the eigenvalues, the eigenvectors, and some information from the solver. see also [`EigsolveResult`](@ref)
 """
 function eigenstates(
     A::QuantumObject{OpType};
-    sparse::Bool = false,
+    sparse::Union{Bool,Val} = Val(false),
     kwargs...,
 ) where {OpType<:Union{Operator,SuperOperator}}
-    if !sparse
-        return eigen(A; kwargs...)
-    else
+    if getVal(sparse)
         return eigsolve(A; kwargs...)
+    else
+        return eigen(A; kwargs...)
     end
 end

src/spectrum.jl

@@ -110,8 +110,7 @@ function _spectrum(
     idxs = findall(x -> abs(x) > solver.tol, amps)
     amps, rates = amps[idxs], rates[idxs]
 
-    # spec = map(ω -> 2 * real(sum(@. amps * (1 / (1im * ω - rates)))), ωlist)
-    amps_rates = zip(amps, rates)
+    amps_rates = zip(amps, complex.(rates))
     spec = map(ω -> 2 * real(sum(x -> x[1] / (1im * ω - x[2]), amps_rates)), ωlist)
 
     return spec

test/core-test/correlations_and_spectrum.jl

@@ -30,12 +30,14 @@
     @test all(spec2 .≈ spec3)
     @test all(spec2 .≈ spec4)
 
-    @testset "Type Inference spectrum" begin
-        @inferred correlation_2op_1t(H, nothing, t_l, c_ops, a', a; progress_bar = Val(false))
-        @inferred spectrum_correlation_fft(t_l, corr1)
-        @inferred spectrum(H, ω_l2, c_ops, a', a)
-        @inferred spectrum(H, ω_l2, c_ops, a', a; solver = PseudoInverse())
-        @inferred spectrum(H, ω_l2, c_ops, a', a; solver = Lanczos())
+    if VERSION >= v"1.11" # eigen is type unstable on v1.10
+        @testset "Type Inference spectrum" begin
+            @inferred correlation_2op_1t(H, nothing, t_l, c_ops, a', a; progress_bar = Val(false))
+            @inferred spectrum_correlation_fft(t_l, corr1)
+            @inferred spectrum(H, ω_l2, c_ops, a', a)
+            @inferred spectrum(H, ω_l2, c_ops, a', a; solver = PseudoInverse())
+            @inferred spectrum(H, ω_l2, c_ops, a', a; solver = Lanczos())
+        end
     end
 
     @testset "Verbose mode Lanczos" begin

test/core-test/eigenvalues_and_operators.jl

@@ -1,19 +1,19 @@
 @testitem "Eigenvalues" begin
     σx = sigmax()
-    result = eigenstates(σx, sparse = false)
+    result = eigenstates(σx, sparse = Val(false))
     λd, ψd, Td = result
     resstring = sprint((t, s) -> show(t, "text/plain", s), result)
     valstring = sprint((t, s) -> show(t, "text/plain", s), result.values)
     vecsstring = sprint((t, s) -> show(t, "text/plain", s), result.vectors)
-    λs, ψs, Ts = eigenstates(σx, sparse = true, eigvals = 2)
-    λs1, ψs1, Ts1 = eigenstates(σx, sparse = true, eigvals = 1)
+    λs, ψs, Ts = eigenstates(σx, sparse = Val(true), eigvals = 2)
+    λs1, ψs1, Ts1 = eigenstates(σx, sparse = Val(true), eigvals = 1)
 
     @test all([ψ.type isa Ket for ψ in ψd])
     @test typeof(Td) <: AbstractMatrix
     @test typeof(Ts) <: AbstractMatrix
     @test typeof(Ts1) <: AbstractMatrix
-    @test all(abs.(eigenenergies(σx, sparse = false)) .≈ abs.(λd))
-    @test all(abs.(eigenenergies(σx, sparse = true, eigvals = 2)) .≈ abs.(λs))
+    @test all(abs.(eigenenergies(σx, sparse = Val(false))) .≈ abs.(λd))
+    @test all(abs.(eigenenergies(σx, sparse = Val(true), eigvals = 2)) .≈ abs.(λs))
     @test resstring ==
           "EigsolveResult:   type=$(Operator())   dims=$(result.dims)\nvalues:\n$(valstring)\nvectors:\n$vecsstring"
 
@@ -33,7 +33,7 @@
 
     vals_d, vecs_d, mat_d = eigenstates(H_d)
     vals_c, vecs_c, mat_c = eigenstates(H_c)
-    vals2, vecs2, mat2 = eigenstates(H_d, sparse = true, sigma = -0.9, eigvals = 10, krylovdim = 30, by = real)
+    vals2, vecs2, mat2 = eigenstates(H_d, sparse = Val(true), sigma = -0.9, eigvals = 10, krylovdim = 30, by = real)
 
     @test real.(vals_d[1:20]) ≈ real.(vals_c[1:20])
     @test real.(vals_d[1:10]) ≈ real.(vals2[1:10])
@@ -67,7 +67,7 @@
     @test vec2mat(vecs[:, 1]) * exp(-1im * angle(vecs[1, 1])) ≈ vec2mat(vecs3[:, 1]) atol=1e-5
 
     # eigen solve for QuantumObject
-    result = eigenstates(L, sparse = true, sigma = 0.01, eigvals = 10, krylovdim = 50)
+    result = eigenstates(L, sparse = Val(true), sigma = 0.01, eigvals = 10, krylovdim = 50)
     vals, vecs = result
     resstring = sprint((t, s) -> show(t, "text/plain", s), result)
     valstring = sprint((t, s) -> show(t, "text/plain", s), result.values)
@@ -105,9 +105,24 @@
         c_ops = [√((1 + n_th) * κ) * a, √κ * b, √(n_th * κ) * a_d]
         L = liouvillian(H, c_ops)
 
-        @inferred eigenstates(H, sparse = false)
-        @inferred eigenstates(H, sparse = true)
-        @inferred eigenstates(L, sparse = true)
+        UnionType = Union{
+            QuantumToolbox.EigsolveResult{
+                Vector{ComplexF64},
+                Matrix{ComplexF64},
+                QuantumToolbox.Operator,
+                QuantumToolbox.Dimensions{2,Tuple{QuantumToolbox.Space,QuantumToolbox.Space}},
+            },
+            QuantumToolbox.EigsolveResult{
+                Vector{Float64},
+                Matrix{ComplexF64},
+                QuantumToolbox.Operator,
+                QuantumToolbox.Dimensions{2,Tuple{QuantumToolbox.Space,QuantumToolbox.Space}},
+            },
+        }
+
+        @inferred UnionType eigenstates(H, sparse = Val(false))
+        @inferred eigenstates(H, sparse = Val(true))
+        @inferred eigenstates(L, sparse = Val(true))
         @inferred eigsolve_al(L, 1 \ (40 * κ), eigvals = 10)
     end
 end

test/core-test/entropy_and_metric.jl

@@ -40,14 +40,16 @@
     @test_throws ArgumentError entropy_mutual(ρAB, (1, 2), 3)
     @test_throws ArgumentError entropy_mutual(ρAB, (1, 2), (1, 3))
 
-    @testset "Type Stability (entropy)" begin
-        @inferred entropy_vn(ρ1)
-        @inferred entropy_vn(ρ1, base = base)
-        @inferred entropy_relative(ρ1, ψ)
-        @inferred entropy_relative(ρ1, σ1, base = base)
-        @inferred entropy_linear(ρ1)
-        @inferred entropy_mutual(ρAB, selA, selB)
-        @inferred entropy_conditional(ρAB, selA)
+    if VERSION >= v"1.11" # eigen is type unstable on v1.10
+        @testset "Type Stability (entropy)" begin
+            @inferred entropy_vn(ρ1)
+            @inferred entropy_vn(ρ1, base = base)
+            @inferred entropy_relative(ρ1, ψ)
+            @inferred entropy_relative(ρ1, σ1, base = base)
+            @inferred entropy_linear(ρ1)
+            @inferred entropy_mutual(ρAB, selA, selB)
+            @inferred entropy_conditional(ρAB, selA)
+        end
     end
 end
 

test/core-test/generalized_master_equation.jl

@@ -44,22 +44,24 @@
 
     @test abs(expect(Xp' * Xp, steadystate(L1)) - n_thermal(1, Tlist[1])) / n_thermal(1, Tlist[1]) < 1e-4
 
-    @testset "Type Inference (liouvillian_generalized)" begin
-        N_c = 30
-        N_trunc = 10
-        tol = 1e-14
+    if VERSION >= v"1.11" # eigen is type unstable on v1.10
+        @testset "Type Inference (liouvillian_generalized)" begin
+            N_c = 30
+            N_trunc = 10
+            tol = 1e-14
 
-        a = kron(destroy(N_c), qeye(2))
-        sm = kron(qeye(N_c), sigmam())
-        sp = sm'
-        sx = sm + sp
-        sz = sp * sm - sm * sp
+            a = kron(destroy(N_c), qeye(2))
+            sm = kron(qeye(N_c), sigmam())
+            sp = sm'
+            sx = sm + sp
+            sz = sp * sm - sm * sp
 
-        H = 1 * a' * a + 1 * sz / 2 + 0.5 * (a + a') * sx
+            H = 1 * a' * a + 1 * sz / 2 + 0.5 * (a + a') * sx
 
-        fields = [sqrt(0.01) * (a + a'), sqrt(0.01) * sx]
-        Tlist = [0, 0.01]
+            fields = [sqrt(0.01) * (a + a'), sqrt(0.01) * sx]
+            Tlist = [0, 0.01]
 
-        @inferred liouvillian_generalized(H, fields, Tlist, N_trunc = N_trunc, tol = tol)
+            @inferred liouvillian_generalized(H, fields, Tlist, N_trunc = N_trunc, tol = tol)
+        end
     end
 end

test/core-test/states_and_operators.jl

@@ -18,7 +18,7 @@
 
     @testset "fock state" begin
         # fock, basis, and fock_dm
-        @test fock_dm(4; dims = (2, 2), sparse = true) ≈ ket2dm(basis(4; dims = (2, 2)))
+        @test fock_dm(4; dims = (2, 2), sparse = Val(true)) ≈ ket2dm(basis(4; dims = (2, 2)))
         @test_throws DimensionMismatch fock(4; dims = 2)
         @test_throws ArgumentError fock(4, 4)
     end
@@ -37,7 +37,7 @@
 
     @testset "thermal state" begin
         ρTd = thermal_dm(5, 0.123)
-        ρTs = thermal_dm(5, 0.123; sparse = true)
+        ρTs = thermal_dm(5, 0.123; sparse = Val(true))
         @test isoper(ρTd)
         @test ρTd.dims == [5]
         @test tr(ρTd) ≈ 1.0
@@ -180,7 +180,7 @@
     end
 
     @testset "tunneling" begin
-        @test tunneling(10, 2) == tunneling(10, 2; sparse = true)
+        @test tunneling(10, 2) == tunneling(10, 2; sparse = Val(true))
         @test_throws ArgumentError tunneling(10, 0)
     end
 

test/ext-test/cpu/arbitrary_precision/arbitrary_precision.jl

@@ -43,8 +43,8 @@
         L = liouvillian(H, c_ops)
         L_big = liouvillian(H_big, c_ops_big)
 
-        vals, vecs = eigenstates(L; sparse = true, sigma = 0.01, eigvals = 7, krylovdim = 30)
-        vals_big, vecs_big = eigenstates(L_big; sparse = true, sigma = 0.01, eigvals = 7, krylovdim = 30)
+        vals, vecs = eigenstates(L; sparse = Val(true), sigma = 0.01, eigvals = 7, krylovdim = 30)
+        vals_big, vecs_big = eigenstates(L_big; sparse = Val(true), sigma = 0.01, eigvals = 7, krylovdim = 30)
 
         # Align eigenvalues
         idxs = [findmin(abs.(vals_big .- val))[2] for val in vals]

test/ext-test/gpu/cuda_ext.jl

@@ -274,16 +274,27 @@ end
 
     a = destroy(N)
     H = Δ * a' * a + U / 2 * a' * a' * a * a + F * (a + a')
+    H_gpu = cu(H)
 
     c_ops = [sqrt(κ) * a]
 
     L = liouvillian(H, c_ops)
     L_gpu = CuSparseMatrixCSR(L)
 
-    vals_cpu, vecs_cpu = eigenstates(L; sparse = true, sigma = 0.01, eigvals = 4, krylovdim = 30)
+    # Dense eigen solver for Hamiltonian
+    vals_H_cpu, vecs_H_cpu = eigenstates(H)
+    vals_H_gpu, vecs_H_gpu = eigenstates(H_gpu)
+
+    @test vals_H_cpu ≈ Array(vals_H_gpu) atol = 1e-8
+    @test all(zip(vecs_H_cpu, vecs_H_gpu)) do (v_cpu, v_gpu)
+        return isapprox(abs(dot(v_cpu.data, Array(v_gpu.data))), 1; atol = 1e-8)
+    end
+
+    # Sparse eigen solver for Liouvillian
+    vals_cpu, vecs_cpu = eigenstates(L; sparse = Val(true), sigma = 0.01, eigvals = 4, krylovdim = 30)
     vals_gpu, vecs_gpu = eigenstates(
         L_gpu;
-        sparse = true,
+        sparse = Val(true),
         sigma = 0.01,
         eigvals = 4,
         krylovdim = 30,