diff --git a/Hamiltonian.jl b/Hamiltonian.jl index bfaed69..321c343 100644 --- a/Hamiltonian.jl +++ b/Hamiltonian.jl @@ -5,39 +5,36 @@ using TensorOperations, KrylovKit, LinearAlgebra, CUDA, CUDA.CUTENSOR "A Hamiltonian that can be applied to a vector" struct Hamiltonian{T} - d::Int - n::Int - N::Int - L::T - μ::T - ∂1 # Matrix{Complex{T}} or Nothing - K_diag # CuTensor{Complex{T}} or Nothing - K_mixed # CuTensor{Complex{T}} or Nothing - Vs # Array{Complex{T}} or CuArray{Complex{T}} + s::system{T} + K_partial::Matrix{Complex{T}} + K_diag::Union{CuTensor{Complex{T}},Nothing} + K_mixed::Union{CuTensor{Complex{T}},Nothing} + Vs::Union{Array{Complex{T}},CuArray{Complex{T}}} hermitian::Bool mode::Hamiltonian_backend - function Hamiltonian{T}(V_twobody::Function, d::Int, n::Int, N::Int, L::T, ϕ::T, μ::T, n_image::Int, mode::Hamiltonian_backend) where {T<:Float} + + 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" - k = -N÷2:N÷2-1 - Vs = calculate_Vs(V_twobody, d, n, N, L, ϕ, n_image) + k = -s.N÷2:s.N÷2-1 + Vs = calculate_Vs(s, V_twobody, convert(T, ϕ), n_image) hermitian = ϕ == 0.0 - if mode == cpu_tensor - ∂1 = exp(-im * ϕ) .* ∂_1DOF.(L, N, k, k') - return new{T}(d, n, N, L, μ, ∂1, nothing, nothing, Vs, hermitian, mode) - elseif mode == gpu_cutensor - K_partial = (exp(-im * ϕ) * im / sqrt(2 * μ)) .* ∂_1DOF.(L, N, k, k') + K_partial = (exp(-im * convert(T, ϕ)) * im / sqrt(2 * s.μ)) .* ∂_1DOF.(Ref(s), k, k') + K_diag = nothing + K_mixed = nothing + if mode == gpu_cutensor K_diag = CuTensor(CuArray(K_partial * K_partial), ['a', 'A']) K_mixed = CuTensor(CuArray(K_partial), ['a', 'A']) * CuTensor(CuArray(K_partial), ['b', 'B']) - return new{T}(d, n, N, L, μ, nothing, K_diag, K_mixed, CuArray(Vs), hermitian, mode) + Vs = CuArray(Vs) end + return new{T}(s, K_partial, K_diag, K_mixed, Vs, hermitian, mode) end end -Base.size(H::Hamiltonian, i::Int)::Int = (i == 1 || i == 2) ? H.N^(H.d * (H.n - 1)) : throw(ArgumentError("Hamiltonian only has 2 dimesions")) +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.N, H.d * (H.n - 1))...) +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} @@ -45,14 +42,13 @@ function LinearAlgebra.mul!(out::Array{Complex{T}}, H::Hamiltonian{T}, v::Array{ # apply V operator @. out = H.Vs * v # apply K opereator - coeff = -1 / (2 * H.μ) - coords = H.n - 1 - nconList_v_template = -collect(1:H.d*(coords)) - for dim = 1:H.d + coords = H.s.n - 1 + nconList_v_template = -collect(1:H.s.d*(coords)) + for dim = 1:H.s.d for coord1 = 1:coords for coord2 = 1:coord1 - i1 = which_index(H.n, dim, coord1) - i2 = which_index(H.n, dim, coord2) + i1 = which_index(H.s, dim, coord1) + i2 = which_index(H.s, dim, coord2) nconList_1 = [-i1, 1] nconList_2 = [-i2, 2] nconList_v = copy(nconList_v_template) @@ -62,8 +58,8 @@ function LinearAlgebra.mul!(out::Array{Complex{T}}, H::Hamiltonian{T}, v::Array{ nconList_v[i1] = 1 end nconList_v[i2] = 2 - v_new = @ncon((H.∂1, H.∂1, v), (nconList_1, nconList_2, nconList_v)) - out = axpy!(coeff, v_new, out) + v_new = @ncon((H.K_partial, H.K_partial, v), (nconList_1, nconList_2, nconList_v)) + out = axpy!(1, v_new, out) end end end @@ -85,15 +81,15 @@ function LinearAlgebra.mul!(out::CuArray{Complex{T}}, H::Hamiltonian{T}, v::CuAr NVTX.@range "V" @. out = H.Vs * v synchronize(ctx) # apply K opereator - coords = H.n - 1 - inds_template = ('a' - 1) .+ collect(1:H.d*(coords)) + 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.d + for dim = 1:H.s.d for coord1 = 1:coords for coord2 = 1:coord1 - i1 = which_index(H.n, dim, coord1) - i2 = which_index(H.n, dim, coord2) + i1 = which_index(H.s, dim, coord1) + i2 = which_index(H.s, dim, coord2) @assert v_t.inds == inds_template "v indices permuted" if i1 == i2 @assert H.K_diag.inds[2] == 'A' "K_diag indices permuted" diff --git a/benchmark.jl b/benchmark.jl index 12324de..fffef78 100644 --- a/benchmark.jl +++ b/benchmark.jl @@ -27,11 +27,11 @@ end N=10 n_image=1 -μ=0.5 -for L::T in 5.0:14.0 +for L in 5.0:14.0 println("Constructing H operator...") - @time H=Hamiltonian{T}(V_test,3,3,N,L,convert(T,0),convert(T,μ),n_image,mode) + s=system{T}(3,3,N,L) + @time H=Hamiltonian{T}(s,V_test,0,n_image,mode) println("Applying H 1000 times...") if GPU_mode v=CUDA.rand(Complex{T},vectorDims(H)...) diff --git a/common.jl b/common.jl index a7ab441..aae050d 100644 --- a/common.jl +++ b/common.jl @@ -1,53 +1,64 @@ Float = Union{Float32,Float64} +"A few-body system defined by its physical parameters" +struct system{T} + d::Int + n::Int + N::Int + L::T + μ::T + + system{T}(d::Int, n::Int, N::Int, L::Real, μ::Real=0.5) where {T<:Float} = new{T}(d, n, N, convert(T, L), convert(T, μ)) +end + norm_square(x::Array{Int})::Int = sum(x .* x) "Eq (46): Partial derivative matrix element for 1 degree of freedom" -function ∂_1DOF(L::T, N::Int, k::Int, l::Int)::Complex{T} where {T<:Float} +function ∂_1DOF(s::system{T}, k::Int, l::Int)::Complex{T} where {T<:Float} if k == l - return -im * (π / L) + return -im * (π / s.L) else - return (π / L) * (-1)^(k - l) * exp(-im * π * (k - l) / N) / sin(π * (k - l) / N) + return (π / s.L) * (-1)^(k - l) * exp(-im * π * (k - l) / s.N) / sin(π * (k - l) / s.N) end end "Which index (dimension of the multidimensional array) corresponds to spatial dimension 'dim' and particle 'p'?" -which_index(n::Int, dim::Int, p::Int)::Int = (dim - 1) * (n - 1) + p +which_index(s::system, dim::Int, p::Int)::Int = (dim - 1) * (s.n - 1) + p "Δk (distance in terms of lattice paramter) between two particles along the given dimension" -function get_Δk(n::Int, N::Int, i::CartesianIndex, dim::Int, p1::Int, p2::Int)::Int +function get_Δk(s::system, i::CartesianIndex, dim::Int, p1::Int, p2::Int)::Int if p1 == p2 return 0 - elseif p1 == n - return -(i[which_index(n, dim, p2)] - N ÷ 2 - 1) - elseif p2 == n - return i[which_index(n, dim, p1)] - N ÷ 2 - 1 + elseif p1 == s.n + return -(i[which_index(s, dim, p2)] - s.N ÷ 2 - 1) + elseif p2 == s.n + return i[which_index(s, dim, p1)] - s.N ÷ 2 - 1 else - return i[which_index(n, dim, p1)] - i[which_index(n, dim, p2)] + return i[which_index(s, dim, p1)] - i[which_index(s, dim, p2)] end end "Calculate diagonal elements of the V matrix" -function calculate_Vs(V_twobody::Function, d::Int, n::Int, N::Int, L::T, ϕ::T, n_image::Int)::Array{Complex{T}} where {T<:Float} - coeff² = (exp(im * ϕ) * L / N)^2 - images = collect.(Iterators.product(fill(-n_image:n_image, d)...)) # TODO: Learn how to use tuples instead of vectors - Vs = zeros(Complex{T}, fill(N, d * (n - 1))...) +function calculate_Vs(s::system{T}, V_twobody::Function, ϕ::T, n_image::Int)::Array{Complex{T}} where {T<:Float} + coeff² = (exp(im * ϕ) * s.L / s.N)^2 + images = collect.(Iterators.product(fill(-n_image:n_image, s.d)...)) # TODO: Learn how to use tuples instead of vectors + Vs = zeros(Complex{T}, fill(s.N, s.d * (s.n - 1))...) Threads.@threads for i in CartesianIndices(Vs) - for p1 in 1:n - for p2 in (p1 + 1):n - min_Δk = Array{Int}(undef, d) - for dim in 1:d - Δk = get_Δk(n, N, i, dim, p1, p2) - if Δk > N ÷ 2 - min_Δk[dim] = Δk - N - elseif Δk < -N ÷ 2 - min_Δk[dim] = Δk + N + for p1 in 1:s.n + for p2 in (p1 + 1):s.n + min_Δk = Array{Int}(undef, s.d) + for dim in 1:s.d + Δk = get_Δk(s, i, dim, p1, p2) + if Δk > s.N ÷ 2 + min_Δk[dim] = Δk - s.N + elseif Δk < -s.N ÷ 2 + min_Δk[dim] = Δk + s.N else min_Δk[dim] = Δk end end for image in images - Δk² = norm_square(min_Δk .- (N .* image)) + Δk² = norm_square(min_Δk .- (s.N .* image)) Vs[i] += V_twobody(Δk² * coeff²) end end diff --git a/example.ipynb b/example.ipynb index b92409e..d46f387 100644 --- a/example.ipynb +++ b/example.ipynb @@ -24,12 +24,12 @@ "d = 3\n", "n = 3\n", "N = 6\n", - "L::T = 12\n", - "ϕ::T = 0.0\n", - "μ::T = 0.5\n", + "L = 12\n", + "ϕ = 0.0\n", "n_imag = 1\n", "\n", - "H = Hamiltonian{T}(V_gauss, d, n, N, L, ϕ, μ, n_imag, mode)\n", + "s = system{T}(d, n, N, L)\n", + "H = Hamiltonian{T}(s, V_gauss, ϕ, n_imag, mode)\n", "@time evals, evecs, info = eig(H, 5)\n", "print(info.numops, \" operations : \")\n", "println(evals)" @@ -49,12 +49,12 @@ "d = 3\n", "n = 2\n", "N = 32\n", - "L::T = 16\n", - "ϕ::T = 0.5\n", - "μ::T = 0.5\n", + "L = 16\n", + "ϕ = 0.5\n", "n_imag = 0\n", "\n", - "H = Hamiltonian{T}(V_gauss, d, n, N, L, ϕ, μ, n_imag, mode)\n", + "s = system{T}(d, n, N, L)\n", + "H = Hamiltonian{T}(s, V_gauss, ϕ, n_imag, mode)\n", "@time evals, evecs, info = eig(H, 20)\n", "print(info.numops, \" operations : \")\n", "print(evals)\n", diff --git a/testing.ipynb b/testing.ipynb new file mode 100644 index 0000000..11a2c5c --- /dev/null +++ b/testing.ipynb @@ -0,0 +1,71 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "include(\"Hamiltonian.jl\")\n", + "\n", + "println(\"Running with \",Threads.nthreads(),\" thread(s)\")\n", + "println(\"Available GPUs:\")\n", + "println.(name.(devices()))\n", + "\n", + "T=Float32\n", + "\n", + "function V_test(r2)\n", + " return -4*exp(-r2/4)\n", + "end\n", + "\n", + "function test(mode)\n", + " for (n,N) in [(2,16),(3,8)]\n", + " println(\"\\n$n-body system with N=$N\")\n", + " n_image=0\n", + " for L::T in 5.0:9.0\n", + " print(\"L=$L\")\n", + " s=system{T}(3,n,N,L)\n", + " H=Hamiltonian{T}(s,V_test,0.0,n_image,mode)\n", + " evals,_,_ = eig(H,5)\n", + " println(real.(evals))\n", + " end\n", + " end\n", + "end" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "test(cpu_tensor)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "test(gpu_cutensor)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Julia 1.8.5", + "language": "julia", + "name": "julia-1.8" + }, + "language_info": { + "file_extension": ".jl", + "mimetype": "application/julia", + "name": "julia", + "version": "1.8.5" + }, + "orig_nbformat": 4 + }, + "nbformat": 4, + "nbformat_minor": 2 +}