Project.toml

@@ -14,6 +14,7 @@ JLD = "4138dd39-2aa7-5051-a626-17a0bb65d9c8"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
+OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 PhysModels = "1ff164ad-7f5d-45f3-982a-6d162405c7c1"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"

logtrexp.jl

-using StatsFuns, ChainRules
-
-""" `symeigen` : manually symmetrize M before the eigen decomposition
-    For CUDA dense matrix eigen solver `CUSOLVER.syevd!` and `CUSOLVER.heevd!`:
-        'N'/'V': return eigenvalues/both eigenvalues and eigenvectors 
-        'U'/'L': Upper/Lower triangle of `M` is stored.
-"""
-function symeigen(M::AbstractMatrix)
-    e, v = M |> symmetrize |> eigen
-    return e, v
-end
-
-"""
-    `ln(Tr[exp(A)])` function
-"""
-function logtrexp(A::AbstractMatrix)
-    val, _ = symeigen(A)
-    return logsumexp(val)
-end
-
-
-"""
-    rrule for logtrexp function, eltype(M) <: Real
-"""    
-function ChainRules.rrule(::typeof(logtrexp), M::AbstractArray{T}) where T<:Real
-    e, v = symeigen(M)
-    y = logsumexp(e)
-    function logtrexp_pullback(ȳ)
-        Λ = map(x->exp(x-y), e) 
-        ∂y_∂M = v * diagm(Λ) * v' |> transpose
-        M̄ = ȳ * ∂y_∂M
-        return ChainRules.NoTangent(), M̄
-    end
-    return y, logtrexp_pullback
-end

src/FiniteTLanczos.jl

@@ -5,13 +5,19 @@ using Parameters, Reexport
 using SparseArrays
 using LinearAlgebra, KrylovKit
 using Random; Random.seed!()
+using CUDA; CUDA.allowscalar(false)
+using OMEinsum
 
-import LinearAlgebra:Eigen
+@reexport import LinearAlgebra:Eigen, diagm
 
 using PhysModels 
 
 #Operator.jl
 export Operator
+#solver.jl
+export Device, CPU, GPU, 
+       solver_function,
+       cpu_solver, gpu_solver
 
 #utilities.jl
 export replicate,
@@ -38,10 +44,20 @@ export Full_ED,
        FullSampling_FTLM
 
 #TraceEstimator.jl
-export TraceEstimator
+export TraceEstimator #, evaluate
+
+#Thermaldynamic.jl
+#export partitian, 
+#       free_energy,
+#       energy,
+#       specific_heat,
+#       entropy,
+#       thermal_average
+
 
 include("Operator.jl")
 
+include("solver.jl")
 include("utilities.jl")
 include("eigensolver.jl")
 

src/FullOrthoLanczos.jl

@@ -64,29 +64,33 @@ Base.hcat(ortho_basis::Nothing, v) = v
 function itFOLM(M, ortho_basis::T; 
                 init_vector::AbstractVector, 
                 Nk::Int = 50) where T<:Union{Nothing, AbstractArray}
-    Nm = size(M, 1)
+    Ns = size(M, 1)
     if T <: Nothing 
         Ne = 0
     elseif T <: AbstractVector 
         Ne = 1 
     else 
-        Ne = size(ortho_basis)[2]
+        Ne = size(ortho_basis,2)
     end
 
-    Nk = min(Nk, Nm - Ne) 
-    dv, ev = zeros(Nk), zeros(Nk - 1)
+    Nk = min(Nk, Ns - Ne) 
+    if Nk == 0 
+        return itFOLMResult(init_vector, SymTridiagonal([0.], [0.]), zeros(Ns,1))
+    else
+        dv, ev = zeros(Nk), zeros(Nk - 1)
 
-    w = icgs(init_vector, ortho_basis); w = w/norm(w)
-    for i = 1 : Nk
-        ortho_basis = hcat(ortho_basis, w) #mind the order
-        r = M * w
-        dv[i] = w' * r
-        r = icgs(r, ortho_basis)
-        b = norm(r)
-        w = r/b
-        if i < Nk ev[i] = b end
+        w = icgs(init_vector, ortho_basis); w = w/norm(w)
+        for i = 1 : Nk
+            ortho_basis = hcat(ortho_basis, w) #mind the order
+            r = M * w
+            dv[i] = w' * r
+            r = icgs(r, ortho_basis)
+            b = norm(r)
+            w = r/b
+            if i < Nk ev[i] = b end
+        end
+        return itFOLMResult(init_vector, SymTridiagonal(dv, ev), ortho_basis[:, Ne+1:end]) #mind the order
     end
