239 lines
9.4 KiB
Julia
239 lines
9.4 KiB
Julia
using SparseArrays
|
||
using QuadGK
|
||
using LRUCache
|
||
include("helper.jl")
|
||
include("math.jl")
|
||
|
||
"1-body HO basis"
|
||
struct ho_basis_1B
|
||
dim::Int # dimensionality of the basis
|
||
Es::Vector{Int}
|
||
ns::Vector{Int}
|
||
ls::Vector{Int}
|
||
|
||
function ho_basis_1B(E_max)
|
||
Es = Int[]
|
||
ns = Int[]
|
||
ls = Int[]
|
||
|
||
# E = 2*n + l
|
||
for E in 0 : E_max
|
||
for n in 0 : E ÷ 2
|
||
l = E - 2*n
|
||
push!(Es, E)
|
||
push!(ns, n)
|
||
push!(ls, l)
|
||
end
|
||
end
|
||
|
||
return new(length(Es), Es, ns, ls)
|
||
end
|
||
end
|
||
|
||
"2-body HO basis"
|
||
struct ho_basis_2B
|
||
E_max::Int
|
||
Λ::Int
|
||
dim::Int # dimensionality of the basis
|
||
Es::Vector{Int}
|
||
n1s::Vector{Int}
|
||
l1s::Vector{Int}
|
||
n2s::Vector{Int}
|
||
l2s::Vector{Int}
|
||
|
||
function ho_basis_2B(E_max, Λ=-1)
|
||
Es = Int[]
|
||
n1s = Int[]
|
||
l1s = Int[]
|
||
n2s = Int[]
|
||
l2s = Int[]
|
||
|
||
# E = 2*n1 + l1 + 2*n2 + l2
|
||
for E in E_max : -2 : 0 # same parity states only
|
||
for n1 in 0 : E ÷ 2
|
||
for n2 in 0 : (E - 2*n1) ÷ 2
|
||
for l1 in 0 : (E - 2*n1 - 2*n2)
|
||
l2 = E - 2*n1 - 2*n2 - l1
|
||
if Λ≥0 && !triangle_ineq(l1, l2, Λ); continue; end
|
||
push!(Es, E)
|
||
push!(n1s, n1)
|
||
push!(l1s, l1)
|
||
push!(n2s, n2)
|
||
push!(l2s, l2)
|
||
end
|
||
end
|
||
end
|
||
end
|
||
|
||
return new(E_max, Λ, length(Es), Es, n1s, l1s, n2s, l2s)
|
||
end
|
||
end
|
||
|
||
"Number of possible distinct matrix elements for a given l."
|
||
function cache_size_l(E_max::Int, l::Int)::Int
|
||
n_max = (E_max - l) ÷ 2
|
||
return (n_max * (n_max + 1)) ÷ 2
|
||
end
|
||
|
||
"Number of possible distinct matrix elements for all l."
|
||
cache_size(E_max::Int)::Int = sum(l -> cache_size_l(E_max, l), 0 : E_max)
|
||
|
||
"Preallocation of cache for PE matrix elements."
|
||
prealloc_V_cache(E_max::Int, dtype::DataType=Float64) = LRU{Tuple{UInt8, UInt8, UInt8}, dtype}(maxsize=cache_size(E_max))
|
||
|
||
function V_numerical(V_of_r, l, n1, n2; μω_gen=1.0, atol=0, maxevals=10^7)
|
||
const_part = sqrt(μω_gen) * ho_basis_const(l, n1) * ho_basis_const(l, n2)
|
||
integrand(r) = ho_basis_func(l, n1, sqrt(μω_gen) * r) * ho_basis_func(l, n2, sqrt(μω_gen) * r) * V_of_r(r)
|
||
(integral, _) = quadgk(integrand, 0, Inf; atol=atol, maxevals=maxevals)
|
||
return const_part * integral
|
||
end
|
||
|
||
"KE matrix for a single DOF. Set kron_deltas=[other quantum numbers] for other DOFs which the operator does not act on.
|
||
E.g. get_sp_T_matrix(n1s, l1s, kron_deltas=[n2s l2s])"
|
||
function get_sp_T_matrix(ns, ls, kron_deltas=[]; μω_gen=1.0, μ=1.0)
|
||
mat = spzeros(length(ns), length(ns))
|
||
for idx in CartesianIndices(mat)
|
||
(i, j) = Tuple(idx)
|
||
all(arr -> arr[i]==arr[j], kron_deltas) || continue # check if all Kronecker deltas are non-zero
|
||
if ls[i] == ls[j]
|
||
if ns[i] == ns[j]
|
||
mat[idx] = ns[j] + ls[i]/2 + 3/4
|
||
elseif abs(ns[i]-ns[j]) == 1
|
||
n_max = max(ns[i], ns[j])
|
||
mat[idx] = -(1/2) * sqrt(n_max * (n_max + ls[i] + 1/2))
|
||
end
|
||
end
|
||
end
|
||
return (μω_gen / μ) .* mat
|
||
end
|
||
|
||
"PE matrix for a single DOF. Set kron_deltas=[other quantum numbers] for other DOFs which the operator does not act on.
|
||
E.g. get_sp_V_matrix(n1s, l1s, kron_deltas=[n2s l2s])
|
||
Providing a preallocated cache is optional. Otherwise, provided E_max value is used to initialize one (defaults to 100)."
|
||
function get_sp_V_matrix(V_l, ns, ls, kron_deltas=[]; dtype=Float64, E_max=100, cache=prealloc_V_cache(E_max, dtype))
|
||
mat = zeros(dtype, length(ns), length(ns))
|
||
Threads.@threads for idx in CartesianIndices(mat)
|
||
(i, j) = Tuple(idx)
|
||
all(arr -> arr[i]==arr[j], kron_deltas) || continue # check if all Kronecker deltas are non-zero
|
||
if ls[i] == ls[j]
|
||
l = UInt8(ls[i])
|
||
n1, n2 = UInt8.(minmax(ns[i], ns[j])) # assuming transpose symmetry
|
||
mat[idx] = (get!(cache, (l, n1, n2)) do; V_l(l, n1, n2); end)
|
||
end
|
||
end
|
||
return sparse(mat)
|
||
end
|
||
|
||
function Moshinsky_transform(basis::ho_basis_2B)
|
||
NQMAX = maximum(basis.Es)
|
||
@assert all(mod.(basis.Es, 2) .== mod(NQMAX, 2)) "Can only admit basis states with same parity"
|
||
|
||
LMIN = basis.Λ
|
||
LMAX = basis.Λ
|
||
CO = 1/sqrt(2)
|
||
SI = 1/sqrt(2)
|
||
|
||
# dimensions BRAC(0:LMAX,0:(NQMAX-LMIN)/2,0:(NQMAX-LMIN)/2,0:(NQMAX-LMIN)/2,0:(NQMAX-LMIN)/2,0:LMAX,0:(NQMAX-LMIN)/2,LMIN:LMAX)
|
||
BRAC = zeros(Float64, 1 + LMAX, 1 + (NQMAX - LMIN) ÷ 2, 1 + (NQMAX - LMIN) ÷ 2, 1 + (NQMAX - LMIN) ÷ 2, 1 + (NQMAX - LMIN) ÷ 2, 1 + LMAX, 1 + (NQMAX-LMIN) ÷ 2, 1 + LMAX-LMIN)
|
||
@ccall "../OSBRACKETS/allosbrac.so".allosbrac_(NQMAX::Ref{Int32},LMIN::Ref{Int32},LMAX::Ref{Int32},CO::Ref{Float64},SI::Ref{Float64},BRAC::Ptr{Array{Float64}})::Cvoid
|
||
|
||
mat = zeros(basis.dim, basis.dim)
|
||
|
||
s = hcat(basis.Es, basis.n1s, basis.l1s, basis.n2s, basis.l2s)
|
||
Threads.@threads for idx in CartesianIndices(mat)
|
||
(i, j) = Tuple(idx)
|
||
(Elhs, N, L, n, l) = s[i, :]
|
||
(Erhs, n1, l1, n2, l2) = s[j, :]
|
||
if Elhs == Erhs && triangle_ineq(L, l, basis.Λ) && triangle_ineq(l1, l2, basis.Λ)
|
||
mat[i, j] = (-1)^(n1 + n2 + N + n) * pick_Moshinsky_bracket(BRAC, n1, l1, n2, l2, N, L, n, l, basis.Λ)
|
||
end
|
||
end
|
||
|
||
return sparse(mat)
|
||
end
|
||
|
||
function pick_Moshinsky_bracket(BRAC, n1′, l1′, n2′, l2′, n1, l1, n2, l2, Λ) # Efros notation -- don't confuse
|
||
ϵ = (l1 + l2 - Λ) % 2
|
||
NP = (l1′ - l2′ + Λ - ϵ) ÷ 2
|
||
MP = (l1′ + l2′ - Λ - ϵ) ÷ 2
|
||
N = (l1 - l2 + Λ - ϵ) ÷ 2
|
||
M = (l1 + l2 - Λ - ϵ) ÷ 2
|
||
|
||
# BRAC(NP,N1P,MP,N1,N2,N,M,L)
|
||
return BRAC[1 + NP, 1 + n1′, 1 + MP, 1 + n1, 1 + n2, 1 + N, 1 + M, 1]
|
||
end
|
||
|
||
function get_jacobi_V_matrix(V_of_r, basis::ho_basis_2B, μ1ω1, μω_global; atol=10^-6, maxevals=10^5)
|
||
V1 = get_jacobi_V1_matrix(V_of_r, basis, μ1ω1; atol=atol, maxevals=maxevals)
|
||
V2 = get_jacobi_V2_matrix(V_of_r, basis, μω_global; atol=atol, maxevals=maxevals)
|
||
return V1 + V2
|
||
end
|
||
|
||
function get_jacobi_V1_matrix(V_of_r, basis::ho_basis_2B, μ1ω1; atol=10^-6, maxevals=10^5)
|
||
V1_elem(l, n1, n2) = V_numerical(V_of_r, l, n1, n2; μω_gen=μ1ω1, atol=atol, maxevals=maxevals)
|
||
V1 = get_sp_V_matrix(V1_elem, basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; dtype=ComplexF64, E_max=basis.E_max)
|
||
return V1
|
||
end
|
||
|
||
function get_jacobi_V2_matrix(V_of_r, basis::ho_basis_2B, μω_global; atol=10^-6, maxevals=10^5)
|
||
V_relative_elem(l, n1, n2) = V_numerical(V_of_r, l, n1, n2; μω_gen=μω_global, atol=atol, maxevals=maxevals)
|
||
V_relative_cache = prealloc_V_cache(basis.E_max, ComplexF64)
|
||
|
||
V_relative = get_sp_V_matrix(V_relative_elem, basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; dtype=ComplexF64, cache=V_relative_cache) + get_sp_V_matrix(V_relative_elem, basis.n2s, basis.l2s, [basis.n1s, basis.l1s]; dtype=ComplexF64, cache=V_relative_cache)
|
||
U = Moshinsky_transform(basis)
|
||
V2 = transpose(U) * V_relative * U
|
||
|
||
return V2
|
||
end
|
||
|
||
function get_2p_p1p2_matrix(basis::ho_basis_2B, μ1ω1, μ2ω2; dtype=Float64)
|
||
# TODO: Cache for integrals
|
||
integral1(np, lp, n, l) = integral_HO(np, lp, n, l, μ1ω1)
|
||
integral2(np, lp, n, l) = integral_HO(np, lp, n, l, μ2ω2)
|
||
|
||
mat = zeros(dtype, basis.dim, basis.dim)
|
||
Threads.@threads for idx in CartesianIndices(mat)
|
||
(i, j) = Tuple(idx)
|
||
val = racahs_reduction_formula(basis.n1s[i], basis.l1s[i], basis.n2s[i], basis.l2s[i], basis.n1s[j], basis.l1s[j], basis.n2s[j], basis.l2s[j], basis.Λ, integral1, integral2)
|
||
if !(val ≈ 0); mat[idx] = val; end
|
||
end
|
||
return sparse(mat)
|
||
end
|
||
|
||
function get_src_V_matrix(V_of_r, basis::ho_basis_2B, μω, μω_global; atol=10^-6, maxevals=10^5)
|
||
V_elem(l, n1, n2) = V_numerical(V_of_r, l, n1, n2; μω_gen=μω, atol=atol, maxevals=maxevals)
|
||
V_cache = prealloc_V_cache(basis.E_max, ComplexF64)
|
||
|
||
V1 = get_sp_V_matrix(V_elem, basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; dtype=ComplexF64, cache=V_cache)
|
||
V2 = get_sp_V_matrix(V_elem, basis.n2s, basis.l2s, [basis.n1s, basis.l1s]; dtype=ComplexF64, cache=V_cache)
|
||
|
||
V12 = get_src_V12_matrix(V_of_r, basis, μω_global; atol=atol, maxevals=maxevals)
|
||
|
||
return V1 + V2 + V12
|
||
end
|
||
|
||
function get_src_V12_matrix(V_of_r, basis::ho_basis_2B, μω_global; atol=10^-6, maxevals=10^5)
|
||
V_relative_elem(l, n1, n2) = V_numerical(V_of_r, l, n1, n2; μω_gen=μω_global, atol=atol, maxevals=maxevals)
|
||
V_relative = get_sp_V_matrix(V_relative_elem, basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; dtype=ComplexF64, E_max=basis.E_max)
|
||
|
||
U = Moshinsky_transform(basis)
|
||
V12 = transpose(U) * V_relative * U
|
||
|
||
return V12
|
||
end
|
||
|
||
"Basis transformation from HO to momentum space"
|
||
function get_W_matrix(basis_p, basis::ho_basis_2B, μ1ω1, μ2ω2=μ1ω1; weights=true)
|
||
W = zeros(ComplexF64, length(basis_p), basis.dim)
|
||
Threads.@threads for idx in CartesianIndices(W)
|
||
(i1, i2) = Tuple(idx)
|
||
((j1, j2), (k1, w1), (k2, w2)) = basis_p[i1]
|
||
if j1 == basis.l1s[i2] && j2 == basis.l2s[i2]
|
||
elem1 = 1/sqrt(sqrt(μ1ω1)) * (-1)^basis.n1s[i2] * ho_basis(j1, basis.n1s[i2], 1/sqrt(μ1ω1) * k1)
|
||
elem2 = 1/sqrt(sqrt(μ2ω2)) * (-1)^basis.n2s[i2] * ho_basis(j2, basis.n2s[i2], 1/sqrt(μ2ω2) * k2)
|
||
W[idx] = elem1 * elem2 * (weights ? w1 * w2 : 1)
|
||
end
|
||
end
|
||
return sparse(W)
|
||
end
|