From 6add50d8216757f8ee311fe2b7de544873869465 Mon Sep 17 00:00:00 2001
From: Alberto Mercurio <61953577+albertomercurio@users.noreply.github.com>
Date: Mon, 9 Sep 2024 12:33:17 +0200
Subject: [PATCH] Fix type instabilities for almost all functions (#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
---
 Project.toml                                  |   2 -
 docs/src/api.md                               |   6 +
 .../QuantumObject/QuantumObject.md            |   4 +-
 src/QuantumToolbox.jl                         |   2 +-
 src/linear_maps.jl                            |  54 +++++
 src/metrics.jl                                |   8 +-
 src/qobj/arithmetic_and_attributes.jl         |  21 +-
 src/qobj/eigsolve.jl                          |  35 ++-
 src/qobj/functions.jl                         |  37 ++-
 src/qobj/operators.jl                         |  42 +++-
 src/qobj/states.jl                            |  69 ++++--
 src/steadystate.jl                            |  16 +-
 src/time_evolution/time_evolution.jl          |  27 ++-
 .../time_evolution_dynamical.jl               |  55 ++---
 test/eigenvalues_and_operators.jl             |  25 +-
 test/entanglement.jl                          |   5 +
 test/generalized_master_equation.jl           |  21 +-
 test/negativity_and_partial_transpose.jl      |  10 +
 test/permutation.jl                           |   7 +-
 test/quantum_objects.jl                       | 222 ++++++++++++++----
 test/runtests.jl                              |   2 +-
 test/states_and_operators.jl                  |  90 ++++++-
 test/steady_state.jl                          |  42 +++-
 ...and_partial_trace.jl => time_evolution.jl} |   2 -
 test/wigner.jl                                |   7 +
 25 files changed, 641 insertions(+), 170 deletions(-)
 create mode 100644 src/linear_maps.jl
 rename test/{time_evolution_and_partial_trace.jl => time_evolution.jl} (97%)

diff --git a/Project.toml b/Project.toml
index 3ea5a0b5..96181a46 100644
--- a/Project.toml
+++ b/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"
diff --git a/docs/src/api.md b/docs/src/api.md
index fb66c49a..20c35bd2 100644
--- a/docs/src/api.md
+++ b/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
diff --git a/docs/src/users_guide/QuantumObject/QuantumObject.md b/docs/src/users_guide/QuantumObject/QuantumObject.md
index 2f53e90c..61e9569f 100644
--- a/docs/src/users_guide/QuantumObject/QuantumObject.md
+++ b/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)
diff --git a/src/QuantumToolbox.jl b/src/QuantumToolbox.jl
index 45b8cbd0..d6735644 100644
--- a/src/QuantumToolbox.jl
+++ b/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")
diff --git a/src/linear_maps.jl b/src/linear_maps.jl
new file mode 100644
index 00000000..8afd88e7
--- /dev/null
+++ b/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]
diff --git a/src/metrics.jl b/src/metrics.jl
index 62f7af6d..d257d24b 100644
--- a/src/metrics.jl
+++ b/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
 
 """
diff --git a/src/qobj/arithmetic_and_attributes.jl b/src/qobj/arithmetic_and_attributes.jl
index 6dc2f240..4ba7d4bc 100644
--- a/src/qobj/arithmetic_and_attributes.jl
+++ b/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],
     )
diff --git a/src/qobj/eigsolve.jl b/src/qobj/eigsolve.jl
index 03e8f852..dee88be2 100644
--- a/src/qobj/eigsolve.jl
+++ b/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)
diff --git a/src/qobj/functions.jl b/src/qobj/functions.jl
index 6dde41d6..667f7930 100644
--- a/src/qobj/functions.jl
+++ b/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"""
diff --git a/src/qobj/operators.jl b/src/qobj/operators.jl
index ac04102b..cf934fe5 100644
--- a/src/qobj/operators.jl
+++ b/src/qobj/operators.jl
@@ -13,7 +13,7 @@ export tunneling
 export qft
 
 @doc raw"""
-    rand_unitary(dimensions, distribution=:haar)
+    rand_unitary(dimensions, distribution=Val(:haar))
 
 Returns a random unitary [`QuantumObject`](@ref).
 
@@ -27,10 +27,14 @@ The `distribution` specifies which of the method used to obtain the unitary matr
 
 # References
 1. [F. Mezzadri, How to generate random matrices from the classical compact groups, arXiv:math-ph/0609050 (2007)](https://arxiv.org/abs/math-ph/0609050)
+
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `rand_unitary(dimensions, Val(distribution))` instead of `rand_unitary(dimensions, distribution)`. Also, put `dimensions` as `Tuple` or `SVector`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
-rand_unitary(dimensions::Int, distribution::Symbol = :haar) = rand_unitary(SVector(dimensions), Val(distribution))
-rand_unitary(dimensions::Union{AbstractVector{Int},Tuple}, distribution::Symbol = :haar) =
-    rand_unitary(dimensions, Val(distribution))
+rand_unitary(dimensions::Int, distribution::Union{Symbol,Val} = Val(:haar)) =
+    rand_unitary(SVector(dimensions), makeVal(distribution))
+rand_unitary(dimensions::Union{AbstractVector{Int},Tuple}, distribution::Union{Symbol,Val} = Val(:haar)) =
+    rand_unitary(dimensions, makeVal(distribution))
 function rand_unitary(dimensions::Union{AbstractVector{Int},Tuple}, ::Val{:haar})
     N = prod(dimensions)
 
@@ -224,7 +228,7 @@ function phase(N::Int, ϕ0::Real = 0)
 end
 
 @doc raw"""
-    jmat(j::Real, which::Symbol)
+    jmat(j::Real, which::Union{Symbol,Val})
 
 Generate higher-order Spin-`j` operators, where `j` is the spin quantum number and can be a non-negative integer or half-integer
 
@@ -250,7 +254,18 @@ Quantum Object:   type=Operator   dims=[2]   size=(2, 2)   ishermitian=false
 2×2 SparseMatrixCSC{ComplexF64, Int64} with 1 stored entry:
      ⋅          ⋅    
  1.0+0.0im      ⋅
+
+julia> jmat(1.5, Val(:z))
+Quantum Object:   type=Operator   dims=[4]   size=(4, 4)   ishermitian=true
+4×4 SparseMatrixCSC{ComplexF64, Int64} with 4 stored entries:
+ 1.5+0.0im      ⋅           ⋅           ⋅    
+     ⋅      0.5+0.0im       ⋅           ⋅    
+     ⋅          ⋅      -0.5+0.0im       ⋅    
+     ⋅          ⋅           ⋅      -1.5+0.0im
 ```
+
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `jmat(j, Val(which))` instead of `jmat(j, which)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 jmat(j::Real, which::Symbol) = jmat(j, Val(which))
 jmat(j::Real) = (jmat(j, Val(:x)), jmat(j, Val(:y)), jmat(j, Val(:z)))
@@ -427,8 +442,8 @@ d_j = \sigma_z^{\otimes j} \otimes \sigma_{-} \otimes I^{\otimes N-j-1}
 
 Note that the site index `j` should satisfy: `0 ≤ j ≤ N - 1`.
 
-> [!IMPORTANT]
-> If you want to keep type stability, it is recommended to use `fdestroy(Val(N), j)` instead of `fdestroy(N, j)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) for more details.
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `fdestroy(Val(N), j)` instead of `fdestroy(N, j)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 fdestroy(N::Union{Int,Val}, j::Int) = _Jordan_Wigner(N, j, sigmam())
 