-    return itFOLMResult(init_vector, SymTridiagonal(dv, ev), ortho_basis[:, Ne+1:end]) #mind the order
 end
 
 itFOLM(M; init_vector::AbstractVector, Nk::Int = 50) = itFOLM(M, nothing, init_vector = init_vector, Nk = Nk)

src/GeneralFTLM.jl

@@ -26,30 +26,42 @@ end
     `Nr`: number of random sampling
     `Nk`: dimension of Krylov subspace
     `Ne`: number of exact eigen states used
+    `which`: order of eigenvalues, default = :SR, 
+        where `:SR` means smallest real part and `:LR` means largest real part.
 """
-@with_kw struct FTLMOptions{Tf<:Function, Tn<:Integer}
-    distr::Tf = Rademacher
+@with_kw struct FTLMOptions{Tdistr<:Function, 
+                            Tdevice<:Device, 
+                            Tn<:Integer}
+    distr::Tdistr = Rademacher
     Nr::Tn = 100
     Nk::Tn = 50
     Ne::Tn = 10
+    which::Symbol = :SR
+    device::Tdevice = CPU
 end
 
 
 """
     Full exact diagonalization method
 """
-function Full_ED(M, expr::Vararg{Function,N}; 
+function Full_ED(M, expr::Vararg{T,N}; 
                  options::FTLMOptions = FTLMOptions()
-                 ) where {N}
+                 ) where {T<:Function, N}
+    @unpack device = options
+    solver = solver_function(device)
+    M = solver(x->x, M)
     val, _ = eigensolver(M)
     func = f -> map(e->f(e), val) |> sum
     return map(func, expr)
 end
 
-function Full_ED(M, Op::Vararg{Operator,N}; 
-                 expr::Function, 
+function Full_ED(M, Op::Vararg{T,N}; 
+                 expr::Te, 
                  options::FTLMOptions=FTLMOptions()
-                 ) where {N}
+                ) where {T<:Operator, N, Te<:Function}
+    @unpack device = options
+    solver = solver_function(device)
+    M = solver(x->x, M)
     val, vec = eigensolver(M)
     Z = map(expr, val) |> sum
     f = x -> map((i -> expr(val[i]) * vec[:,i]' * x.O * vec[:,i]), 1:length(val)) |> sum
@@ -61,14 +73,18 @@ end
 """
     Simple FTLM algorithom, where `M` is required to be hermitian.
 """
-function simple_FTLM(M, expr::Vararg{Function,N}; 
+function simple_FTLM(M, expr::Vararg{T,N}; 
                      options::FTLMOptions = FTLMOptions()
-                    ) where {N}
-    @unpack distr, Nr, Nk = options
+                    ) where {T<:Function, N}
+    @unpack distr, Nr, Nk, device = options
+    solver = solver_function(device)
+    M = solver(x->x, M)
+
     Ns = size(M, 1)
-    res = Tuple(zeros(N))
+    res = zeros(N)
     for r = 1: Nr
         v0 = random_unit_vector(Ns, distr)
+        v0 = solver(x->x, v0)
         @unpack values, weight = itFOLM(M, init_vector = v0, Nk = Nk) |> eigensolver
         func = f -> map((e,w)->f(e)* w * w', values, weight) |> sum
         res = map(+, res, map(func, expr))
@@ -76,16 +92,20 @@ function simple_FTLM(M, expr::Vararg{Function,N};
     return map(x -> x * Ns/Nr, res)
 end
 
-function simple_FTLM(M, Op::Vararg{Operator,N};
-                     expr::Function,
+function simple_FTLM(M, Op::Vararg{T,N};
+                     expr::Te,
                      options::FTLMOptions = FTLMOptions()
-                    ) where {N}
-    @unpack distr, Nr, Nk = options
+                    ) where {T<:Operator,N,Te<:Function}
+    @unpack distr, Nr, Nk, device = options
+    solver = solver_function(device)
+    M = solver(x->x, M)
+
     Ns = size(M,1)
-    res = Tuple(zeros(N))
+    res = zeros(N)
     Z = 0.0
     for r = 1: Nr
         v0 = random_unit_vector(Ns, distr)
+        v0 = solver(x->x, v0)
         @unpack weight, values, vectors = itFOLM(M, init_vector = v0, Nk = Nk) |> eigensolver
         Z += map((e,w)-> expr(e) * w * w', values, weight) |> sum
         operator_average = x -> begin
@@ -102,20 +122,26 @@ end
 """
     Replaced FTLM algorithom
 """
-function replaced_FTLM(M, expr::Vararg{Function,N}; 
+function replaced_FTLM(M, expr::Vararg{T,N}; 
                        options::FTLMOptions = FTLMOptions()
-                       ) where {N}
-    @unpack distr, Nr, Nk, Ne = options
+                       ) where {T<:Function,N}
+    @unpack distr, Nr, Nk, Ne, device, which = options
+    solver = solver_function(device)
+    M = solver(x->x, M)
+
     Ns = size(M, 1)
-    res = Tuple(zeros(N))
+    res = zeros(N)
 
-    Ne = min(Ne, Ns-1)
-    evals, evecs, _ = eigsolve(M, Ns, Ne, :SR, ishermitian = true)
+    Ne = min(Ne, Ns)
+    krylovdim = max(30, Ne+1)
+    evals, evecs, _ = eigsolve(M, Ns, Ne, which, ishermitian = true, krylovdim=krylovdim)
     evals = evals[1:Ne]
     evecs = hcat(evecs[1:Ne]...)
 
+    Nk = max(Nk, Ne)
     for r = 1: Nr
         v0 = random_unit_vector(Ns, distr)
+        v0 = solver(x->x, v0)
         @unpack weight, values = itFOLM(M, init_vector = v0, Nk = Nk) |> eigensolver
         values = vcat(evals, values[Ne+1:end])
         eweight = map(i-> v0' * evecs[:,i], 1:Ne)
@@ -128,22 +154,28 @@ function replaced_FTLM(M, expr::Vararg{Function,N};
 end
 
 
-function replaced_FTLM(M, Op::Vararg{Operator,N};
-                       expr::Function,
+function replaced_FTLM(M, Op::Vararg{T,N};
+                       expr::Te,
                        options::FTLMOptions=FTLMOptions()
-                       ) where {N}
-    @unpack distr, Nr, Nk, Ne = options
+                       ) where {T<:Operator,N,Te<:Function}
+    @unpack distr, Nr, Nk, Ne, device, which = options
+    solver = solver_function(device)
+    M = solver(x->x, M)
+
     Ns = size(M, 1)
-    res = Tuple(zeros(N))
+    res = zeros(N)
 
-    Ne = min(Ne, Ns-1)
-    evals, evecs, _ = eigsolve(M, Ns, Ne, :SR, ishermitian = true)
+    Ne = min(Ne, Ns)
+    krylovdim = max(30, Ne+1)
+    evals, evecs, _ = eigsolve(M, Ns, Ne, which, ishermitian = true, krylovdim=krylovdim)
     evals = evals[1:Ne]
     evecs = hcat(evecs[1:Ne]...)
 
     Z = 0.0
+    Nk = max(Nk, Ne)
     for r = 1: Nr
         v0 = random_unit_vector(Ns, distr)
+        v0 = solver(x->x, v0)
         @unpack weight, values, vectors = itFOLM(M, init_vector = v0, Nk = Nk) |> eigensolver
         values = vcat(evals, values[Ne+1:end])
         vectors = hcat(evecs, vectors[:, Ne+1:end])
@@ -166,37 +198,46 @@ end
     Orthogonalized FTLM algorithom
     `M`: Symmetric matrix
 """
