Avoid to forcing the type of eigen result, and use Val to choose between sparse or dense solvers

Author: albertomercurio

src/qobj/eigsolve.jl

@@ -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},
@@ -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

test/core-test/eigenvalues_and_operators.jl

@@ -105,9 +105,11 @@
         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/ext-test/gpu/Project.toml

@@ -4,5 +4,8 @@ CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 QuantumToolbox = "6c2fb7c5-b903-41d2-bc5e-5a7c320b9fab"
 
+[sources]
+QuantumToolbox = {path = "/home/alberto/.julia/dev/QuantumToolbox"}
+
 [compat]
 CUDA = "5"