@@ -444,8 +459,8 @@ d_j^\dagger = \sigma_z^{\otimes j} \otimes \sigma_{+} \otimes I^{\otimes N-j-1}
 
 Note that the site index `j` should satisfy: `0 ≤ j ≤ N - 1`.
 
-> [!IMPORTANT]
-> If you want to keep type stability, it is recommended to use `fcreate(Val(N), j)` instead of `fcreate(N, j)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) for more details.
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `fcreate(Val(N), j)` instead of `fcreate(N, j)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 fcreate(N::Union{Int,Val}, j::Int) = _Jordan_Wigner(N, j, sigmap())
 
@@ -470,7 +485,7 @@ end
 
 Generates the projection operator ``\hat{O} = \dyad{i}{j}`` with Hilbert space dimension `N`.
 """
-projection(N::Int, i::Int, j::Int) = QuantumObject(sparse([i + 1], [j + 1], [1.0 + 0.0im], N, N))
+projection(N::Int, i::Int, j::Int) = QuantumObject(sparse([i + 1], [j + 1], [1.0 + 0.0im], N, N), type = Operator)
 
 @doc raw"""
     tunneling(N::Int, m::Int=1; sparse::Union{Bool,Val{<:Bool}}=Val(false))
@@ -485,8 +500,8 @@ where ``N`` is the number of basis states in the Hilbert space, and ``m`` is the
 
 If `sparse=true`, the operator is returned as a sparse matrix, otherwise a dense matrix is returned.
 
-> [!IMPORTANT]
-> If you want to keep type stability, it is recommended to use `tunneling(N, m, Val(sparse))` instead of `tunneling(N, m, sparse)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) for more details.
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `tunneling(N, m, Val(sparse))` instead of `tunneling(N, m, sparse)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 function tunneling(N::Int, m::Int = 1; sparse::Union{Bool,Val} = Val(false))
     (m < 1) && throw(ArgumentError("The number of excitations (m) cannot be less than 1"))
@@ -522,6 +537,9 @@ The `dimensions` can be either the following types:
 ```
 
 where ``\omega = \exp(\frac{2 \pi i}{N})``.
+
+!!! warning "Beware of type-stability!"
+    It is highly recommended to use `qft(dimensions)` with `dimensions` as `Tuple` or `SVector` to keep type stability. See the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 qft(dimensions::Int) = QuantumObject(_qft_op(dimensions), Operator, dimensions)
 qft(dimensions::Union{AbstractVector{T},Tuple}) where {T} =
diff --git a/src/qobj/states.jl b/src/qobj/states.jl
index d1b8a0e9..04a3cdcc 100644
--- a/src/qobj/states.jl
+++ b/src/qobj/states.jl
@@ -15,41 +15,45 @@ Returns a zero [`Ket`](@ref) vector with given argument `dimensions`.
 The `dimensions` can be either the following types:
 - `dimensions::Int`: Number of basis states in the Hilbert space.
 - `dimensions::Union{AbstractVector{Int}, Tuple}`: list of dimensions representing the each number of basis in the subsystems.
+
+!!! warning "Beware of type-stability!"
+    It is highly recommended to use `zero_ket(dimensions)` with `dimensions` as `Tuple` or `SVector` to keep type stability. See the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 zero_ket(dimensions::Int) = QuantumObject(zeros(ComplexF64, dimensions), Ket, dimensions)
 zero_ket(dimensions::Union{AbstractVector{Int},Tuple}) =
     QuantumObject(zeros(ComplexF64, prod(dimensions)), Ket, dimensions)
 
 @doc raw"""
-    fock(N::Int, pos::Int=0; dims::Union{Int,AbstractVector{Int},Tuple}=N, sparse::Union{Bool,Val}=Val(false))
+    fock(N::Int, j::Int=0; dims::Union{Int,AbstractVector{Int},Tuple}=N, sparse::Union{Bool,Val}=Val(false))
 
 Generates a fock state ``\ket{\psi}`` of dimension `N`. 
 
 It is also possible to specify the list of dimensions `dims` if different subsystems are present.
+
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `fock(N, j, dims=dims, sparse=Val(sparse))` instead of `fock(N, j, dims=dims, sparse=sparse)`. Consider also to use `dims` as a `Tuple` or `SVector` instead of `Vector`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
-function fock(
-    N::Int,
-    pos::Int = 0;
-    dims::Union{Int,AbstractVector{Int},Tuple} = N,
-    sparse::Union{Bool,Val} = Val(false),
-)
+function fock(N::Int, j::Int = 0; dims::Union{Int,AbstractVector{Int},Tuple} = N, sparse::Union{Bool,Val} = Val(false))
     if getVal(makeVal(sparse))
-        array = sparsevec([pos + 1], [1.0 + 0im], N)
+        array = sparsevec([j + 1], [1.0 + 0im], N)
     else
         array = zeros(ComplexF64, N)
-        array[pos+1] = 1
+        array[j+1] = 1
     end
     return QuantumObject(array; type = Ket, dims = dims)
 end
 
 @doc raw"""
-    basis(N::Int, pos::Int = 0; dims::Union{Int,AbstractVector{Int},Tuple}=N)
+    basis(N::Int, j::Int = 0; dims::Union{Int,AbstractVector{Int},Tuple}=N)
 
 Generates a fock state like [`fock`](@ref).
 
 It is also possible to specify the list of dimensions `dims` if different subsystems are present.
+
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `basis(N, j, dims=dims)` with `dims` as a `Tuple` or `SVector` instead of `Vector`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
-basis(N::Int, pos::Int = 0; dims::Union{Int,AbstractVector{Int},Tuple} = N) = fock(N, pos, dims = dims)
+basis(N::Int, j::Int = 0; dims::Union{Int,AbstractVector{Int},Tuple} = N) = fock(N, j, dims = dims)
 
 @doc raw"""
     coherent(N::Int, α::Number)
@@ -68,6 +72,9 @@ Generate a random normalized [`Ket`](@ref) vector with given argument `dimension
 The `dimensions` can be either the following types:
 - `dimensions::Int`: Number of basis states in the Hilbert space.
 - `dimensions::Union{AbstractVector{Int},Tuple}`: list of dimensions representing the each number of basis in the subsystems.
+
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `rand_ket(dimensions)` with `dimensions` as `Tuple` or `SVector` to keep type stability. See the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 rand_ket(dimensions::Int) = rand_ket(SVector(dimensions))
 function rand_ket(dimensions::Union{AbstractVector{Int},Tuple})