-function orthogonalized_FTLM(M, expr::Vararg{Function,N}; 
+function orthogonalized_FTLM(M, expr::Vararg{T,N}; 
                              options::FTLMOptions=FTLMOptions()
-                            ) where {N}
-    @unpack distr, Nr, Nk, Ne = options
+                            ) where {T<:Function,N}
+    @unpack distr, Nr, Nk, Ne, device, which = options
+    solver = solver_function(device)
+    M = solver(x->x, M)
+
     Ns = size(M, 1)
-    Ne = min(Ne, Ns-1)
-    evals, evecs, _ = eigsolve(M, Ns, Ne, :SR, ishermitian = true)
+    Ne = min(Ne, Ns)
+    krylovdim = max(30, Ne+1)
+    evals, evecs, _ = eigsolve(M, Ns, Ne, which, ishermitian = true, krylovdim=krylovdim)
     evals = evals[1:Ne]
     evecs = hcat(evecs[1:Ne]...)
 
     func1 = f -> map(e->f(e), evals) |> sum
     res1 = map(func1, expr)
 
-    res2 = Tuple(zeros(N))
+    res2 = zeros(N)
     for r = 1: Nr
         v0 = random_unit_vector(Ns, distr)
+        v0 = solver(x->x, v0)
         @unpack weight, values = itFOLM(M, evecs, init_vector = v0, Nk = Nk) |> eigensolver
         func2 = f -> map((e,w)->f(e)* w * w', values, weight) |> sum
         res2 = map(+, res2, map(func2, expr))
     end
-    return map((x1, x2) -> x1 + x2 * Ns/Nr, res1, res2)
+    return map((x1, x2) -> x1 + x2 * (Ns-Ne)/Nr, res1, res2)
 end
 
-function orthogonalized_FTLM(M, Op::Vararg{Operator,N};
-                             expr::Function,
+function orthogonalized_FTLM(M, Op::Vararg{T,N};
+                             expr::Te,
                              options::FTLMOptions=FTLMOptions()
-                            ) where {N}
-    @unpack distr, Nr, Nk, Ne = options
+                            ) where {T<:Operator,N,Te<:Function}
+    @unpack distr, Nr, Nk, Ne, device, which = options
+    solver = solver_function(device)
+    M = solver(x->x, M)
+
     Ns = size(M, 1)
-    Ne = min(Ne, Ns-1)
-    evals, evecs, _ = eigsolve(M, Ns, Ne, :SR, ishermitian = true)
+    Ne = min(Ne, Ns)
+    krylovdim = max(30, Ne+1)
+    evals, evecs, _ = eigsolve(M, Ns, Ne, which, ishermitian = true, krylovdim=krylovdim)
     evals = evals[1:Ne]
     evecs = hcat(evecs[1:Ne]...)
 
@@ -205,9 +246,10 @@ function orthogonalized_FTLM(M, Op::Vararg{Operator,N};
     res1 = map(func1, Op)
 
     Z2 = 0.0
-    res2 = Tuple(zeros(N))
+    res2 = zeros(N)
     for r = 1: Nr
         v0 = random_unit_vector(Ns, distr)
+        v0 = solver(x->x, v0)
         @unpack weight, values, vectors = itFOLM(M, evecs, init_vector = v0, Nk = Nk) |> eigensolver
         Z2 += map((e,w)-> expr(e) * w * w', values, weight) |> sum
         operator_average = x -> begin
@@ -216,71 +258,59 @@ function orthogonalized_FTLM(M, Op::Vararg{Operator,N};
         end
         res2 = map(+, res2, map(operator_average, Op))
     end
-    Z = Z1 + Z2 * Ns/Nr
-    return map((x1, x2) -> x1/Z + x2/Z * Ns/Nr, res1, res2)
+    Z = Z1 + Z2 * (Ns-Ne)/Nr
+    return map((x1, x2) -> (x1 + x2*(Ns-Ne)/Nr)/Z, res1, res2)
 end
 
 
 """
     Full sampling simple_FTLM
 """
-function FullSampling_FTLM(M, expr::Vararg{Function,N};
+function FullSampling_FTLM(M, expr::Vararg{T,N};
                            options::FTLMOptions = FTLMOptions()
-                          ) where {N}
-    ortho_basis = basis_generate(Ns)
-    @unpack Nk, Ne = options
-
+                          ) where {T<:Function,N}
     Ns = size(M, 1)
-    Ne = min(Ne, Ns-1)
-    evals, evecs, _ = eigsolve(M, Ns, Ne, :SR, ishermitian = true)
-    evals = evals[1:Ne]
-    evecs = hcat(evecs[1:Ne]...)
-
-    func1 = f -> map(e->f(e), evals) |> sum
-    res1 = map(func1, expr)
+    ortho_basis = basis_generate(Ns)
+    @unpack Nk, which, device = options
+    solver = solver_function(device)
+    M = solver(x->x, M)
 
     #Full FullSampling_FTLM
-    res2 = Tuple(zeros(N))
+    res = zeros(N)
     for i = 1: Ns
         v0 = ortho_basis[:, i]
-        @unpack weight, values = itFOLM(M, evecs, init_vector = v0, Nk = Nk) |> eigensolver
+        v0 = solver(x->x, v0)
+        @unpack weight, values = itFOLM(M, init_vector = v0, Nk = Nk) |> eigensolver
         func2 = f -> map((e,w)->f(e)* w * w', values, weight) |> sum
-        res2 = map(+, res2, map(func2, expr))
+        res = map(+, res, map(func2, expr))
     end
-    return map(+, res1, res2)
+    return res
 end
 
-function FullSampling_FTLM(M, Op::Vararg{Operator,N}; 
-                           expr::Function,
+function FullSampling_FTLM(M, Op::Vararg{T,N}; 
+                           expr::Te,
                            options::FTLMOptions = FTLMOptions()
-                          ) where {N}
-    ortho_basis = basis_generate(Ns)
-    @unpack Nk, Ne = options
-
+                          ) where {T<:Operator,N,Te<:Function}
     Ns = size(M, 1)
-    Ne = min(Ne, Ns-1)
-    evals, evecs, _ = eigsolve(M, Ns, Ne, :SR, ishermitian = true)
-    evals = evals[1:Ne]
-    evecs = hcat(evecs[1:Ne]...)
-
-    Z1 = map(func, evals) |> sum
-    func1 = x -> map((i -> expr(evals[i]) * evecs[:,i]' * x.O * evecs[:,i]), 1:length(evals)) |> sum
-    res1 = map(func1, Op)
+    ortho_basis = basis_generate(Ns)
+    @unpack Nk, which, device = options
+    solver = solver_function(device)
+    M = solver(x->x, M)
 
-    Z2 = 0.0
-    res2 = Tuple(zeros(N))
+    Z = 0.0
+    res = zeros(N)
     for i = 1: Ns
         v0 = ortho_basis[:, i]
-        @unpack weight, values, vectors = itFOLM(M, evecs, init_vector = v0, Nk = Nk) |> eigensolver
-        Z2 += map((e,w)-> expr(e) * w * w', values, weight) |> sum
+        v0 = solver(x->x, v0)
+        @unpack weight, values, vectors = itFOLM(M, init_vector = v0, Nk = Nk) |> eigensolver
+        Z += map((e,w)-> expr(e) * w * w', values, weight) |> sum
         operator_average = x -> begin
             weight2 = map(i -> v0' * x.O * vectors[:,i], 1:Nk)
             return map((e, w1, w2)-> expr(e)* w1 * w2, val, weight, weight2) |> sum
         end
-        res2 = map(+, res2, map(operator_average, Op))
+        res = map(+, res, map(operator_average, Op))
     end
-    Z = Z1 + Z2
-    return map((x1, x2) -> x1/Z + x2/Z, res1, res2)
+    return map(x -> x/Z, res)
 end
 
 #end #module methods

src/TraceEstimator.jl

@@ -7,14 +7,14 @@
 end
 
 
-function evaluate(method::TraceEstimator, M, expr::Vararg{Function,N}) where {N}
+function evaluate(method::TraceEstimator, M, expr::Vararg{T,N}) where {T<:Function,N}
     @unpack estimator, options = method
-    return estimator(M, expr, options=options)
+    return estimator(M, expr..., options=options)
 end
 
-function evaluate(method::TraceEstimator, M, Op::Vararg{Operator,N}; expr::Function) where {N}
+function evaluate(method::TraceEstimator, M, Op::Vararg{T,N}; expr::Te) where {T<:Operator,N,Te<:Function}
     @unpack estimator, options = method
-    return estimator(M, Op, expr = expr, options=method.options)
+    return estimator(M, Op..., expr = expr, options=method.options)
 end
 
 

src/eigensolver.jl

@@ -20,4 +20,24 @@ function eigensolver(M::T; which::Symbol = :SR, kwargs...) where T<:AbstractSpar
     return Eigen(val, vec)
 end
 
+for elty in (:Float32, :Float64)
+    """ 
+    eigensolver for CuMatrix
+    For CUDA dense matrix eigen solver `CUSOLVER.syevd!` and `CUSOLVER.heevd!`:
+        'N'/'V': return eigenvalues/both eigenvalues and eigenvectors 
+        'U'/'L': Upper/Lower triangle of `M` is stored.
+    """
+    @eval eigensolver(M::CuMatrix{$elty}) = CUSOLVER.syevd!('V', 'U', copy(M))
+end
+
+for elty in (:ComplexF32, :ComplexF64)
+    """ 
+    eigensolver for CuMatrix
+    For CUDA dense matrix eigen solver `CUSOLVER.syevd!` and `CUSOLVER.heevd!`:
+        'N'/'V': return eigenvalues/both eigenvalues and eigenvectors 
+        'U'/'L': Upper/Lower triangle of `M` is stored.
+    """
+    @eval eigensolver(M::CuArray{$elty}) = CUSOLVER.heevd!('V', 'U', copy(M))
+end
+
 

src/solver.jl

+@enum Device CPU GPU
+function cpu_solver end
+function gpu_solver end
+function solver_function(d::Device)
+    d == CPU ? (return cpu_solver) : (return cpu_solver)
+end
+
+function cpu_solver(f::Function, M::AbstractArray...)
+    M = map(x->convert(Array, x), M)
+    return f(M...)
+end
+
+function gpu_solver(f::Function, M::AbstractArray...)
+    M = map(x->convert(CuArray, x), M)
+    return f(M...)
+end
+

src/utilities.jl

@@ -6,6 +6,10 @@ macro replicate(ex, n)
     return :(map(_ -> $(esc(ex)), 1:$(esc(n))))
 end
 
+"""
+    diagm(v) function for CuVector
+"""
+LinearAlgebra.diagm(v::CuVector) = ein"i->ii"(v)
 
 """
     Symmetrize a matrix `M`

test/FullOrthoLanczos.jl

 using Test, FiniteTLanczos
-using LinearAlgebra
+using LinearAlgebra, Parameters
 using Random; Random.seed!()
+solver = solver_function(device)
 
 @testset "icgs algorithm" begin
     N = 256
     ortho = Vector{Bool}(undef, N-1)
     for dist in [Gaussian, Rademacher, rand]
         Q = random_unit_vector(N, dist)
+        Q = solver(x->x, Q)
         for i = 1:N-1
             v = rand(N); v = v/norm(v)
+            v = solver(x->x, v)
             u = icgs(v, Q)
             u = u/norm(u)
             #test Orthogonality
@@ -18,7 +21,7 @@ using Random; Random.seed!()
             Q = hcat(Q, u)
         end
         @test all(ortho)
-        @test Q' * Q ≈ Matrix(1.0I, N, N)
+        @test Q' * Q ≈ solver(x->x, Matrix(1.0I, N, N))
     end
 end
 
@@ -26,31 +29,35 @@ end
 @testset "itFOLM algorithm" begin
     N = 256; Nv = 10
     A = randn(N, N) |> symmetrize
-    val, vec = eigensolver(A)
+    A = solver(x->x, A)
+    vals, vecs = eigensolver(A)
 
     for dist in [Gaussian, Rademacher, rand]
         v0 = random_unit_vector(N, dist)
-        res1 = itFOLM(A, init_vector = v0, ncv = N)
+        v0 = solver(x->x, v0)
+        res1 = itFOLM(A, init_vector = v0, Nk = N)
         e1, v1 = eigensolver(res1.symtridiagonal)
-        @test res1.symtridiagonal ≈ res1.krylov_basis' * A * res1.krylov_basis
-        @test e1 ≈ val
+        @test Matrix(res1.symtridiagonal) ≈ Matrix(res1.krylov_basis' * A * res1.krylov_basis)
+        @test solver(x->x, e1) ≈ vals
 
         for Nv in [1, 10]
-            res2 = itFOLM(A, vec[:,1:Nv], init_vector = v0, ncv = N)
+            res2 = itFOLM(A, vecs[:,1:Nv], init_vector = v0, Nk = N)
             e2, v2 = eigensolver(res2.symtridiagonal)
-            @test res2.symtridiagonal ≈ res2.krylov_basis' * A * res2.krylov_basis
-            @test e2 ≈ val[Nv+1 : end]
+            @test Matrix(res2.symtridiagonal) ≈ Matrix(res2.krylov_basis' * A * res2.krylov_basis)
+            @test e2 ≈ Vector(vals)[Nv+1 : end]
         end
     end
 end
 
 
 @testset "itFOLMResult weight: <v0|ψ> " begin
-    N = 256; Ncv = 100
+    N = 256; Nk = 100
     A = randn(N, N) |> symmetrize
+    A = solver(x->x, A)
     for dist in [Gaussian, Rademacher, rand]
         v0 = random_unit_vector(N, dist)
-        res = itFOLM(A, init_vector = v0, ncv = Ncv) |> eigensolver 
+        v0 = solver(x->x, v0)
+        res = itFOLM(A, init_vector = v0, Nk = Nk) |> eigensolver 
         Q1 = res.weight
         Q2 = res.init_vector' * res.vectors |> transpose
         @test Q1 ≈ Q2
@@ -58,23 +65,24 @@ end
 end
 
 
-@testset "Full Rank itFOLM generate eigenvalues and eigenvectors" begin
+@testset "Full Rank itFOLM generate exact eigen pairs" begin
     N = 100
-    R = 100
-    RandM = rand(N, N) |> symmetrize
-    val, vec = eigensolver(RandM)
+    A = rand(N, N) |> symmetrize
+    A = solver(x->x, A)
+    
+    vals, vecs = eigensolver(A)
     for dist in [Gaussian, Rademacher, rand]
-        iseigenval = Vector{Bool}(undef, R)
-        iseigenvec = Vector{Bool}(undef, R)
-        for r=1:R
+        iseigenval = Vector{Bool}(undef, N)
+        iseigenvec = Vector{Bool}(undef, N)
+        for r = 1:N
             v0 = random_unit_vector(N, dist)
-            res = itFOLM(RandM, init_vector = v0, ncv = N) |> eigensolver 
-            kval = res.values
-            kvec = res.vectors
-            iseigenval[r] = (val ≈ kval)
-            iseigenvec[r] = (kvec * diagm(kval) * kvec' ≈ RandM)
+            v0 = solver(x->x, v0)
+            res = itFOLM(A, init_vector = v0, Nk = N) |> eigensolver 
+            @unpack values, vectors = res
+            iseigenval[r] = (vals ≈ values)
+            iseigenvec[r] = (vectors * diagm(values) * vectors' ≈ A)
         end
         @test all(iseigenval)
         @test all(iseigenvec)
     end
-end
\ No newline at end of file
+end

test/GeneralFTLM.jl

 using Test, FiniteTLanczos
-using LinearAlgebra, KrylovKit, LogExpFunctions
+using LinearAlgebra, KrylovKit, LogExpFunctions, Parameters
 using Random; Random.seed!()
+solver = solver_function(device)
 
-function logtrexp(A::AbstractMatrix)
-    val, _ = eigensolver(A)
-    return logsumexp(val)
+
+function logtrexp(M::AbstractMatrix)
+    vals, _ = eigensolver(M)
+    return logsumexp(vals)
 end
 
-function logtrexp(M, method::Function)
+function logtrexp(M, estimator::TraceEstimator)
     e0, _, _ = eigsolve(M, size(M,1), 1, :LR, ishermitian = true)
     e0 = e0[1]
     expr = e -> exp(e .- e0)
-    res = method(M, expr)[1]
+    @unpack estimator, options = estimator
+    options = FTLMOptions(options, which = :LR)
+    res = FiniteTLanczos.evaluate(TraceEstimator(estimator, options), M, expr)[1]
     return log(res) + e0
 end
 
 N = 256
 M = rand(N, N) |> symmetrize
+M = solver(x->x, M)
 
-options = FTLMOptions()
-rtol = 0.01
+log_trexp = logtrexp(M)
+@testset "Full_ED method" begin
+    estimator = TraceEstimator(Full_ED, FTLMOptions(FTLMOptions(), device = device))
+    res = logtrexp(M, estimator)
+    @test res ≈ log_trexp
+end
 
-ftlm_function = (M, expr) -> simple_FTLM(M,expr, options=options)
-rftlm_function = (M, expr) -> replaced_FTLM(M, expr, options=options)
-oftlm_function = (M, expr) -> replaced_FTLM(M, expr, options = options)
+@testset "simple_FTLM method" begin
+    Ns = size(M,1)
+    options = FTLMOptions(FTLMOptions(), Nk=Ns, device = device)
+    estimator = TraceEstimator(FullSampling_FTLM, options)
+    res = logtrexp(M, estimator)
+    @test res ≈ log_trexp
 
-y0 = logtrexp(M)
-y_fed = logtrexp(M, Full_ED)
-y_fslm = logtrexp(M, FullSampling_FTLM)
-y_ftlm = logtrexp(M, ftlm_function)
-y_rftlm = logtrexp(M, rftlm_function)
-y_oftlm = logtrexp(M, oftlm_function)
+    options = FTLMOptions(FTLMOptions(), device = device)
+    estimator = TraceEstimator(simple_FTLM, options)
+    res = logtrexp(M, estimator)
 
-@testset "Full_ED" begin 
-    @test y_fed == y0 
+    @unpack Nr = options
+    rtol = 1/sqrt(min(Nr, Ns))
+    @test ≈(res, log_trexp, rtol = rtol)
 end
 
-@testset "FullSampling_FTLM" begin 
-    @test y_fslm ≈ y0 
-end
+@testset "orthogonalized_FTLM method" begin
+    Ns = size(M,1)
+    options = FTLMOptions(FTLMOptions(), Ne = Ns , device = device)
+    estimator = TraceEstimator(orthogonalized_FTLM, options)
+    res = logtrexp(M, estimator)
+    @test res ≈ log_trexp
 
-@testset "simple_FTLM" begin 
-    @test ≈(y_ftlm, y0, rtol = rtol) 
-end
+    for Ne in [1, div(Ns,3), div(Ns,2), Ns-1]
+        options = FTLMOptions(FTLMOptions(), Ne = Ne, device = device)
+        estimator = TraceEstimator(orthogonalized_FTLM, options)
+        res = logtrexp(M, estimator)
 
-@testset "replaced_FTLM" begin 
-    @test ≈(y_rftlm, y0, rtol = rtol) 
+        @unpack Nr = options
+        rtol = 1/Nr
+        @test ≈(res, log_trexp, rtol = rtol)
+    end
 end
 
-@testset "orthogonalized_FTLM" begin 
-    @test ≈(y_oftlm, y0, rtol = rtol) 
+
+@testset "replaced_FTLM method" begin
+    Ns = size(M,1)
+    options = FTLMOptions(FTLMOptions(), Ne = Ns , device = device)
+    estimator = TraceEstimator(replaced_FTLM, options)
+    res = logtrexp(M, estimator)
+    @test ≈(res, log_trexp, rtol = 0.1)
+
+    for Ne in [1, div(Ns,3), div(Ns,2), Ns-1]
+        options = FTLMOptions(FTLMOptions(), Ne = Ne, device = device)
+        estimator = TraceEstimator(replaced_FTLM, options)
+        res = logtrexp(M, estimator)
+
+        @unpack Nr = options
+        rtol = 1/sqrt(min(Nr, Ns))
+        @test ≈(res, log_trexp, rtol = rtol)
+    end
 end
 
+
+
+
+

test/cpu_test.jl

+using Test, FiniteTLanczos
+device = CPU
+
+@testset "FullOrthoLanczos.jl" begin
+    include("FullOrthoLanczos.jl")
+end
+
+@testset "GeneralFTLM.jl" begin
+    include("GeneralFTLM.jl")
+end
\ No newline at end of file

test/gpu_test.jl

+using CUDA; CUDA.allowscalar(false)
+using Test, FiniteTLanczos
+device = GPU
+
+@testset "FullOrthoLanczos.jl" begin
+    include("FullOrthoLanczos.jl")
+end
+
+@testset "GeneralFTLM.jl" begin
+    include("GeneralFTLM.jl")
+end
\ No newline at end of file

test/runtests.jl

-using Test
-using FiniteTLanczos
-using Random; Random.seed!()
-
-include("FullOrthoLanczos.jl")
-include("GeneralFTLM.jl")
+using Test, FiniteTLanczos
 
+@testset "cpu_test.jl" begin
+    include("cpu_test.jl")
+end
 
+@testset "gpu_test.jl" begin
+    include("gpu_test.jl")
+end