119 lines
4.7 KiB
Julia
119 lines
4.7 KiB
Julia
include("common.jl")
|
|
using TensorOperations, KrylovKit, LinearAlgebra, CUDA, cuTENSOR, NVTX
|
|
|
|
@enum Hamiltonian_backend cpu_tensor gpu_cutensor
|
|
|
|
"A Hamiltonian that can be applied to a vector"
|
|
struct Hamiltonian{T}
|
|
s::system{T}
|
|
K::Union{CuTensor{Complex{T}}, Matrix{Complex{T}}}
|
|
Vs::Union{Array{Complex{T}}, CuArray{Complex{T}}}
|
|
hermitian::Bool
|
|
mode::Hamiltonian_backend
|
|
|
|
function Hamiltonian{T}(s::system{T}, V_twobody::Function, ϕ::Real, n_image::Int, mode::Hamiltonian_backend) where {T<:Float}
|
|
@assert mode != gpu_cutensor || CUDA.functional() && CUDA.has_cuda() && CUDA.has_cuda_gpu() "CUDA not available"
|
|
hermitian = ϕ == 0.0
|
|
Vs = calculate_Vs(s, V_twobody, convert(T, ϕ), n_image)
|
|
k = -s.N÷2:s.N÷2-1
|
|
∂ = ∂_1DOF.(Ref(s), k, k')
|
|
K = exp(-2im * convert(T, ϕ)) .* (∂ * ∂) # TODO: Calculate K matrix elements directly
|
|
if mode == gpu_cutensor
|
|
K = CuTensor(CuArray(K), ['a', 'A'])
|
|
Vs = CuArray(Vs)
|
|
end
|
|
return new{T}(s, K, Vs, hermitian, mode)
|
|
end
|
|
end
|
|
|
|
Base.size(H::Hamiltonian, i::Int)::Int = (i == 1 || i == 2) ? H.s.N^(H.s.d * (H.s.n - 1)) : throw(ArgumentError("Hamiltonian only has 2 dimesions"))
|
|
Base.size(H::Hamiltonian)::Dims{2} = (size(H, 1), size(H, 2))
|
|
|
|
"Dimensions of a vector to which 'H' can be applied"
|
|
vectorDims(H::Hamiltonian)::Dims = tuple(fill(H.s.N, H.s.d * (H.s.n - 1))...)
|
|
|
|
"Apply 'H' on 'v' and store the result in 'out' using the 'cpu_tensor' backend"
|
|
function LinearAlgebra.mul!(out::Array{Complex{T}}, H::Hamiltonian{T}, v::Array{Complex{T}})::Array{Complex{T}} where {T<:Float}
|
|
#LinearMaps.check_dim_mul(out,H,v) --- dimensions don't match
|
|
# apply V operator
|
|
@. out = H.Vs * v
|
|
# apply K opereator
|
|
coords = H.s.n - 1
|
|
nconList_v_template = -collect(1:H.s.d*(coords))
|
|
for dim = 1:H.s.d
|
|
for coord = 1:coords
|
|
i = which_index(H.s, dim, coord)
|
|
nconList_K = [-i, 1]
|
|
nconList_v = copy(nconList_v_template)
|
|
nconList_v[i] = 1
|
|
v_new = @ncon((H.K, v), (nconList_K, nconList_v))
|
|
coeff = -1 / (2 * H.s.μs[coord])
|
|
out = axpy!(coeff, v_new, out)
|
|
end
|
|
end
|
|
return out
|
|
end
|
|
|
|
"cuTENSOR contraction and accumulation (C = A * B + C)"
|
|
function contract_accumulate!(alpha::Number, C::CuTensor, A::CuTensor, B::CuTensor)::CuTensor
|
|
cuTENSOR.contraction!(alpha, A.data, A.inds, cuTENSOR.CUTENSOR_OP_IDENTITY, B.data, B.inds, cuTENSOR.CUTENSOR_OP_IDENTITY,
|
|
one(eltype(C)), C.data, C.inds, cuTENSOR.CUTENSOR_OP_IDENTITY, cuTENSOR.CUTENSOR_OP_IDENTITY)
|
|
return C
|
|
end
|
|
|
|
"Apply 'H' on 'v' and store the result in 'out' using the 'gpu_cutensor' backend"
|
|
function LinearAlgebra.mul!(out::CuArray{Complex{T}}, H::Hamiltonian{T}, v::CuArray{Complex{T}})::CuArray{Complex{T}} where {T<:Float}
|
|
#LinearMaps.check_dim_mul(out,H,v) --- dimensions don't match
|
|
ctx = context()
|
|
# apply V operator
|
|
NVTX.@range "V" @. out = H.Vs * v
|
|
synchronize(ctx)
|
|
# apply K opereator
|
|
coords = H.s.n - 1
|
|
inds_template = ('a' - 1) .+ collect(1:H.s.d*(coords))
|
|
v_t = CuTensor(v, copy(inds_template))
|
|
out_t = CuTensor(out, copy(inds_template))
|
|
for dim = 1:H.s.d
|
|
for coord = 1:coords
|
|
i = which_index(H.s, dim, coord)
|
|
@assert v_t.inds == inds_template "v indices permuted"
|
|
@assert H.K.inds[2] == 'A' "K_diag indices permuted"
|
|
H.K.inds[1] = 'a' - 1 + i
|
|
v_t.inds[i] = 'A'
|
|
#synchronize(ctx)
|
|
coeff = -1 / (2 * H.s.μs[coord])
|
|
NVTX.@range "K" out_t = contract_accumulate!(coeff, out_t, H.K, v_t)
|
|
v_t.inds[i] = 'a' - 1 + i
|
|
end
|
|
end
|
|
@assert out_t.inds == inds_template "out indices permuted"
|
|
synchronize(ctx)
|
|
return out_t.data
|
|
end
|
|
|
|
"Apply 'H' on 'v' and return the result"
|
|
function (H::Hamiltonian)(v)
|
|
out = similar(v)
|
|
return mul!(out, H, v)
|
|
end
|
|
|
|
tolerance = 1e-6
|
|
|
|
"Wrapper for KrylovKit.eigsolve"
|
|
function eig(H::Hamiltonian{T}, levels::Int; resonances = !H.hermitian)::Tuple{Vector,Vector,KrylovKit.ConvergenceInfo} where {T<:Float}
|
|
if H.mode == cpu_tensor
|
|
x₀ = rand(Complex{T}, vectorDims(H)...)
|
|
elseif H.mode == gpu_cutensor
|
|
x₀ = CUDA.rand(Complex{T}, vectorDims(H)...)
|
|
synchronize()
|
|
end
|
|
evals, evecs, info = eigsolve(H, x₀, levels, resonances ? :LI : :SR; ishermitian = H.hermitian, tol = tolerance, krylovdim = levels * 8)
|
|
info.converged < levels && throw(error("Not enough convergence"))
|
|
if H.hermitian evals = real.(evals) end
|
|
if H.mode == gpu_cutensor # to avoid possible GPU memory leak
|
|
CUDA.reclaim()
|
|
GC.gc(true)
|
|
end
|
|
return evals, evecs, info
|
|
end
|