First working implementation

This commit is contained in:
ysyapa 2023-03-31 07:38:38 -04:00
parent 526ef37578
commit 153fb770c3
4 changed files with 72 additions and 27 deletions

19
CPU.jl
View File

@ -9,12 +9,13 @@ struct HOperator{T}
L::T
μ::T
∂1::Matrix{Complex{T}}
Vs::Array{T}
function HOperator{T}(V_twobody::Function, d::Int, n::Int, N::Int, L::T, μ::T, n_image::Int) where {T<:Float}
Vs::Array{Complex{T}}
hermitian::Bool
function HOperator{T}(V_twobody::Function, d::Int, n::Int, N::Int, L::T, ϕ::T, μ::T, n_image::Int) where {T<:Float}
k = -N÷2:N÷2-1
∂1 = ∂_1DOF.(L, N, k, k')
Vs = calculate_Vs(V_twobody, d, n, N, L, n_image)
return new{T}(d, n, N, L, μ, ∂1, Vs)
∂1 = exp(-im * ϕ) .* ∂_1DOF.(L, N, k, k')
Vs = calculate_Vs(V_twobody, d, n, N, L, ϕ, n_image)
return new{T}(d, n, N, L, μ, ∂1, Vs, ϕ == 0.0)
end
end
@ -64,9 +65,9 @@ end
tolerance = 1e-6
"Wrapper for KrylovKit.eigsolve"
function eig(H::HOperator{T}, levels::Int)::Tuple{Vector{T},Any,Any} where {T<:Float}
function eig(H::HOperator{T}, levels::Int; resonances = !H.hermitian)::Tuple{Vector{Complex{T}},Any,Any} where {T<:Float}
x₀ = rand(Complex{T}, vectorDims(H))
evals, evecs, info = eigsolve(H, x₀, levels, :SR; ishermitian = true, tol = tolerance)
info.converged < levels && throw(error("Not enough convergence"))
return real.(evals), evecs, info
evals, evecs, info = eigsolve(H, x₀, levels, resonances ? :LI : :SR; ishermitian = H.hermitian, tol = tolerance)
resonances || info.converged < levels && throw(error("Not enough convergence")) # don't check convergence for resonances
return evals, evecs, info
end

19
GPU.jl
View File

@ -10,14 +10,15 @@ struct HOperator{T}
N::Int
K_diag::CuTensor{Complex{T}}
K_mixed::CuTensor{Complex{T}}
Vs::CuArray{T}
function HOperator{T}(V_twobody::Function, d::Int, n::Int, N::Int, L::T, μ::T, n_image::Int) where {T<:Float}
Vs::CuArray{Complex{T}}
hermitian::Bool
function HOperator{T}(V_twobody::Function, d::Int, n::Int, N::Int, L::T, ϕ::T, μ::T, n_image::Int) where {T<:Float}
k = -N÷2:N÷2-1
K_partial = (im / sqrt(2 * μ)) .* ∂_1DOF.(L, N, k, k')
K_partial = (exp(-im * ϕ) * im / sqrt(2 * μ)) .* ∂_1DOF.(L, N, k, k')
K_diag = CuTensor(CuArray(K_partial * K_partial), ['a', 'A'])
K_mixed = CuTensor(CuArray(K_partial), ['a', 'A']) * CuTensor(CuArray(K_partial), ['b', 'B'])
Vs = calculate_Vs(V_twobody, d, n, N, L, n_image)
return new{T}(d, n, N, K_diag, K_mixed, Vs)
Vs = calculate_Vs(V_twobody, d, n, N, L, ϕ, n_image)
return new{T}(d, n, N, K_diag, K_mixed, Vs, ϕ == 0.0)
end
end
@ -88,10 +89,10 @@ end
tolerance = 1e-6
"Wrapper for KrylovKit.eigsolve"
function eig(H::HOperator{T}, levels::Int)::Tuple{Vector{T},Any,Any} where {T<:Float}
function eig(H::HOperator{T}, levels::Int; resonances = !H.hermitian)::Tuple{Vector{Complex{T}},Any,Any} where {T<:Float}
x₀ = CUDA.rand(Complex{T}, vectorDims(H)...) # ... added
synchronize()
evals, evecs, info = eigsolve(H, x₀, levels, :SR; ishermitian = true, tol = tolerance)
info.converged < levels && throw(error("Not enough convergence"))
return real.(evals), evecs, info
evals, evecs, info = eigsolve(H, x₀, levels, resonances ? :LI : :SR; ishermitian = H.hermitian, tol = tolerance)
resonances || info.converged < levels && throw(error("Not enough convergence")) # don't check convergence for resonances
return evals, evecs, info
end

View File

@ -28,10 +28,10 @@ function get_Δk(n::Int, N::Int, i::CartesianIndex, dim::Int, p1::Int, p2::Int):
end
"Calculate diagonal elements of the V matrix"
function calculate_Vs(V_twobody::Function, d::Int, n::Int, N::Int, L::T, n_image::Int)::Array{T} where {T<:Float}
L²_over_N² = (L / N)^2
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(T, fill(N, d * (n - 1))...)
Vs = zeros(Complex{T}, fill(N, d * (n - 1))...)
Threads.@threads for i in CartesianIndices(Vs)
for p1 in 1:n
for p2 in (p1 + 1):n
@ -48,7 +48,7 @@ function calculate_Vs(V_twobody::Function, d::Int, n::Int, N::Int, L::T, n_image
end
for image in images
Δk² = norm_square(min_Δk .- (N .* image))
Vs[i] += V_twobody(Δk² * L²_over_N²)
Vs[i] += V_twobody(Δk² * coeff²)
end
end
end

View File

@ -7,25 +7,68 @@
"outputs": [],
"source": [
"# prerequisite packages: KrylovKit, TensorOperations, LinearAlgebra, CUDA#tb/cutensor\n",
"\n",
"include(\"CPU.jl\") # using CPU mode\n",
"T = Float32\n",
"\n",
"V_gauss(r2::T)::T =\n",
"T = Float32 # single-precision mode"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"V_gauss(r2) =\n",
" -4 * exp(-r2 / 4)\n",
"\n",
"d = 3\n",
"n = 3\n",
"N = 6\n",
"L::T = 12\n",
"ϕ::T = 0.0\n",
"mu::T = 0.5\n",
"n_imag = 1\n",
"\n",
"H = HOperator{T}(V_gauss, d, n, N, L, mu, n_imag)\n",
"H = HOperator{T}(V_gauss, d, n, N, L, ϕ, mu, n_imag)\n",
"@time evals, evecs, info = eig(H, 5)\n",
"print(info.numops, \" operations : \")\n",
"println(evals)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"using Plots\n",
"\n",
"V_gauss(r2) =\n",
" -4 * exp(-r2 / 4)\n",
"\n",
"d = 3\n",
"n = 2\n",
"N = 32\n",
"L::T = 16\n",
"ϕ::T = 0.5\n",
"mu::T = 0.5\n",
"n_imag = 0\n",
"\n",
"H = HOperator{T}(V_gauss, d, n, N, L, ϕ, mu, n_imag)\n",
"@time evals, evecs, info = eig(H, 20)\n",
"print(info.numops, \" operations : \")\n",
"print(evals)\n",
"\n",
"scatter(real.(evals), imag.(evals); legend=false)\n",
"xlabel!(\"Re E\")\n",
"ylabel!(\"Im E\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {