From 0ceb379be712b000fbb3c122ee84852b2e0747de Mon Sep 17 00:00:00 2001
From: Nuwan Yapa <yapanuwan@gmail.com>
Date: Wed, 15 Jan 2025 03:25:46 -0500
Subject: [PATCH] HO systems refactored (breaks EC calculations)

---
 ho_basis.jl                 | 42 +++++++++++++++++++++++++++++++++++++
 ho_basis_3body.jl           | 32 +++-------------------------
 ho_basis_3body_resonance.jl | 32 ++++++----------------------
 3 files changed, 51 insertions(+), 55 deletions(-)

diff --git a/ho_basis.jl b/ho_basis.jl
index 7641cd5..f624815 100644
--- a/ho_basis.jl
+++ b/ho_basis.jl
@@ -236,3 +236,45 @@ function get_W_matrix(basis_p, basis::ho_basis_2B, μ1ω1, μ2ω2=μ1ω1; weight
     end
     return sparse(W)
 end
+
+@enum coordinate_system jacobi src
+
+function get_3b_H_matrix(coord_system::coordinate_system, V_of_r, μω_global, E_max, Λ, m = 1.0; atol = 10^-5, maxevals = 10^5, verbose=true)
+
+    if coord_system == jacobi
+        μ1ω1 = μω_global * 1/2
+        μ2ω2 = μω_global * 2
+        μ1 = m * 1/2
+        μ2 = m * 2/3        
+    elseif coord_system == src
+        μ1ω1 = μ2ω2 = μω = μω_global * 2
+        μ1 = μ2 = μ = m/2
+    end
+
+    verbose || println("No of threads = ", Threads.nthreads())
+
+    @time "Basis" basis = ho_basis_2B(E_max, Λ)
+
+    verbose || println("Basis size = ", basis.dim)
+
+    verbose || println("Constructing KE matrices")
+    @time "T1" T1 = get_sp_T_matrix(basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; μω_gen=μ1ω1, μ=μ1)
+    @time "T2" T2 = get_sp_T_matrix(basis.n2s, basis.l2s, [basis.n1s, basis.l1s]; μω_gen=μ2ω2, μ=μ2)
+    if coord_system == src
+        @time "T_cross" T_cross = get_2p_p1p2_matrix(basis, μ1ω1, μ2ω2; dtype=ComplexF64) ./ (2*μ)
+    else
+        T_cross = spzeros(basis.dim, basis.dim)
+    end
+    
+    verbose || println("Constructing PE matrices")
+    if coord_system == jacobi
+        @time "V" V = get_jacobi_V_matrix(V_of_r, basis, μ1ω1, μω_global; atol=atol, maxevals=maxevals)
+    elseif coord_system == src
+        @time "V" V = get_src_V_matrix(V_of_r, basis, μω, μω_global; atol=atol, maxevals=maxevals)
+    end
+
+    verbose || println("Adding up all matrices")
+    @time "H" H = T1 + T2 + T_cross + V
+
+    return H
+end
\ No newline at end of file
diff --git a/ho_basis_3body.jl b/ho_basis_3body.jl
index 5797bbd..b0a4a41 100644
--- a/ho_basis_3body.jl
+++ b/ho_basis_3body.jl
@@ -1,40 +1,14 @@
-using Arpack, SparseArrays
+using Arpack
 include("ho_basis.jl")
 
 Λ = 0
 m = 1.0
-Va = -2
-Ra = 2
+V_of_r(r) = -2 * exp(-r^2 / 4)
 
 E_max = 40
 μω_global = 0.3
 
-# due to simple relative coordinates
-μω = μω_global * 2
-μ = m/2
+H = get_3b_H_matrix(src, V_of_r, μω_global, E_max, Λ, m)
 
-println("No of threads = ", Threads.nthreads())
-
-@time "Basis" basis = ho_basis_2B(E_max, Λ)
-
-println("Basis size = ", basis.dim)
-
-println("Constructing KE matrices")
-@time "T1" T1 = get_sp_T_matrix(basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; μω_gen=μω, μ=μ)
-@time "T2" T2 = get_sp_T_matrix(basis.n2s, basis.l2s, [basis.n1s, basis.l1s]; μω_gen=μω, μ=μ)
-@time "T_cross" T_cross = get_2p_p1p2_matrix(basis, μω, μω) ./ (2*μ)
-
-println("Constructing PE matrices")
-V_elem(l, n1, n2) = Va * V_Gaussian(Ra, l, n1, n2; μω_gen=μω)
-V_relative_elem(l, n1, n2) = Va * V_Gaussian(Ra, l, n1, n2; μω_gen=μω_global)
-@time "V1" V1 = get_sp_V_matrix(V_elem, basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; E_max=E_max)
-@time "V2" V2 = get_sp_V_matrix(V_elem, basis.n2s, basis.l2s, [basis.n1s, basis.l1s]; E_max=E_max)
-@time "V relative" V_relative = get_sp_V_matrix(V_relative_elem, basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; E_max=E_max)
-@time "Moshinsky brackets" U = Moshinsky_transform(basis)
-@time "V12" V12 =  transpose(U) * V_relative * U
-
-println("Calculating spectrum")
-@time "H" H = T1 + T2 + T_cross + V1 + V2 + V12
 @time "Eigenvalues" evals, _ = eigs(H, nev=3, ncv=30, which=:SR, maxiter=5000, tol=1e-5, ritzvec=false, check=1)
-
 display(evals)
\ No newline at end of file
diff --git a/ho_basis_3body_resonance.jl b/ho_basis_3body_resonance.jl
index 998621a..0fa2e43 100644
--- a/ho_basis_3body_resonance.jl
+++ b/ho_basis_3body_resonance.jl
@@ -1,35 +1,15 @@
-using Arpack, SparseArrays, LRUCache
+using Arpack
 include("ho_basis.jl")
 
+target_E = 4.07656088827514 - 0.012743522750966718im
+V_of_r(r) = 2 * exp(-(r-3)^2 / (1.5)^2)
 Λ = 0
 m = 1.0
-V_of_r(r) = 2 * exp(-(r-3)^2 / (1.5)^2)
 
-E_max = 30
 μω_global = 0.5 * exp(-2im * pi / 9)
+E_max = 30
 
-# due to simple relative coordinates
-μω = μω_global * 2
-μ = m/2
-
-println("No of threads = ", Threads.nthreads())
-
-@time "Basis" basis = ho_basis_2B(E_max, Λ)
-
-println("Basis size = ", basis.dim)
-
-println("Constructing KE matrices")
-
-@time "T1" T1 = get_sp_T_matrix(basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; μω_gen=μω, μ=μ)
-@time "T2" T2 = get_sp_T_matrix(basis.n2s, basis.l2s, [basis.n1s, basis.l1s]; μω_gen=μω, μ=μ)
-@time "T_cross" T_cross = get_2p_p1p2_matrix(basis, μω, μω; dtype=ComplexF64) ./ (2*μ)
-
-println("Constructing PE matrices")
-
-@time "V" V = get_src_V_matrix(V_of_r, basis, μω, μω_global)
-
-println("Calculating spectrum")
-@time "H" H = T1 + T2 + T_cross + V
-@time "Eigenvalues" evals, _ = eigs(H, nev=5, ncv=50, which=:LI, maxiter=5000, tol=1e-5, ritzvec=false, check=1)
+H = get_3b_H_matrix(src, V_of_r, μω_global, E_max, Λ, m)
 
+@time "Eigenvalues" evals, _ = eigs(H, nev=5, ncv=50, sigma=target_E, maxiter=5000, tol=1e-5, ritzvec=false, check=1)
 display(evals)
\ No newline at end of file