@@ -77,19 +84,22 @@ function rand_ket(dimensions::Union{AbstractVector{Int},Tuple})
 end
 
 @doc raw"""
-    fock_dm(N::Int, pos::Int=0; dims::Union{Int,AbstractVector{Int},Tuple}=N, sparse::Union{Bool,Val}=Val(false))
+    fock_dm(N::Int, j::Int=0; dims::Union{Int,AbstractVector{Int},Tuple}=N, sparse::Union{Bool,Val}=Val(false))
 
 Density matrix representation of a Fock state.
 
 Constructed via outer product of [`fock`](@ref).
+
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `fock_dm(N, j, dims=dims, sparse=Val(sparse))` instead of `fock_dm(N, j, dims=dims, sparse=sparse)`. Consider also to use `dims` as a `Tuple` or `SVector` instead of `Vector`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 function fock_dm(
     N::Int,
-    pos::Int = 0;
+    j::Int = 0;
     dims::Union{Int,AbstractVector{Int},Tuple} = N,
     sparse::Union{Bool,Val} = Val(false),
 )
-    ψ = fock(N, pos; dims = dims, sparse = sparse)
+    ψ = fock(N, j; dims = dims, sparse = sparse)
     return ket2dm(ψ)
 end
 
@@ -113,8 +123,8 @@ Density matrix for a thermal state (generating thermal state probabilities) with
 - `n::Real`: Expectation value for number of particles in the thermal state.
 - `sparse::Union{Bool,Val}`: If `true`, return a sparse matrix representation.
 
-> [!IMPORTANT]
-> If you want to keep type stability, it is recommended to use `thermal_dm(N, n, Val(sparse))` instead of `thermal_dm(N, n, sparse)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) for more details.
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `thermal_dm(N, n, sparse=Val(sparse))` instead of `thermal_dm(N, n, sparse=sparse)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 function thermal_dm(N::Int, n::Real; sparse::Union{Bool,Val} = Val(false))
     β = log(1.0 / n + 1.0)
@@ -135,6 +145,9 @@ Returns the maximally mixed density matrix with given argument `dimensions`.
 The `dimensions` can be either the following types:
 - `dimensions::Int`: Number of basis states in the Hilbert space.
 - `dimensions::Union{AbstractVector{Int},Tuple}`: list of dimensions representing the each number of basis in the subsystems.
+
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `maximally_mixed_dm(dimensions)` with `dimensions` as `Tuple` or `SVector` to keep type stability. See the [related Section](@ref doc:Type-Stability) about type stability for more details.
 """
 maximally_mixed_dm(dimensions::Int) = QuantumObject(I(dimensions) / complex(dimensions), Operator, SVector(dimensions))
 function maximally_mixed_dm(dimensions::Union{AbstractVector{Int},Tuple})
@@ -153,6 +166,9 @@ The `dimensions` can be either the following types:
 
 The default keyword argument `rank = prod(dimensions)` (full rank).
 
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `rand_dm(dimensions; rank=rank)` with `dimensions` as `Tuple` or `SVector` instead of `Vector`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details.
+
 # References
 - [J. Ginibre, Statistical ensembles of complex, quaternion, and real matrices, Journal of Mathematical Physics 6.3 (1965): 440-449](https://doi.org/10.1063/1.1704292)
 - [K. Życzkowski, et al., Generating random density matrices, Journal of Mathematical Physics 52, 062201 (2011)](http://dx.doi.org/10.1063/1.3595693)
@@ -222,7 +238,7 @@ function spin_coherent(j::Real, θ::Real, ϕ::Real)
 end
 
 @doc raw"""
-    bell_state(x::Int, z::Int)
+    bell_state(x::Union{Int}, z::Union{Int})
 
 Return the [Bell state](https://en.wikipedia.org/wiki/Bell_state) depending on the arguments `(x, z)`:
 - `(0, 0)`: ``| \Phi^+ \rangle = ( |00\rangle + |11\rangle ) / \sqrt{2}``
@@ -242,7 +258,18 @@ Quantum Object:   type=Ket   dims=[2, 2]   size=(4,)
                 0.0 + 0.0im
                 0.0 + 0.0im
  0.7071067811865475 + 0.0im
+
+julia> bell_state(Val(1), Val(0))
+Quantum Object:   type=Ket   dims=[2, 2]   size=(4,)
+4-element Vector{ComplexF64}:
+                0.0 + 0.0im
+ 0.7071067811865475 + 0.0im
+ 0.7071067811865475 + 0.0im
+                0.0 + 0.0im
 ```
+
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `bell_state(Val(x), Val(z))` instead of `bell_state(x, z)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) for more details.
 """
 bell_state(x::Int, z::Int) = bell_state(Val(x), Val(z))
 bell_state(::Val{0}, ::Val{0}) = QuantumObject(ComplexF64[1, 0, 0, 1] / sqrt(2), Ket, (2, 2))
@@ -284,8 +311,8 @@ Returns the `n`-qubit [W-state](https://en.wikipedia.org/wiki/W_state):
 \frac{1}{\sqrt{n}} \left( |100...0\rangle + |010...0\rangle + \cdots + |00...01\rangle \right)
 ```
 
-> [!IMPORTANT]
-> If you want to keep type stability, it is recommended to use `w_state(Val(n))` instead of `w_state(n)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) for more details.
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `w_state(Val(n))` instead of `w_state(n)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) for more details.
 """
 function w_state(::Val{n}) where {n}
     nzind = 2 .^ (0:(n-1)) .+ 1
@@ -305,8 +332,8 @@ Returns the generalized `n`-qudit [Greenberger–Horne–Zeilinger (GHZ) state](
 
 Here, `d` specifies the dimension of each qudit. Default to `d=2` (qubit).
 
-> [!IMPORTANT]
-> If you want to keep type stability, it is recommended to use `ghz_state(Val(n); kwargs...)` instead of `ghz_state(n; kwargs...)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) for more details.
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `ghz_state(Val(n))` instead of `ghz_state(n)`. See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) for more details.
 """
 function ghz_state(::Val{n}; d::Int = 2) where {n}
     nzind = collect((0:(d-1)) .* Int((d^n - 1) / (d - 1)) .+ 1)
diff --git a/src/steadystate.jl b/src/steadystate.jl
index eafcf924..7cdd0f14 100644
--- a/src/steadystate.jl
+++ b/src/steadystate.jl
@@ -31,10 +31,14 @@ end
 
 struct SteadyStateDirectSolver <: SteadyStateSolver end
 struct SteadyStateEigenSolver <: SteadyStateSolver end
-Base.@kwdef struct SteadyStateLinearSolver{MT<:Union{SciMLLinearSolveAlgorithm,Nothing}} <: SteadyStateSolver
-    alg::MT = nothing
-    Pl::Union{Function,Nothing} = nothing
-    Pr::Union{Function,Nothing} = nothing
+Base.@kwdef struct SteadyStateLinearSolver{
+    MT<:Union{SciMLLinearSolveAlgorithm,Nothing},
+    PlT<:Union{Function,Nothing},
+    PrT<:Union{Function,Nothing},
+} <: SteadyStateSolver
+    alg::MT = KrylovJL_GMRES()
+    Pl::PlT = nothing
+    Pr::PrT = nothing
 end
 
 @doc raw"""
@@ -363,13 +367,13 @@ function _steadystate_floquet(
     ρ0 = reshape(ρtot[offset1+1:offset2], Ns, Ns)
     ρ0_tr = tr(ρ0)
     ρ0 = ρ0 / ρ0_tr
-    ρ0 = QuantumObject((ρ0 + ρ0') / 2, dims = L_0.dims)
+    ρ0 = QuantumObject((ρ0 + ρ0') / 2, type = Operator, dims = L_0.dims)
     ρtot = ρtot / ρ0_tr
 
     ρ_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, dims = L_0.dims)
+        ρi_m = QuantumObject(ρi_m, type = Operator, dims = L_0.dims)
         push!(ρ_list, ρi_m)
     end
 
diff --git a/src/time_evolution/time_evolution.jl b/src/time_evolution/time_evolution.jl
index 8417ee7a..6751a3da 100644
--- a/src/time_evolution/time_evolution.jl
+++ b/src/time_evolution/time_evolution.jl
@@ -197,10 +197,10 @@ end
 
 function liouvillian_floquet(
     H::QuantumObject{<:AbstractArray{T1},OpType1},
-    c_ops::AbstractVector,
     Hₚ::QuantumObject{<:AbstractArray{T2},OpType2},
     Hₘ::QuantumObject{<:AbstractArray{T3},OpType3},
-    ω::Real;
+    ω::Real,
+    c_ops::Union{AbstractVector,Nothing} = nothing;
     n_max::Int = 3,
     tol::Real = 1e-15,
 ) where {
@@ -226,30 +226,31 @@ function liouvillian_generalized(
     H::QuantumObject{MT,OperatorQuantumObject},
     fields::Vector,
     T_list::Vector{<:Real};
-    N_trunc::Int = size(H, 1),
+    N_trunc::Union{Int,Nothing} = nothing,
     tol::Real = 1e-12,
     σ_filter::Union{Nothing,Real} = nothing,
 ) where {MT<:AbstractMatrix}
     (length(fields) == length(T_list)) || throw(DimensionMismatch("The number of fields, ωs and Ts must be the same."))
 
-    dims = N_trunc == size(H, 1) ? H.dims : SVector(N_trunc)
+    dims = (N_trunc isa Nothing) ? H.dims : SVector(N_trunc)
+    final_size = prod(dims)
     result = eigen(H)
-    E = real.(result.values[1:N_trunc])
+    E = real.(result.values[1:final_size])
     U = QuantumObject(result.vectors, result.type, result.dims)
 
-    H_d = QuantumObject(Diagonal(complex(E)), dims = dims)
+    H_d = QuantumObject(Diagonal(complex(E)), type = Operator, dims = dims)
 
     Ω = E' .- E
     Ωp = triu(dense_to_sparse(Ω, tol), 1)
 
     # Filter in the Hilbert space
     σ = isnothing(σ_filter) ? 500 * maximum([norm(field) / length(field) for field in fields]) : σ_filter
-    F1 = QuantumObject(gaussian.(Ω, 0, σ), dims = dims)
+    F1 = QuantumObject(gaussian.(Ω, 0, σ), type = Operator, dims = dims)
     F1 = dense_to_sparse(F1, tol)
 
     # Filter in the Liouville space
-    # M1 = ones(N_trunc, N_trunc)
-    M1 = similar(E, N_trunc, N_trunc)
+    # M1 = ones(final_size, final_size)
+    M1 = similar(E, final_size, final_size)
     M1 .= 1
     Ω1 = kron(Ω, M1)
     Ω2 = kron(M1, Ω)
@@ -261,16 +262,16 @@ function liouvillian_generalized(
 
     for i in eachindex(fields)
         # The operator that couples the system to the bath in the eigenbasis
-        X_op = dense_to_sparse((U'*fields[i]*U).data[1:N_trunc, 1:N_trunc], tol)
+        X_op = dense_to_sparse((U'*fields[i]*U).data[1:final_size, 1:final_size], tol)
         if ishermitian(fields[i])
             X_op = (X_op + X_op') / 2 # Make sure it's hermitian
         end
 
         # Ohmic reservoir
         N_th = n_th.(Ωp, T_list[i])
-        Sp₀ = QuantumObject(triu(X_op, 1), dims = dims)
-        Sp₁ = QuantumObject(droptol!((@. Ωp * N_th * Sp₀.data), tol), dims = dims)
-        Sp₂ = QuantumObject(droptol!((@. Ωp * (1 + N_th) * Sp₀.data), tol), dims = dims)
+        Sp₀ = QuantumObject(triu(X_op, 1), type = Operator, dims = dims)
+        Sp₁ = QuantumObject(droptol!((@. Ωp * N_th * Sp₀.data), tol), type = Operator, dims = dims)
+        Sp₂ = QuantumObject(droptol!((@. Ωp * (1 + N_th) * Sp₀.data), tol), type = Operator, dims = dims)
         # S0 = QuantumObject( spdiagm(diag(X_op)), dims=dims )
 
         L +=
diff --git a/src/time_evolution/time_evolution_dynamical.jl b/src/time_evolution/time_evolution_dynamical.jl
index 643e6847..9ee68fac 100644
--- a/src/time_evolution/time_evolution_dynamical.jl
+++ b/src/time_evolution/time_evolution_dynamical.jl
@@ -8,20 +8,19 @@ function _reduce_dims(
     sel,
     reduce,
 ) where {T,N,DT<:Integer}
-    rd = dims
-    nd = length(rd)
-    rd_new = zero(rd)
-    rd_new[sel] .= reduce
-    @. rd_new = rd - rd_new
+    nd = length(dims)
+    dims_new = zero(dims)
+    dims_new[sel] .= reduce
+    @. dims_new = dims - dims_new
 
     if nd == 1
-        ρmat = similar(QO, rd_new[1], rd_new[1])
-        copyto!(ρmat, view(QO, 1:rd_new[1], 1:rd_new[1]))
+        ρmat = similar(QO, dims_new[1], dims_new[1])
+        copyto!(ρmat, view(QO, 1:dims_new[1], 1:dims_new[1]))
     else
-        ρmat = reshape(QO, reverse!(repeat(rd, 2))...)
-        ρmat2 = similar(QO, reverse!(repeat(rd_new, 2))...)
-        copyto!(ρmat2, view(ρmat, reverse!(repeat([1:n for n in rd_new], 2))...))
-        ρmat = reshape(ρmat2, prod(rd_new), prod(rd_new))
+        ρmat = reshape(QO, reverse(vcat(dims, dims))...)
+        ρmat2 = similar(QO, reverse(vcat(dims_new, dims_new))...)
+        copyto!(ρmat2, view(ρmat, reverse!(repeat([1:n for n in dims_new], 2))...))
+        ρmat = reshape(ρmat2, prod(dims_new), prod(dims_new))
     end
 
     return ρmat
@@ -33,26 +32,25 @@ function _increase_dims(
     sel,
     increase,
 ) where {T,N,DT<:Integer}
-    rd = dims
-    nd = length(rd)
-    rd_new = MVector(zero(rd)) # Mutable SVector
-    rd_new[sel] .= increase
-    @. rd_new = rd + rd_new
+    nd = length(dims)
+    dims_new = MVector(zero(dims)) # Mutable SVector
+    dims_new[sel] .= increase
+    @. dims_new = dims + dims_new
 
     if nd == 1
-        ρmat = similar(QO, rd_new[1], rd_new[1])
-        fill!(selectdim(ρmat, 1, rd[1]+1:rd_new[1]), 0)
-        fill!(selectdim(ρmat, 2, rd[1]+1:rd_new[1]), 0)
-        copyto!(view(ρmat, 1:rd[1], 1:rd[1]), QO)
+        ρmat = similar(QO, dims_new[1], dims_new[1])
+        fill!(selectdim(ρmat, 1, dims[1]+1:dims_new[1]), 0)
+        fill!(selectdim(ρmat, 2, dims[1]+1:dims_new[1]), 0)
+        copyto!(view(ρmat, 1:dims[1], 1:dims[1]), QO)
     else
-        ρmat2 = similar(QO, reverse!(repeat(rd_new, 2))...)
-        ρmat = reshape(QO, reverse!(repeat(rd, 2))...)
+        ρmat2 = similar(QO, reverse(vcat(dims_new, dims_new))...)
+        ρmat = reshape(QO, reverse(vcat(dims, dims))...)
         for i in eachindex(sel)
-            fill!(selectdim(ρmat2, nd - sel[i] + 1, rd[sel[i]]+1:rd_new[sel[i]]), 0)
-            fill!(selectdim(ρmat2, 2 * nd - sel[i] + 1, rd[sel[i]]+1:rd_new[sel[i]]), 0)
+            fill!(selectdim(ρmat2, nd - sel[i] + 1, dims[sel[i]]+1:dims_new[sel[i]]), 0)
+            fill!(selectdim(ρmat2, 2 * nd - sel[i] + 1, dims[sel[i]]+1:dims_new[sel[i]]), 0)
         end
-        copyto!(view(ρmat2, reverse!(repeat([1:n for n in rd], 2))...), ρmat)
-        ρmat = reshape(ρmat2, prod(rd_new), prod(rd_new))
+        copyto!(view(ρmat2, reverse!(repeat([1:n for n in dims], 2))...), ρmat)
+        ρmat = reshape(ρmat2, prod(dims_new), prod(dims_new))
     end
 
     return ρmat
@@ -131,7 +129,9 @@ function _DFDIncreaseReduceAffect!(integrator)
 
     resize!(integrator, size(L, 1))
     copyto!(integrator.u, mat2vec(ρt))
-    return integrator.p = merge(internal_params, (L = L, e_ops = e_ops2, dfd_ρt_cache = similar(integrator.u)))
+    integrator.p = merge(internal_params, (L = L, e_ops = e_ops2, dfd_ρt_cache = similar(integrator.u)))
+
+    return nothing
 end
 
 function dfd_mesolveProblem(
@@ -241,6 +241,7 @@ function dfd_mesolve(
     ρt = map(
         i -> QuantumObject(
             vec2mat(sol.u[i]),
+            type = Operator,
             dims = sol.prob.p.dim_list_evo[searchsortedlast(sol.prob.p.dim_list_evo_times, sol.t[i])],
         ),
         eachindex(sol.t),
diff --git a/test/eigenvalues_and_operators.jl b/test/eigenvalues_and_operators.jl
index 6f281208..9e8ad76a 100644
--- a/test/eigenvalues_and_operators.jl
+++ b/test/eigenvalues_and_operators.jl
@@ -59,7 +59,7 @@
     # eigen solve for general matrices
     vals, _, vecs = eigsolve(L.data, sigma = 0.01, k = 10, krylovdim = 50)
     vals2, vecs2 = eigen(sparse_to_dense(L.data))
-    vals3, state3, vecs3 = eigsolve_al(liouvillian(H, c_ops), 1 \ (40 * κ), k = 10, krylovdim = 50)
+    vals3, state3, vecs3 = eigsolve_al(L, 1 \ (40 * κ), k = 10, krylovdim = 50)
     idxs = sortperm(vals2, by = abs)
     vals2 = vals2[idxs][1:10]
     vecs2 = vecs2[:, idxs][:, 1:10]
@@ -91,4 +91,27 @@
     @test isapprox(abs2(vals2[1]), abs2(vals3[1]), atol = 1e-7)
     @test isapprox(vec2mat(vecs[1]).data * exp(-1im * angle(vecs[1][1])), vec2mat(vecs2[1]).data, atol = 1e-7)
     @test isapprox(vec2mat(vecs[1]).data * exp(-1im * angle(vecs[1][1])), vec2mat(state3[1]).data, atol = 1e-5)
+
+    @testset "Type Inference (eigen)" begin
+        N = 5
+        a = kron(destroy(N), qeye(N))
+        a_d = a'
+        b = kron(qeye(N), destroy(N))
+        b_d = b'
+
+        ωc = 1
+        ωb = 1
+        g = 0.01
+        κ = 0.1
+        n_thermal = 0.01
+
+        H = ωc * a_d * a + ωb * b_d * b + g * (a + a_d) * (b + b_d)
+        c_ops = [√((1 + n_thermal) * κ) * a, √κ * b, √(n_thermal * κ) * a_d]
+        L = liouvillian(H, c_ops)
+
+        @inferred eigenstates(H, sparse = false)
+        @inferred eigenstates(H, sparse = true)
+        @inferred eigenstates(L, sparse = true)
+        @inferred eigsolve_al(L, 1 \ (40 * κ), k = 10)
+    end
 end
diff --git a/test/entanglement.jl b/test/entanglement.jl
index 860e168c..c9df3e20 100644
--- a/test/entanglement.jl
+++ b/test/entanglement.jl
@@ -5,4 +5,9 @@
     rho = state * state'
     @test entanglement(state, 1) / log(2) ≈ 1
     @test entanglement(rho, 1) / log(2) ≈ 1
+
+    @testset "Type Stability (entanglement)" begin
+        @inferred entanglement(state, 1)
+        @inferred entanglement(rho, 1)
+    end
 end
diff --git a/test/generalized_master_equation.jl b/test/generalized_master_equation.jl
index 965206f9..080ab17e 100644
--- a/test/generalized_master_equation.jl
+++ b/test/generalized_master_equation.jl
@@ -26,7 +26,7 @@
     c_ops = [sqrt(0.01) * a2, sqrt(0.01) * sm2]
     L2 = liouvillian(H_d, c_ops)
 
-    @test (expect(Xp' * Xp, steadystate(L1)) < 1e-10 && expect(Xp' * Xp, steadystate(L2)) > 1e-3)
+    @test (expect(Hermitian(Xp' * Xp), steadystate(L1)) < 1e-10 && expect(Hermitian(Xp' * Xp), steadystate(L2)) > 1e-3)
 
     H = 1 * a' * a + 1 * sz / 2 + 1e-5 * (a * sp + a' * sm)
 
@@ -41,4 +41,23 @@
     sm2 = Qobj(dense_to_sparse((U'*sm*U).data[1:N_trunc, 1:N_trunc], tol))
 
     @test abs(expect(Xp' * Xp, steadystate(L1)) - n_th(1, Tlist[1])) / n_th(1, Tlist[1]) < 1e-4
+
+    @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
+
+        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]
+
+        @inferred liouvillian_generalized(H, fields, Tlist, N_trunc = N_trunc, tol = tol)
+    end
 end
diff --git a/test/negativity_and_partial_transpose.jl b/test/negativity_and_partial_transpose.jl
index d729d0f4..ba10db89 100644
--- a/test/negativity_and_partial_transpose.jl
+++ b/test/negativity_and_partial_transpose.jl
@@ -14,6 +14,11 @@
         @test negativity(rho, 2) ≈ Neg
         @test negativity(rho, 1; logarithmic = true) ≈ log2(2 * Neg + 1)
         @test_throws ArgumentError negativity(rho, 3)
+
+        @testset "Type Inference (negativity)" begin
+            @inferred negativity(rho, 1)
+            @inferred negativity(rho, 1; logarithmic = true)
+        end
     end
 
     @testset "partial_transpose" begin
@@ -30,5 +35,10 @@
             end
         end
         @test_throws ArgumentError partial_transpose(A_dense, [true])
+
+        @testset "Type Inference (partial_transpose)" begin
+            @inferred partial_transpose(A_dense, [true, false, true])
+            @inferred partial_transpose(A_sparse, [true, false, true])
+        end
     end
 end
diff --git a/test/permutation.jl b/test/permutation.jl
index 3b28d9bf..6f089254 100644
--- a/test/permutation.jl
+++ b/test/permutation.jl
@@ -7,14 +7,13 @@
     θ = atan(tg)
     U = sin(θ)
     κ2 = cos(θ)
-    κ1 = 0.0
     κϕ = 1e-3
     nth = 0.0
 
     a = destroy(N)
     ad = create(N)
     H = -Δ * ad * a + G / 2 * (ad^2 + a^2) + U / 2 * (ad^2 * a^2)
-    c_ops = [√(κ2) * a^2, √(κ1 * (nth + 1)) * a, √(κ1 * nth) * ad, √(κϕ) * ad * a]
+    c_ops = [√(κ2) * a^2, √(κϕ) * ad * a]
     L = liouvillian(H, c_ops)
 
     P, L_bd, block_sizes = bdf(L)
@@ -24,4 +23,8 @@
     @test length(blocks_list) == 4
     @test length(block_indices) == 4
     @test sum(block_sizes .== 100) == 4
+
+    @testset "Type Inference (bdf)" begin
+        @inferred bdf(L)
+    end
 end
diff --git a/test/quantum_objects.jl b/test/quantum_objects.jl
index 1ef56ae4..e61ff996 100644
--- a/test/quantum_objects.jl
+++ b/test/quantum_objects.jl
@@ -1,28 +1,4 @@
 @testset "Quantum Objects" verbose = true begin
-    @testset "Type Inference" begin
-        for T in [ComplexF32, ComplexF64]
-            N = 4
-            a = rand(T, N)
-            @inferred QuantumObject{typeof(a),KetQuantumObject} Qobj(a)
-            for type in [Ket, OperatorKet]
-                @inferred Qobj(a, type = type)
-            end
-
-            UnionType = Union{QuantumObject{Matrix{T},BraQuantumObject},QuantumObject{Matrix{T},OperatorQuantumObject}}
-            a = rand(T, 1, N)
-            @inferred UnionType Qobj(a)
-            for type in [Bra, OperatorBra]
-                @inferred Qobj(a, type = type)
-            end
-
-            a = rand(T, N, N)
-            @inferred UnionType Qobj(a)
-            for type in [Operator, SuperOperator]
-                @inferred Qobj(a, type = type)
-            end
-        end
-    end
-
     # ArgumentError: type is incompatible with vector or matrix
     @testset "ArgumentError" begin
         a = rand(ComplexF64, 2)
@@ -294,40 +270,102 @@
         @test typeof(SparseMatrixCSC{ComplexF64}(Ms).data) == SparseMatrixCSC{ComplexF64,Int64}
     end
 
+    @testset "Type Inference (QuantumObject)" begin
+        for T in [ComplexF32, ComplexF64]
+            N = 4
+            a = rand(T, N)
+            @inferred QuantumObject{typeof(a),KetQuantumObject} Qobj(a)
+            for type in [Ket, OperatorKet]
+                @inferred Qobj(a, type = type)
+            end
+
+            UnionType =
+                Union{QuantumObject{Matrix{T},BraQuantumObject,1},QuantumObject{Matrix{T},OperatorQuantumObject,1}}
+            a = rand(T, 1, N)
+            @inferred UnionType Qobj(a)
+            for type in [Bra, OperatorBra]
+                @inferred Qobj(a, type = type)
+            end
+
+            a = rand(T, N, N)
+            @inferred UnionType Qobj(a)
+            for type in [Operator, SuperOperator]
+                @inferred Qobj(a, type = type)
+            end
+        end
+
+        @testset "Math Operation" begin
+            a = destroy(20)
+            σx = sigmax()
+            @inferred a + a
+            @inferred a + a'
+            @inferred a + 2
+            @inferred 2 * a
+            @inferred a / 2
+            @inferred a^2
+            @inferred a .+ 2
+            @inferred a .* 2
+            @inferred a ./ 2
+            @inferred a .^ 2
+            @inferred a * a
+            @inferred a * a'
+            @inferred kron(a, σx)
+            @inferred kron(a, eye(2))
+        end
+    end
+
     @testset "projection" begin
         N = 10
-        a = fock(N, 3)
-        @test proj(a) ≈ proj(a') ≈ sparse(ket2dm(a)) ≈ projection(N, 3, 3)
-        @test isket(a') == false
-        @test isbra(a') == true
-        @test shape(a) == (N,)
-        @test shape(a') == (1, N)
-        @test norm(a) ≈ 1
-        @test norm(a') ≈ 1
+        ψ = fock(N, 3)
+        @test proj(ψ) ≈ proj(ψ') ≈ sparse(ket2dm(ψ)) ≈ projection(N, 3, 3)
+        @test isket(ψ') == false
+        @test isbra(ψ') == true
+        @test shape(ψ) == (N,)
+        @test shape(ψ') == (1, N)
+        @test norm(ψ) ≈ 1
+        @test norm(ψ') ≈ 1
+
+        @testset "Type Inference (proj)" begin
+            @inferred proj(ψ)
+            @inferred proj(ψ')
+        end
     end
 
     @testset "dot product" begin
         ψ = rand_ket(10)
         @test dot(ψ, ψ) ≈ ψ' * ψ ≈ norm(ψ) ≈ 1.0
+
+        @testset "Type Inference (dot)" begin
+            @inferred dot(ψ, ψ)
+            @inferred ψ' * ψ
+            @inferred norm(ψ)
+        end
     end
 
     @testset "normalization" begin
         # normalize, normalize!, unit
         N = 10
-        a = Qobj(rand(ComplexF64, N))
-        M = a * a'
-        @test (norm(a) ≈ 1) == false
+        ψ = Qobj(rand(ComplexF64, N))
+        M = ψ * ψ'
+        @test (norm(ψ) ≈ 1) == false
         @test (norm(M) ≈ 1) == false
-        @test (norm(unit(a)) ≈ 1) == true
+        @test (norm(unit(ψ)) ≈ 1) == true
         @test (norm(unit(M)) ≈ 1) == true
-        @test (norm(a) ≈ 1) == false # Again, to be sure that it is still non-normalized
+        @test (norm(ψ) ≈ 1) == false # Again, to be sure that it is still non-normalized
         @test (norm(M) ≈ 1) == false # Again, to be sure that it is still non-normalized
-        normalize!(a)
+        normalize!(ψ)
         normalize!(M)
-        @test (norm(a) ≈ 1) == true
+        @test (norm(ψ) ≈ 1) == true
         @test (norm(M) ≈ 1) == true
-        @test M ≈ a * a'
+        @test M ≈ ψ * ψ'
         @test (unit(qeye(N)) ≈ (qeye(N) / N)) == true
+
+        @testset "Type Inference (normalize)" begin
+            ψ = Qobj(rand(ComplexF64, N))
+            M = ket2dm(ψ)
+            @inferred normalize(ψ)
+            @inferred normalize(M)
+        end
     end
 
     @testset "expectation value" begin
@@ -341,6 +379,19 @@
         ψ = fock(N, 3)
         @test norm(ψ' * a) ≈ 2
         @test expect(a' * a, ψ' * a) ≈ 16
+
+        ρ = rand_dm(N)
+        @test expect(a, ρ) ≈ tr(a * ρ)
+        @test variance(a, ρ) ≈ tr(a^2 * ρ) - tr(a * ρ)^2
+
+        @testset "Type Inference (expect)" begin
+            @inferred expect(a, ψ)
+            @inferred expect(a, ψ')
+            @inferred variance(a, ψ)
+            @inferred variance(a, ψ')
+            @inferred expect(a, ρ)
+            @inferred variance(a, ρ)
+        end
     end
 
     @testset "get coherence" begin
@@ -350,6 +401,11 @@
         ρ = ket2dm(ψ)
         α, δρ = get_coherence(ρ)
         @test isapprox(abs(α), 3, atol = 1e-5) && abs2(δρ[1, 1]) > 0.999
+
+        @testset "Type Inference (get_coherence)" begin
+            @inferred get_coherence(ψ)
+            @inferred get_coherence(ρ)
+        end
     end
 
     @testset "SVD and Schatten p-norm" begin
@@ -361,6 +417,13 @@
         @test svdvals(vs)[1] ≈ √(vs' * vs)
         @test norm(Md, 1) ≈ sum(sqrt, abs.(eigenenergies(Md' * Md))) atol = 1e-6
         @test norm(Ms, 1) ≈ sum(sqrt, abs.(eigenenergies(Ms' * Ms))) atol = 1e-6
+
+        @testset "Type Inference (SVD and Schatten p-norm)" begin
+            @inferred svdvals(vd)
+            @inferred svdvals(vs)
+            @inferred norm(Md, 1)
+            @inferred norm(Ms, 1)
+        end
     end
 
     @testset "purity" begin
@@ -374,6 +437,13 @@
         @test purity(ψd) ≈ norm(ψd)^2 ≈ 1.0
         @test purity(ρ1) ≈ 1.0
         @test (1.0 / N) <= purity(ρ2) <= 1.0
+
+        @testset "Type Inference (purity)" begin
+            @inferred purity(ψ)
+            @inferred purity(ψd)
+            @inferred purity(ρ1)
+            @inferred purity(ρ2)
+        end
     end
 
     @testset "trace distance" begin
@@ -386,6 +456,13 @@
         @test tracedist(ρz0, ψz1) ≈ 1.0
         @test tracedist(ψz1, ρz0) ≈ 1.0
         @test tracedist(ρz0, ρz1) ≈ 1.0
+
+        @testset "Type Inference (trace distance)" begin
+            @inferred tracedist(ψz0, ψx0)
+            @inferred tracedist(ρz0, ψz1)
+            @inferred tracedist(ψz1, ρz0)
+            @inferred tracedist(ρz0, ρz1)
+        end
     end
 
     @testset "sqrt and fidelity" begin
@@ -397,6 +474,13 @@
         @test sqrtm(M0) ≈ sqrtm(sparse_to_dense(M0))
         @test isapprox(fidelity(M0, M1), fidelity(ψ1, M0); atol = 1e-6)
         @test isapprox(fidelity(ψ1, ψ2), fidelity(ket2dm(ψ1), ket2dm(ψ2)); atol = 1e-6)
+
+        @testset "Type Inference (sqrt and fidelity)" begin
+            @inferred sqrtm(M0)
+            @inferred fidelity(M0, M1)
+            @inferred fidelity(ψ1, M0)
+            @inferred fidelity(ψ1, ψ2)
+        end
     end
 
     @testset "log, exp, sinm, cosm" begin
@@ -448,6 +532,58 @@
         ρ2 = dense_to_sparse(ρ1)
         @test tidyup(ρ2, tol) != ρ2
         @test dense_to_sparse(tidyup(ρ1, tol)) == tidyup(ρ2, tol)
+
+        @testset "Type Inference (tidyup)" begin
+            @inferred tidyup(ψ1, tol)
+            @inferred tidyup(ρ1, tol)
+            @inferred tidyup(ψ2, tol)
+            @inferred tidyup(ρ2, tol)
+        end
+    end
+
+    @testset "ptrace" begin
+        g = fock(2, 1)
+        e = fock(2, 0)
+        α = sqrt(0.7)
+        β = sqrt(0.3) * 1im
+        ψ = α * kron(g, e) + β * kron(e, g)
+
+        ρ1 = ptrace(ψ, 1)
+        ρ2 = ptrace(ψ, 2)
+        @test ρ1.data ≈ [0.3 0.0; 0.0 0.7] atol = 1e-10
+        @test ρ2.data ≈ [0.7 0.0; 0.0 0.3] atol = 1e-10
+
+        ψ_d = ψ'
+
+        ρ1 = ptrace(ψ_d, 1)
+        ρ2 = ptrace(ψ_d, 2)
+        @test ρ1.data ≈ [0.3 0.0; 0.0 0.7] atol = 1e-10
+        @test ρ2.data ≈ [0.7 0.0; 0.0 0.3] atol = 1e-10
+
+        ρ = ket2dm(ψ)
+        ρ1 = ptrace(ρ, 1)
+        ρ2 = ptrace(ρ, 2)
+        @test ρ1.data ≈ [0.3 0.0; 0.0 0.7] atol = 1e-10
+        @test ρ2.data ≈ [0.7 0.0; 0.0 0.3] atol = 1e-10
+
+        ψ1 = normalize(g + 1im * e)
+        ψ2 = normalize(g + e)
+        ρ1 = ket2dm(ψ1)
+        ρ2 = ket2dm(ψ2)
+        ρ = kron(ρ1, ρ2)
+        ρ1_ptr = ptrace(ρ, 1)
+        ρ2_ptr = ptrace(ρ, 2)
+        @test ρ1.data ≈ ρ1_ptr.data atol = 1e-10
+        @test ρ2.data ≈ ρ2_ptr.data atol = 1e-10
+
+        @testset "Type Inference (ptrace)" begin
+            @inferred ptrace(ρ, 1)
+            @inferred ptrace(ρ, 2)
+            @inferred ptrace(ψ_d, 1)
+            @inferred ptrace(ψ_d, 2)
+            @inferred ptrace(ψ, 1)
+            @inferred ptrace(ψ, 2)
+        end
     end
 
     @testset "permute" begin
@@ -484,5 +620,11 @@
         @test_throws ArgumentError permute(bra_bdca, wrong_order2)
         @test_throws ArgumentError permute(op_bdca, wrong_order1)
         @test_throws ArgumentError permute(op_bdca, wrong_order2)
+
+        @testset "Type Inference (permute)" begin
+            @inferred permute(ket_bdca, (2, 4, 3, 1))
+            @inferred permute(bra_bdca, (2, 4, 3, 1))
+            @inferred permute(op_bdca, (2, 4, 3, 1))
+        end
     end
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index a1f9f8df..ea61379d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -22,7 +22,7 @@ core_tests = [
     "quantum_objects.jl",
     "states_and_operators.jl",
     "steady_state.jl",
-    "time_evolution_and_partial_trace.jl",
+    "time_evolution.jl",
     "wigner.jl",
 ]
 
diff --git a/test/states_and_operators.jl b/test/states_and_operators.jl
index 8bf3718a..08aa0c95 100644
--- a/test/states_and_operators.jl
+++ b/test/states_and_operators.jl
@@ -202,10 +202,6 @@
         @test U3.dims == U4.dims == [5, 5]
 
         @test_throws ArgumentError rand_unitary(20, :wrong)
-        @testset "Type Inference" begin
-            @inferred rand_unitary(10, :haar)
-            @inferred rand_unitary(10, :exp)
-        end
     end
 
     @testset "Spin-j operators" begin
@@ -349,4 +345,90 @@
         @test_throws DimensionMismatch sprepost(A_wrong1, A_wrong2)
         @test_throws DimensionMismatch sprepost(A_wrong1, A_wrong3)
     end
+
+    @testset "Type Inference (States)" begin
+        @inferred zero_ket(4)
+        @inferred zero_ket((2, 2))
+
+        @inferred fock(4, 1, dims = (2, 2), sparse = Val(true))
+        @inferred fock(4, 1, dims = (2, 2), sparse = Val(false))
+        @inferred basis(4, 1, dims = (2, 2))
+
+        @inferred coherent(5, 0.25im)
+
+        @inferred rand_ket(10)
+        @inferred rand_ket((2, 5))
+
+        @inferred fock_dm(4; dims = (2, 2), sparse = Val(true))
+        @inferred fock_dm(4; dims = (2, 2), sparse = Val(false))
+
+        @inferred coherent_dm(5, 0.25im)
+
+        @inferred thermal_dm(10, 0.123)
+        @inferred thermal_dm(10, 0.123; sparse = Val(true))
+
+        @inferred maximally_mixed_dm(4)
+        @inferred maximally_mixed_dm((2, 2))
+
+        @inferred rand_dm((2, 2))
+        @inferred rand_dm(10, rank = 5)
+
+        @inferred spin_state(3.5, 1.5)
+        @inferred spin_coherent(0.5, π, 2π)
+
+        @inferred bell_state(Val(0), Val(0))
+        @inferred bell_state(Val(0), Val(1))
+        @inferred bell_state(Val(1), Val(0))
+        @inferred bell_state(Val(1), Val(1))
+        @inferred singlet_state()
+        @inferred triplet_states()
+        @inferred w_state(Val(2))
+
+        @inferred ghz_state(Val(2))
+        @inferred ghz_state(Val(3))
+        @inferred ghz_state(Val(3); d = 3)
+    end
+
+    @testset "Type Inference (Operators)" begin
+        @inferred rand_unitary(10, Val(:haar))
+        @inferred rand_unitary(10, Val(:exp))
+
+        a = destroy(20)
+        a_d = create(20)
+        @inferred commutator(a, a_d)
+
+        @inferred destroy(20)
+        @inferred create(20)
+        @inferred num(20)
+        @inferred displace(20, 0.5 + 0.5im)
+        @inferred squeeze(20, 0.5 + 0.5im)
+        @inferred position(20)
+        @inferred momentum(20)
+
+        @inferred phase(20, 0.5)
+
+        @inferred jmat(2.5)
+        @inferred jmat(2.5, Val(:x))
+        @inferred jmat(2.5, Val(:y))
+        @inferred jmat(2.5, Val(:z))
+
+        @inferred sigmap()
+        @inferred sigmam()
+        @inferred sigmax()
+        @inferred sigmay()
+        @inferred sigmaz()
+
+        @inferred eye(20)
+
+        @inferred fdestroy(Val(10), 4)
+        @inferred fcreate(Val(10), 4)
+
+        @inferred projection(20, 5, 3)
+
+        @inferred tunneling(20, 1, sparse = Val(false))
+        @inferred tunneling(20, 1, sparse = Val(true))
+
+        @inferred qft(20)
+        @inferred qft((2, 10))
+    end
 end
diff --git a/test/steady_state.jl b/test/steady_state.jl
index 1f497906..924cc090 100644
--- a/test/steady_state.jl
+++ b/test/steady_state.jl
@@ -22,8 +22,7 @@
     ρ_ss = steadystate(H, c_ops, solver = solver)
     @test tracedist(rho_me, ρ_ss) < 1e-4
 
-    import LinearSolve: KrylovJL_GMRES
-    solver = SteadyStateLinearSolver(alg = KrylovJL_GMRES())
+    solver = SteadyStateLinearSolver()
     ρ_ss = steadystate(H, c_ops, solver = solver)
     @test tracedist(rho_me, ρ_ss) < 1e-4
 
@@ -31,6 +30,33 @@
     ρ_ss = steadystate(H, c_ops, solver = solver)
     @test tracedist(rho_me, ρ_ss) < 1e-4
 
+    @testset "Type Inference (steadystate)" begin
+        L = liouvillian(H, c_ops)
+
+        solver = SteadyStateODESolver()
+        @inferred steadystate(H, psi0, t_l[end], c_ops, solver = solver)
+        @inferred steadystate(L, psi0, t_l[end], solver = solver)
+
+        solver = SteadyStateDirectSolver()
+        @inferred steadystate(H, c_ops, solver = solver)
+        @inferred steadystate(L, solver = solver)
+
+        solver = SteadyStateLinearSolver()
+        @inferred steadystate(H, c_ops, solver = solver)
+        @inferred steadystate(L, solver = solver)
+
+        solver = SteadyStateLinearSolver()
+        @inferred steadystate(H, c_ops, solver = solver)
+        @inferred steadystate(L, solver = solver)
+
+        solver = SteadyStateEigenSolver()
+        @inferred steadystate(H, c_ops, solver = solver)
+        @inferred steadystate(L, solver = solver)
+
+        @inferred steadystate(H, c_ops)
+        @inferred steadystate(L)
+    end
+
     H = a_d * a
     H_t = 0.1 * (a + a_d)
     c_ops = [sqrt(0.1) * a]
@@ -45,4 +71,16 @@
 
     @test abs(sum(sol_me.expect[1, end-100:end]) / 101 - expect(e_ops[1], ρ_ss1)) < 1e-3
     @test abs(sum(sol_me.expect[1, end-100:end]) / 101 - expect(e_ops[1], ρ_ss2)) < 1e-3
+
+    @testset "Type Inference (steadystate_floquet)" begin
+        @inferred steadystate_floquet(H, -1im * 0.5 * H_t, 1im * 0.5 * H_t, 1, c_ops, solver = SSFloquetLinearSystem())
+        @inferred steadystate_floquet(
+            H,
+            -1im * 0.5 * H_t,
+            1im * 0.5 * H_t,
+            1,
+            c_ops,
+            solver = SSFloquetEffectiveLiouvillian(),
+        )
+    end
 end
diff --git a/test/time_evolution_and_partial_trace.jl b/test/time_evolution.jl
similarity index 97%
rename from test/time_evolution_and_partial_trace.jl
rename to test/time_evolution.jl
index bb0fe569..b4d097de 100644
--- a/test/time_evolution_and_partial_trace.jl
+++ b/test/time_evolution.jl
@@ -16,8 +16,6 @@
         sol3 = sesolve(H, psi0, t_l, e_ops = e_ops, saveat = t_l, progress_bar = Val(false))
         sol_string = sprint((t, s) -> show(t, "text/plain", s), sol)
         @test sum(abs.(sol.expect[1, :] .- sin.(η * t_l) .^ 2)) / length(t_l) < 0.1
-        @test ptrace(sol.states[end], 1) ≈ ptrace(ket2dm(sol.states[end]), 1)
-        @test ptrace(sol.states[end]', 1) ≈ ptrace(sol.states[end], 1)
         @test length(sol.states) == 1
         @test size(sol.expect) == (length(e_ops), length(t_l))
         @test length(sol2.states) == length(t_l)
diff --git a/test/wigner.jl b/test/wigner.jl
index 3a4d3ee2..ae30b188 100644
--- a/test/wigner.jl
+++ b/test/wigner.jl
@@ -20,4 +20,11 @@
     wig2 = maximum(wig) * reshape(kron(wig_tmp1, wig_tmp2), 300, 300)
 
     @test sqrt(sum(abs.(wig2 .- wig)) / length(wig)) < 0.1
+
+    @testset "Type Inference (wigner)" begin
+        @inferred wigner(ψ, xvec, yvec, solver = WignerLaguerre(tol = 1e-6))
+        @inferred wigner(ρ, xvec, yvec, solver = WignerLaguerre(parallel = false))
+        @inferred wigner(ρ, xvec, yvec, solver = WignerLaguerre(parallel = true))
+        @inferred wigner(ψ, xvec, yvec, solver = WignerClenshaw())
+    end
 end