Implemented LRUCache for V matrix elements

This commit is contained in:
Nuwan Yapa 2024-04-24 15:14:32 -04:00
parent aac1a2e431
commit 6a3590cd85
3 changed files with 11 additions and 14 deletions

View File

@ -1,4 +1,4 @@
using Arpack, SparseArrays using Arpack, SparseArrays, LRUCache
include("ho_basis.jl") include("ho_basis.jl")
include("p_space.jl") include("p_space.jl")
@ -38,8 +38,8 @@ atol = 10^-6
maxevals = 10^5 maxevals = 10^5
V1_elem(l, n1, n2) = V_numerical(V_of_r, l, n1, n2; μω_gen=μ1ω1, atol=atol, maxevals=maxevals) V1_elem(l, n1, n2) = V_numerical(V_of_r, l, n1, n2; μω_gen=μ1ω1, atol=atol, maxevals=maxevals)
V_relative_elem(l, n1, n2) = V_numerical(V_of_r, l, n1, n2; μω_gen=μω_global, atol=atol, maxevals=maxevals) V_relative_elem(l, n1, n2) = V_numerical(V_of_r, l, n1, n2; μω_gen=μω_global, atol=atol, maxevals=maxevals)
V1_cache = fill(Complex(NaN), 1+l_max, 1+n_max, 1+n_max) V1_cache = LRU{Tuple{UInt8, UInt8, UInt8}, ComplexF64}(maxsize=(1+l_max)*(1+n_max)^2)
V_relative_cache = fill(Complex(NaN), 1+l_max, 1+n_max, 1+n_max) V_relative_cache = LRU{Tuple{UInt8, UInt8, UInt8}, ComplexF64}(maxsize=(1+l_max)*(1+n_max)^2)
@time "V1" V1 = sp_V_matrix(V1_elem, n1s, l1s; mask=mask1, dtype=ComplexF64, cache=V1_cache) @time "V1" V1 = sp_V_matrix(V1_elem, n1s, l1s; mask=mask1, dtype=ComplexF64, cache=V1_cache)
@time "V relative" V_relative = sp_V_matrix(V_relative_elem, n1s, l1s; mask=mask1, dtype=ComplexF64, cache=V_relative_cache) + sp_V_matrix(V_relative_elem, n2s, l2s; mask=mask2, dtype=ComplexF64, cache=V_relative_cache) @time "V relative" V_relative = sp_V_matrix(V_relative_elem, n1s, l1s; mask=mask1, dtype=ComplexF64, cache=V_relative_cache) + sp_V_matrix(V_relative_elem, n2s, l2s; mask=mask2, dtype=ComplexF64, cache=V_relative_cache)

View File

@ -1,6 +1,7 @@
using SparseArrays using SparseArrays
using SpecialFunctions using SpecialFunctions
using QuadGK using QuadGK
using LRUCache
include("helper.jl") include("helper.jl")
# Gaussian potentials in HO space # Gaussian potentials in HO space
@ -71,19 +72,15 @@ function sp_T_matrix(ns, ls; mask=trues(length(ns),length(ns)), μω_gen=1.0, μ
return (μω_gen / μ) .* mat return (μω_gen / μ) .* mat
end end
function sp_V_matrix(V_l, ns, ls; mask=trues(length(ns),length(ns)), dtype=Float64, cache=fill(convert(dtype, NaN), 1+maximum(ls), 1+maximum(ns), 1+maximum(ns))) function sp_V_matrix(V_l, ns, ls; mask=trues(length(ns),length(ns)), dtype=Float64, cache=LRU{Tuple{UInt8, UInt8, UInt8}, dtype}(maxsize=(1+maximum(ns))^2))
mat = zeros(dtype, length(ns), length(ns)) mat = zeros(dtype, length(ns), length(ns))
Threads.@threads for idx in CartesianIndices(mat) Threads.@threads for idx in CartesianIndices(mat)
if !mask[idx]; continue; end if !mask[idx]; continue; end
(i, j) = Tuple(idx) (i, j) = Tuple(idx)
if ls[i] == ls[j] if ls[i] == ls[j]
l = ls[i] l = UInt8(ls[i])
n1, n2 = minmax(ns[i], ns[j]) # assuming transpose symmetry n1, n2 = UInt8.(minmax(ns[i], ns[j])) # assuming transpose symmetry
if isnan(cache[1+l, 1+n1, 1+n2]) mat[idx] = (get!(cache, (l, n1, n2)) do; V_l(l, n1, n2); end)
cache[1+l, 1+n1, 1+n2] = V_l(l, n1, n2) # hopefully no race condition
@assert !isnan(cache[1+l, 1+n1, 1+n2]) "V matrix element returned NaN"
end
mat[idx] = cache[1+l, 1+n1, 1+n2]
end end
end end
return sparse(mat) return sparse(mat)

View File

@ -1,4 +1,4 @@
using Arpack, SparseArrays using Arpack, SparseArrays, LRUCache
include("ho_basis.jl") include("ho_basis.jl")
include("p_space.jl") include("p_space.jl")
@ -35,8 +35,8 @@ println("Constructing KE matrices")
println("Constructing PE matrices") println("Constructing PE matrices")
V1_elem(l, n1, n2) = Va * V_Gaussian(Ra, l, n1, n2; μω_gen=μ1ω1) V1_elem(l, n1, n2) = Va * V_Gaussian(Ra, l, n1, n2; μω_gen=μ1ω1)
V_relative_elem(l, n1, n2) = Va * V_Gaussian(Ra, l, n1, n2; μω_gen=μω_global) V_relative_elem(l, n1, n2) = Va * V_Gaussian(Ra, l, n1, n2; μω_gen=μω_global)
V1_cache = fill(NaN, 1+l_max, 1+n_max, 1+n_max) V1_cache = LRU{Tuple{UInt8, UInt8, UInt8}, Float64}(maxsize=(1+l_max)*(1+n_max)^2)
V_relative_cache = fill(NaN, 1+l_max, 1+n_max, 1+n_max) V_relative_cache = LRU{Tuple{UInt8, UInt8, UInt8}, Float64}(maxsize=(1+l_max)*(1+n_max)^2)
@time "V1" V1 = sp_V_matrix(V1_elem, n1s, l1s; mask=mask1, cache=V1_cache) @time "V1" V1 = sp_V_matrix(V1_elem, n1s, l1s; mask=mask1, cache=V1_cache)
@time "V relative" V_relative = sp_V_matrix(V_relative_elem, n1s, l1s; mask=mask1, cache=V_relative_cache) + sp_V_matrix(V_relative_elem, n2s, l2s; mask=mask2, cache=V_relative_cache) @time "V relative" V_relative = sp_V_matrix(V_relative_elem, n1s, l1s; mask=mask1, cache=V_relative_cache) + sp_V_matrix(V_relative_elem, n2s, l2s; mask=mask2, cache=V_relative_cache)
@time "Moshinsky brackets" U = Moshinsky_transform(Es, n1s, l1s, n2s, l2s, Λ) @time "Moshinsky brackets" U = Moshinsky_transform(Es, n1s, l1s, n2s, l2s, Λ)