DVR-jl/common.jl

96 lines
3.2 KiB
Julia

include("irrep.jl")
Float = Union{Float32,Float64}
@enum rep all A1
"A few-body system defined by its physical parameters"
struct system{T}
d::Int
n::Int
N::Int
L::T
μ::T
sym::rep
unique_i::Array{Int}
unique_point::Array{Int}
multiplicity::Array{Int}
function system{T}(d::Int, n::Int, N::Int, L::Real, μ::Real=0.5, sym::rep=all) where {T<:Float}
@assert d == 3 "Only supports 3D"
if sym == all
unique_i, unique_point, multiplicity = calculate_all_data(N)
elseif sym == A1
unique_i, unique_point, multiplicity = calculate_A1_data(N)
else
throw(ArgumentError("Symmetry not yet implemented"))
end
return new{T}(d, n, N, convert(T, L), convert(T, μ), sym, unique_i, unique_point, multiplicity)
end
end
norm_square(x::Array{Int})::Int = sum(x .* x)
"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' and particle 'p'?"
which_index(s::system, dim::Int, p::Int)::Int = (dim - 1) * (s.n - 2) + p + 1
"Δk (distance in terms of lattice paramter) between two particles along the given dimension"
function get_k(s::system, i::CartesianIndex, dim::Int, p::Int)::Int
if p == 1
s.unique_point[i[1], dim]
else
return i[which_index(s, dim, p)] - s.N ÷ 2 - 1
end
end
"Δk (distance in terms of lattice paramter) between two particles along the given dimension"
function get_Δk(s::system, i::CartesianIndex, dim::Int, p1::Int, p2::Int)::Int
if p1 == p2
return 0
elseif p1 == s.n
return -get_k(s, i, dim, p2)
elseif p2 == s.n
return get_k(s, i, dim, p1)
else
return get_k(s, i, dim, p1) - get_k(s, i, dim, p2)
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}, length(s.unique_i), fill(s.N, s.d * (s.n - 2))...)
Threads.@threads for i in CartesianIndices(Vs)
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 .- (s.N .* image))
Vs[i] += V_twobody(Δk² * coeff²)
end
end
end
end
return Vs
end