Fix type instabilities for almost all functions (qutip#221)

* Improve c_ops handling

* Format code

* Fix ptrace and operators

* Make states stable

* Fix type instabilities for qobj methods

* FIx eigenvalues

* Other fixes and format

* Minor changes to dfd_mesolve

* Fix typo

* Remove version condition of runtest

Author: albertomercurio

Project.toml

@@ -11,7 +11,6 @@ FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-LinearMaps = "7a12625a-238d-50fd-b39a-03d52299707e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
@@ -39,7 +38,6 @@ FFTW = "1.5"
 Graphs = "1.7"
 IncompleteLU = "0.2"
 LinearAlgebra = "<0.0.1, 1"
-LinearMaps = "3"
 LinearSolve = "2"
 OrdinaryDiffEqCore = "1"
 OrdinaryDiffEqTsit5 = "1"

docs/src/api.md

@@ -219,6 +219,12 @@ wigner
 negativity
 ```
 
+## [Linear Maps](@id doc-API:Linear-Maps)
+
+```@docs
+AbstractLinearMap
+```
+
 ## [Utility functions](@id doc-API:Utility-functions)
 
 ```@docs

docs/src/users_guide/QuantumObject/QuantumObject.md

@@ -45,8 +45,8 @@ Qobj(M, dims = (2, 2))  # dims as Tuple (recommended)
 Qobj(M, dims = SVector(2, 2)) # dims as StaticArrays.SVector (recommended)
 ```
 
-> [!IMPORTANT]
-> Please note that here we put the `dims` as a tuple `(2, 2)`. Although it supports also `Vector` type (`dims = [2, 2]`), it is recommended to use `Tuple` or `SVector` from [`StaticArrays.jl`](https://github.com/JuliaArrays/StaticArrays.jl) to improve performance. For a brief explanation on the impact of the type of `dims`, see the Section [The Importance of Type-Stability](@ref doc:Type-Stability).
+!!! warning "Beware of type-stability!"
+    Please note that here we put the `dims` as a tuple `(2, 2)`. Although it supports also `Vector` type (`dims = [2, 2]`), it is recommended to use `Tuple` or `SVector` from [`StaticArrays.jl`](https://github.com/JuliaArrays/StaticArrays.jl) to improve performance. For a brief explanation on the impact of the type of `dims`, see the Section [The Importance of Type-Stability](@ref doc:Type-Stability).
 
 ```@example Qobj
 Qobj(rand(4, 4), type = SuperOperator)

src/QuantumToolbox.jl

@@ -40,7 +40,6 @@ import ArrayInterface: allowed_getindex, allowed_setindex!
 import FFTW: fft, fftshift
 import Graphs: connected_components, DiGraph
 import IncompleteLU: ilu
-import LinearMaps: LinearMap
 import Pkg
 import Random
 import SpecialFunctions: loggamma
@@ -54,6 +53,7 @@ BLAS.set_num_threads(1)
 include("utilities.jl")
 include("versioninfo.jl")
 include("progress_bar.jl")
+include("linear_maps.jl")
 
 # Quantum Object
 include("qobj/quantum_object.jl")

src/linear_maps.jl

@@ -0,0 +1,54 @@
+export AbstractLinearMap
+
+@doc raw"""
+    AbstractLinearMap{T, TS}
+
+Represents a general linear map with element type `T` and size `TS`.
+
+## Overview
+
+A **linear map** is a transformation `L` that satisfies:
+
+- **Additivity**: 
+    ```math
+    L(u + v) = L(u) + L(v)
+    ```
+- **Homogeneity**: 
+    ```math
+    L(cu) = cL(u)
+    ```
+
+It is typically represented as a matrix with dimensions given by `size`, and this abtract type helps to define this map when the matrix is not explicitly available.
+
+## Methods
+
+- `Base.eltype(A)`: Returns the element type `T`.
+- `Base.size(A)`: Returns the size `A.size`.
+- `Base.size(A, i)`: Returns the `i`-th dimension.
+
+## Example
+
+As an example, we now define the linear map used in the [`eigsolve_al`](@ref) function for Arnoldi-Lindblad eigenvalue solver:
+
+```julia-repl
+struct ArnoldiLindbladIntegratorMap{T,TS,TI} <: AbstractLinearMap{T,TS}
+    elty::Type{T}
+    size::TS
+    integrator::TI
+end
+
+function LinearAlgebra.mul!(y::AbstractVector, A::ArnoldiLindbladIntegratorMap, x::AbstractVector)
+    reinit!(A.integrator, x)
+    solve!(A.integrator)
+    return copyto!(y, A.integrator.u)
+end
+```
+
+where `integrator` is the ODE integrator for the time-evolution. In this way, we can diagonalize this linear map using the [`eigsolve`](@ref) function.
+"""
+abstract type AbstractLinearMap{T,TS} end
+
+Base.eltype(A::AbstractLinearMap{T}) where {T} = T
+
+Base.size(A::AbstractLinearMap) = A.size
+Base.size(A::AbstractLinearMap, i::Int) = A.size[i]

src/metrics.jl

@@ -53,12 +53,12 @@ function entropy_vn(
     base::Int = 0,
     tol::Real = 1e-15,
 ) where {T}
-    vals = eigvals(ρ)
-    indexes = abs.(vals) .> tol
-    1 ∉ indexes && return 0
+    vals = eigenenergies(ρ)
+    indexes = findall(x -> abs(x) > tol, vals)
+    length(indexes) == 0 && return zero(real(T))
     nzvals = vals[indexes]
     logvals = base != 0 ? log.(base, Complex.(nzvals)) : log.(Complex.(nzvals))
-    return -real(sum(nzvals .* logvals))
+    return -real(mapreduce(*, +, nzvals, logvals))
 end
 
 """

src/qobj/arithmetic_and_attributes.jl

@@ -239,8 +239,10 @@ julia> tr(a' * a)
 """
 LinearAlgebra.tr(
     A::QuantumObject{<:AbstractArray{T},OpType},
-) where {T,OpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject}} =
-    ishermitian(A) ? real(tr(A.data)) : tr(A.data)
+) where {T,OpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject}} = tr(A.data)
+LinearAlgebra.tr(
+    A::QuantumObject{<:Union{<:Hermitian{TF},Symmetric{TR}},OpType},
+) where {TF<:BlasFloat,TR<:Real,OpType<:OperatorQuantumObject} = real(tr(A.data))
 
 @doc raw"""
     svdvals(A::QuantumObject)
@@ -515,7 +517,7 @@ function ptrace(QO::QuantumObject{<:AbstractArray,KetQuantumObject}, sel::Union{
     length(QO.dims) == 1 && return QO
 
     ρtr, dkeep = _ptrace_ket(QO.data, QO.dims, SVector(sel))
-    return QuantumObject(ρtr, dims = dkeep)
+    return QuantumObject(ρtr, type = Operator, dims = dkeep)
 end
 
 ptrace(QO::QuantumObject{<:AbstractArray,BraQuantumObject}, sel::Union{AbstractVector{Int},Tuple}) = ptrace(QO', sel)
@@ -524,7 +526,7 @@ function ptrace(QO::QuantumObject{<:AbstractArray,OperatorQuantumObject}, sel::U
     length(QO.dims) == 1 && return QO
 
     ρtr, dkeep = _ptrace_oper(QO.data, QO.dims, SVector(sel))
-    return QuantumObject(ρtr, dims = dkeep)
+    return QuantumObject(ρtr, type = Operator, dims = dkeep)
 end
 ptrace(QO::QuantumObject, sel::Int) = ptrace(QO, SVector(sel))
 
@@ -547,7 +549,7 @@ function _ptrace_ket(QO::AbstractArray, dims::Union{SVector,MVector}, sel)
 
     vmat = reshape(QO, reverse(dims)...)
     topermute = nd + 1 .- sel_qtrace
-    vmat = PermutedDimsArray(vmat, topermute)
+    vmat = permutedims(vmat, topermute) # TODO: use PermutedDimsArray when Julia v1.11.0 is released
     vmat = reshape(vmat, prod(dkeep), prod(dtrace))
 
     return vmat * vmat', dkeep
@@ -576,14 +578,14 @@ function _ptrace_oper(QO::AbstractArray, dims::Union{SVector,MVector}, sel)
 
     ρmat = reshape(QO, reverse(vcat(dims, dims))...)
     topermute = 2 * nd + 1 .- qtrace_sel
-    ρmat = PermutedDimsArray(ρmat, reverse(topermute))
+    ρmat = permutedims(ρmat, reverse(topermute)) # TODO: use PermutedDimsArray when Julia v1.11.0 is released
 
     ## TODO: Check if it works always
 
     # ρmat = row_major_reshape(ρmat, prod(dtrace), prod(dtrace), prod(dkeep), prod(dkeep))
     # res = dropdims(mapslices(tr, ρmat, dims=(1,2)), dims=(1,2))
     ρmat = reshape(ρmat, prod(dkeep), prod(dkeep), prod(dtrace), prod(dtrace))
-    res = dropdims(mapslices(tr, ρmat, dims = (3, 4)), dims = (3, 4))
+    res = map(tr, eachslice(ρmat, dims = (1, 2)))
 
     return res, dkeep
 end
@@ -673,6 +675,9 @@ julia> ψ_123 = tensor(ψ1, ψ2, ψ3)
 julia> permute(ψ_123, [2, 1, 3]) ≈ tensor(ψ2, ψ1, ψ3)
 true
 ```
+
+!!! warning "Beware of type-stability!"
+    It is highly recommended to use `permute(A, order)` with `order` as `Tuple` or `SVector` to keep type stability. See the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 function permute(
     A::QuantumObject{<:AbstractArray{T},ObjType},
@@ -691,7 +696,7 @@ function permute(
     dims, perm = _dims_and_perm(A.type, A.dims, order_svector, length(order_svector))
 
     return QuantumObject(
-        reshape(PermutedDimsArray(reshape(A.data, dims...), Tuple(perm)), size(A)),
+        reshape(permutedims(reshape(A.data, dims...), Tuple(perm)), size(A)),
         A.type,
         A.dims[order_svector],
     )

src/qobj/eigsolve.jl

@@ -84,9 +84,27 @@ function Base.show(io::IO, res::EigsolveResult)
     return show(io, MIME("text/plain"), res.vectors)
 end
 
-function _map_ldiv(linsolve, y, x)
-    linsolve.b .= x
-    return y .= solve!(linsolve).u
+struct ArnoldiLindbladIntegratorMap{T,TS,TI} <: AbstractLinearMap{T,TS}
+    elty::Type{T}
+    size::TS
+    integrator::TI
+end
+
+function LinearAlgebra.mul!(y::AbstractVector, A::ArnoldiLindbladIntegratorMap, x::AbstractVector)
+    reinit!(A.integrator, x)
+    solve!(A.integrator)
+    return copyto!(y, A.integrator.u)
+end
+
+struct EigsolveInverseMap{T,TS,TI} <: AbstractLinearMap{T,TS}
+    elty::Type{T}
+    size::TS
+    linsolve::TI
+end
+
+function LinearAlgebra.mul!(y::AbstractVector, A::EigsolveInverseMap, x::AbstractVector)
+    A.linsolve.b .= x
+    return copyto!(y, solve!(A.linsolve).u)
 end
 
 function _permuteschur!(
@@ -293,7 +311,8 @@ function eigsolve(
 
         prob = LinearProblem(Aₛ, v0)
         linsolve = init(prob, solver; kwargs2...)
-        Amap = LinearMap{T}((y, x) -> _map_ldiv(linsolve, y, x), length(v0))
+
+        Amap = EigsolveInverseMap(T, size(A), linsolve)
 
         res = _eigsolve(Amap, v0, type, dims, k, krylovdim, tol = tol, maxiter = maxiter)
         vals = @. (1 + sigma * res.values) / res.values
@@ -370,13 +389,7 @@ function eigsolve_al(
 
     # prog = ProgressUnknown(desc="Applications:", showspeed = true, enabled=progress)
 
-    function arnoldi_lindblad_solve(ρ)
-        reinit!(integrator, ρ)
-        solve!(integrator)
-        return integrator.u
-    end
-
-    Lmap = LinearMap{eltype(MT1)}(arnoldi_lindblad_solve, size(L, 1), ismutating = false)
+    Lmap = ArnoldiLindbladIntegratorMap(eltype(MT1), size(L), integrator)
 
     res = _eigsolve(Lmap, mat2vec(ρ0), L.type, L.dims, k, krylovdim, maxiter = maxiter, tol = eigstol)
     # finish!(prog)

src/qobj/functions.jl

@@ -25,7 +25,7 @@ If `ψ` is a [`Ket`](@ref) or [`Bra`](@ref), the function calculates ``\langle\p
 
 If `ψ` is a density matrix ([`Operator`](@ref)), the function calculates ``\textrm{Tr} \left[ \hat{O} \hat{\psi} \right]``
 
-The function returns a real number if `O` is hermitian, and returns a complex number otherwise.
+The function returns a real number if `O` is of `Hermitian` type or `Symmetric` type, and returns a complex number otherwise. You can make an operator `O` hermitian by using `Hermitian(O)`.
 
 # Examples
 
@@ -34,31 +34,48 @@ julia> ψ = 1 / √2 * (fock(10,2) + fock(10,4));
 
 julia> a = destroy(10);
 
-julia> expect(a' * a, ψ) ≈ 3
-true
+julia> expect(a' * a, ψ) |> round
+3.0 + 0.0im
+
+julia> expect(Hermitian(a' * a), ψ) |> round
+3.0
 ```
 """
 function expect(
     O::QuantumObject{<:AbstractArray{T1},OperatorQuantumObject},
     ψ::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
 ) where {T1,T2}
-    ψd = ψ.data
-    Od = O.data
-    return ishermitian(O) ? real(dot(ψd, Od, ψd)) : dot(ψd, Od, ψd)
+    return dot(ψ.data, O.data, ψ.data)
 end
 function expect(
     O::QuantumObject{<:AbstractArray{T1},OperatorQuantumObject},
     ψ::QuantumObject{<:AbstractArray{T2},BraQuantumObject},
 ) where {T1,T2}
-    ψd = ψ.data'
-    Od = O.data
-    return ishermitian(O) ? real(dot(ψd, Od, ψd)) : dot(ψd, Od, ψd)
+    return expect(O, ψ')
 end
 function expect(
     O::QuantumObject{<:AbstractArray{T1},OperatorQuantumObject},
     ρ::QuantumObject{<:AbstractArray{T2},OperatorQuantumObject},
 ) where {T1,T2}
-    return ishermitian(O) ? real(tr(O * ρ)) : tr(O * ρ)
+    return tr(O * ρ)
+end
+function expect(
+    O::QuantumObject{<:Union{<:Hermitian{TF},<:Symmetric{TR}},OperatorQuantumObject},
+    ψ::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
+) where {TF<:Number,TR<:Real,T2}
+    return real(dot(ψ.data, O.data, ψ.data))
+end
+function expect(
+    O::QuantumObject{<:Union{<:Hermitian{TF},<:Symmetric{TR}},OperatorQuantumObject},
+    ψ::QuantumObject{<:AbstractArray{T2},BraQuantumObject},
+) where {TF<:Number,TR<:Real,T2}
+    return real(expect(O, ψ'))
+end
+function expect(
+    O::QuantumObject{<:Union{<:Hermitian{TF},<:Symmetric{TR}},OperatorQuantumObject},
+    ρ::QuantumObject{<:AbstractArray{T2},OperatorQuantumObject},
+) where {TF<:Number,TR<:Real,T2}
+    return real(tr(O * ρ))
 end
 
 @doc raw"""