84 lines
2.6 KiB
Julia
84 lines
2.6 KiB
Julia
Float = Union{Float32,Float64}
|
|
|
|
norm_square(x) = sum(x .* x)
|
|
reducedMass(m1, m2) = 1 / (1/m1 + 1/m2)
|
|
|
|
"A few-body system defined by its physical parameters"
|
|
struct system{T}
|
|
d::Int
|
|
n::Int
|
|
N::Int
|
|
L::T
|
|
μs::Vector{Int}
|
|
invU::Matrix{Int}
|
|
|
|
function system{T}(d::Int, n::Int, N::Int, L::Real) where {T<:Float}
|
|
μs = [Int((coord + 1)^2 * reducedMass(coord, 1)) for coord in 1:(n - 1)]
|
|
|
|
# TODO: Optimize
|
|
invU = Matrix{Int}(undef, n, n - 1)
|
|
for i in CartesianIndices(invU)
|
|
if i[1] - 1 == i[2]
|
|
invU[i] = -i[2]
|
|
elseif i[1] > i[2]
|
|
invU[i] = 0
|
|
else
|
|
invU[i] = 1
|
|
end
|
|
end
|
|
|
|
return new{T}(d, n, N, convert(T, L), μs, invU)
|
|
end
|
|
end
|
|
|
|
"Eq (46): Partial derivative matrix element for 1 degree of freedom"
|
|
function ∂_1DOF(s::system{T}, k::Int, l::Int)::Complex{T} where {T<:Float}
|
|
if k == l
|
|
return -im * (π / s.L)
|
|
else
|
|
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' of coordinate 'coord'?"
|
|
which_index(s::system, dim::Int, coord::Int)::Int = (dim - 1) * (s.n - 1) + coord
|
|
|
|
"Get the distance to the nearest image of the particle"
|
|
function nearest(s::system, Δk::Int)::Int
|
|
# TODO: Optimize
|
|
while true
|
|
if Δk > s.N ÷ 2
|
|
Δk -= s.N
|
|
elseif Δk < -s.N ÷ 2
|
|
Δk += s.N
|
|
else
|
|
return Δk
|
|
end
|
|
end
|
|
end
|
|
|
|
"Calculate diagonal elements of the V matrix"
|
|
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)
|
|
xs = reshape(collect(Tuple(i)), s.n - 1, s.d) .- (s.N ÷ 2 - 1)
|
|
rs = s.invU * xs
|
|
for p1 in 1:s.n
|
|
for p2 in 1:(p1 - 1)
|
|
Δk = Array{Int}(undef, s.d)
|
|
for dim in 1:s.d
|
|
Δk_temp = Int(rs[p1, dim] - rs[p2, dim])
|
|
Δk[dim] = nearest(s, Δk_temp)
|
|
end
|
|
for image in images
|
|
Δk² = norm_square(Δk .- (s.N .* image))
|
|
Vs[i] += V_twobody(Δk² * coeff²)
|
|
end
|
|
end
|
|
end
|
|
end
|
|
return Vs
|
|
end
|