3-body XZ in p-space (not working)

EC.jl

@@ -1,4 +1,4 @@
-using SparseArrays, LinearAlgebra, Arpack, Plots
+using Statistics, SparseArrays, LinearAlgebra, Arpack, Plots
 include("common.jl")
 
 "EC model for a Hamiltonian family H(c) = H0 + c * H1"
@@ -12,20 +12,25 @@ mutable struct affine_EC
     H1_EC
     N_EC
 
+    ensemble_size::Int
+    H0_EC_ensemble
+    H1_EC_ensemble
+    N_EC_ensemble
+
     training_E::Vector{ComplexF64}
     exact_E::Vector{ComplexF64}
     extrapolated_E::Vector{ComplexF64}
+    extrapolated_CI::Vector{ComplexF64}
 
-    affine_EC(H0::AbstractMatrix{ComplexF64}, H1::AbstractMatrix{ComplexF64}, weights::Vector{ComplexF64}=ones(ComplexF64, size(H0, 1))) = new(H0, H1, weights, false, nothing, nothing, nothing, ComplexF64[], ComplexF64[], ComplexF64[])
+    affine_EC(H0::AbstractMatrix{ComplexF64}, H1::AbstractMatrix{ComplexF64}, weights::Vector{ComplexF64}=ones(ComplexF64, size(H0, 1)); ensemble_size=0) =
+        new(H0, H1, weights, false, nothing, nothing, nothing, ensemble_size, Matrix[], Matrix[], Matrix[], ComplexF64[], ComplexF64[], ComplexF64[], ComplexF64[])
 end
 
 "Train an EC model for a given range of c values.
 If a list is provided for ref_eval, they are used as reference values for picking the closest eigenvalues at each sampling point.
 If a single number is provided for ref_eval, it is used as a reference for the first point, and the previous eigenvalue is used as the reference for each successive point.
-If pseudo_inv_rtol > 0, the GEVP is avoided using Moore-Penrose psuedoinverse, using this value as the relative tolerance for dropping redundant vectors.
-If orthonormalize_threshold > 0, Gram-Schmidt orthonormalization is performed, overriding Moore-Penrose, using this value as the threshold for dropping redundant vectors.
-If both are = 0, GEVP is solved instead."
-function train!(EC::affine_EC, c_vals; ref_eval=-10.0, pseudo_inv_rtol=1e-6, orthonormalize_threshold=0, CAEC=false, verbose=true, tol=1e-5)
+If orthonormalize_threshold > 0, Gram-Schmidt orthonormalization is performed, using this value as the threshold for dropping redundant vectors."
+function train!(EC::affine_EC, c_vals; ref_eval=-10.0, CAEC=false, gram_schmidt_threshold=0, verbose=true, tol=1e-5)
     training_vecs = Vector{ComplexF64}[]
 
     for c in c_vals
@@ -50,33 +55,47 @@ function train!(EC::affine_EC, c_vals; ref_eval=-10.0, pseudo_inv_rtol=1e-6, ort
 
     CAEC && append!(training_vecs, conj.(training_vecs))
 
-    if orthonormalize_threshold ≠ 0
-        training_vecs = gram_schmidt!(training_vecs, EC.weights, orthonormalize_threshold; verbose=verbose)
-    end
+    (EC.H0_EC, EC.H1_EC, EC.N_EC) = get_reduced_matrices(EC, training_vecs, gram_schmidt_threshold; verbose=true)
 
-    EC_basis = hcat(training_vecs...)
-    weights_mat = spdiagm(EC.weights)
-    EC.H0_EC = transpose(EC_basis) * weights_mat * EC.H0 * EC_basis
-    EC.H1_EC = transpose(EC_basis) * weights_mat * EC.H1 * EC_basis
-
-    if orthonormalize_threshold == 0
-        EC.N_EC = transpose(EC_basis) * weights_mat * EC_basis
-        if pseudo_inv_rtol > 0
-            inv_N_EC = pinv(EC.N_EC; rtol=pseudo_inv_rtol)
-            EC.H0_EC = inv_N_EC * EC.H0_EC
-            EC.H1_EC = inv_N_EC * EC.H1_EC
-            EC.N_EC = nothing
+    for _ in 1:EC.ensemble_size
+        subsample = resample(length(training_vecs))
+        if gram_schmidt_threshold > 0
+            (H0_EC, H1_EC, N_EC) = get_reduced_matrices(EC, training_vecs, gram_schmidt_threshold, subsample; verbose=false)
+            push!(EC.H0_EC_ensemble, H0_EC)
+            push!(EC.H1_EC_ensemble, H1_EC)
+            push!(EC.N_EC_ensemble, N_EC)
+        else
+            push!(EC.H0_EC_ensemble, EC.H0_EC[subsample, subsample])
+            push!(EC.H1_EC_ensemble, EC.H1_EC[subsample, subsample])
+            push!(EC.N_EC_ensemble, EC.N_EC[subsample, subsample])
         end
     end
 
     EC.trained = true
 end
 
+function get_reduced_matrices(EC::affine_EC, training_vecs, gram_schmidt_threshold, subsample=1:length(training_vecs); verbose=false)
+    vecs = deepcopy(training_vecs[subsample])
+
+    if gram_schmidt_threshold > 0; vecs = gram_schmidt!(vecs, EC.weights, gram_schmidt_threshold; verbose=verbose); end
+
+    EC_basis = hcat(vecs...)
+    weights_mat = spdiagm(EC.weights)
+    H0_EC = transpose(EC_basis) * weights_mat * EC.H0 * EC_basis
+    H1_EC = transpose(EC_basis) * weights_mat * EC.H1 * EC_basis
+    N_EC = transpose(EC_basis) * weights_mat * EC_basis
+
+    return (H0_EC, H1_EC, N_EC)
+end
+
+resample(n::Int) = rand(1:n, n) |> unique |> sort
+
 "Extrapolate using a trained EC model for a given range of c values
 If a list is provided for ref_eval, they are used as reference values for picking the closest eigenvalues at each point.
 If a single number is provided for ref_eval, it is used as a reference for the first point, and the previous eigenvalue is used as the reference for each successive point.
-If precalculated_exact_E is provided, ref_eval is ignored."
-function extrapolate!(EC::affine_EC, c_vals; ref_eval=EC.training_E[end], verbose=true, tol=1e-5, precalculated_exact_E=nothing)
+If precalculated_exact_E is provided, ref_eval is ignored.
+If pseudo_inv_tol > 0, the GEVP is avoided using Moore-Penrose psuedoinverse, using this value as the relative tolerance for dropping redundant vectors."
+function extrapolate!(EC::affine_EC, c_vals; ref_eval=EC.training_E[end], pseudo_inv_tol=0, verbose=true, tol=1e-5, precalculated_exact_E=nothing)
     @assert EC.trained "EC model must be trained using train() before extrapolation"
 
     for c in c_vals
@@ -103,9 +122,31 @@ function extrapolate!(EC::affine_EC, c_vals; ref_eval=EC.training_E[end], verbos
 
         verbose && println("Extrapolating for c = $c")
 
-        H_EC = EC.H0_EC + c .* EC.H1_EC
-        evals = isnothing(EC.N_EC) ? eigvals(H_EC) : eigvals(H_EC, EC.N_EC)
+        evals = get_extrapolated_evals(EC.H0_EC, EC.H1_EC, EC.N_EC, c, pseudo_inv_tol)
         push!(EC.extrapolated_E, nearest(evals, current_E))
+
+        if EC.ensemble_size > 0
+            E_ensemble = zeros(ComplexF64, EC.ensemble_size)
+            for i in 1:EC.ensemble_size
+                evals = get_extrapolated_evals(EC.H0_EC_ensemble[i], EC.H1_EC_ensemble[i], EC.N_EC_ensemble[i], c, pseudo_inv_tol)
+                E_ensemble[i] = nearest(evals, current_E)
+            end
+            re_CI = std(real.(E_ensemble))
+            im_CI = std(imag.(E_ensemble))
+            push!(EC.extrapolated_CI, complex(re_CI, im_CI))
+        end
+    end
+end
+
+"Solve the GEVP with or without Moore-Penrose psuedoinverse"
+function get_extrapolated_evals(H0_EC, H1_EC, N_EC, c, pseudo_inv_tol)
+    H_EC = H0_EC + c .* H1_EC
+    if pseudo_inv_tol > 0
+        inv_N_EC = pinv(N_EC; atol=pseudo_inv_tol)
+        H_EC = inv_N_EC * H_EC
+        return eigvals(H_EC)
+    else
+        return eigvals(H_EC, N_EC)
     end
 end
 
@@ -115,8 +156,12 @@ exportCSV(EC::affine_EC, filename) = exportCSV(filename, (EC.training_E, EC.exac
 "Plot EC data and optionally save figure to a file"
 function plot(EC::affine_EC, save_fig_filename=nothing; basis_points=nothing, basis_contour=nothing, xlims=nothing, ylims=nothing)
     scatter(real.(EC.training_E), imag.(EC.training_E), label="training")
-    scatter!(real.(EC.exact_E), imag.(EC.exact_E), label="exact")
-    scatter!(real.(EC.extrapolated_E), imag.(EC.extrapolated_E), label="extrapolated")
+    scatter!(real.(EC.exact_E), imag.(EC.exact_E), label="exact", markercolor=:white)
+    if EC.ensemble_size > 0
+        scatter!(real.(EC.extrapolated_E), imag.(EC.extrapolated_E), xerror=real.(EC.extrapolated_CI), yerror=imag.(EC.extrapolated_CI), label="extrapolated", m=:x)
+    else
+        scatter!(real.(EC.extrapolated_E), imag.(EC.extrapolated_E), label="extrapolated", m=:x)
+    end
 
     isnothing(basis_points) || scatter!(real.(basis_points), imag.(basis_points), m=:x, label="basis")
     isnothing(basis_contour) || plot!(real.(basis_contour), imag.(basis_contour), label="contour")

Project.toml

@@ -1,5 +1,7 @@
 [deps]
 Arpack = "7d9fca2a-8960-54d3-9f78-7d1dccf2cb97"
+CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
@@ -7,6 +9,7 @@ LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
+Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 WignerSymbols = "9f57e263-0b3d-5e2e-b1be-24f2bb48858b"

calculations/2body_data.jl

@@ -0,0 +1,33 @@
+using Roots, DelimitedFiles
+
+include("../EC.jl")
+include("../common.jl")
+include("../p_space.jl")
+
+μ = 0.5
+V_system(c) = (p, q) -> c*(-5*g0(sqrt(3), p, q) + 2*g0(sqrt(10), p, q)) # ResonanceEC: Eq. (20)
+
+# determining c0 with EC
+temp_c = range(1.1, 0.9, 3)
+p, w = get_mesh([0, 8], [256])
+H0 = get_T_matrix(p, μ)
+V = get_V_matrix(V_system(1), p, w)
+EC = affine_EC(H0, V, w)
+train!(EC, temp_c; ref_eval=-0.2, CAEC=false, verbose=false)
+quick_extrapolate(c) = minimum(abs2, get_extrapolated_evals(EC.H0_EC, EC.H1_EC, EC.N_EC, c, 0))
+c0 = find_zero(quick_extrapolate, 0.85)
+
+training_c = range(1.2, 0.9, 9) # original: range(1.35, 0.9, 5)
+extrapolating_c = range(0.78, 0.45, 7) # original: range(0.75, 0.40, 8)
+
+data_c = vcat(training_c, extrapolating_c)
+data_E = [quick_pole_E(V_system(c)) for c in data_c]
+
+# export to CSV
+file = "temp/2body_data.csv"
+delim = ','
+open(file, "w") do f
+    writedlm(f, ["c" "re_E" "im_E"], delim)
+    writedlm(f, [c0 0 0], delim) # first entry for the threshold
+    writedlm(f, hcat(data_c, real.(data_E), imag.(data_E)), delim)
+end

calculations/3body_Berggren_B2R_EC.jl

@@ -5,9 +5,10 @@ include("../EC.jl")
 m = 1.0
 V_of_r(r) =  2 * exp(-(r-3)^2 / (1.5)^2)
 
-vertices = [0, 2 - 0.2im, 3, 4]
-subdivisions = [15, 10, 10]
-jmax = 4
+ϕ = 0.1
+vertices = [0, 4 * exp(-1im * ϕ), 5, 6]
+subdivisions = [40, 12, 15]
+jmax = 6
 
 E_max = 40
 μω_global = 0.5

calculations/3body_CSM_B2R_EC.jl

@@ -5,7 +5,8 @@ include("../EC.jl")
 m = 1.0
 V_of_r(r) =  2 * exp(-(r-3)^2 / (1.5)^2)
 
-vertices = [0, 6 - 0.6im]
+ϕ = 0.1
+vertices = [0, 6 * exp(-1im * ϕ)]
 subdivisions = [50]
 jmax = 4
 
@@ -23,13 +24,13 @@ extrapolating_c = 0.0 : 0.2 : 1.2
 
 training_ref = -2.22 # complete list not needed because identification is simple
 
-exact_E = [4.077092809998592-0.01206085331850782im,
-3.613579042377367-0.006920188044987599im,
-3.145489628680764-0.003757512658877539im,
-2.673033482861357-0.001939576896512454im,
-2.196539134888566-0.0009597849595725841im,
-1.7163902133045392-0.000456595029296216im,
-1.2329696647679096-0.00019879325231064393im]
+exact_E = [4.076662025307587-0.012709842443350328im,
+    3.613318119833891-0.007335804709990623im,
+    3.1453431847006783-0.004030580410326795im,
+    2.672967129943755-0.00211498327461944im,
+    2.196542557810288-0.0010719835443437104im,
+    1.7164583929199813-0.0005455212208182736im,
+    1.233088227541505-0.0003070320106485624im]
 
 EC = affine_EC(H0, Vp, weights)
 train!(EC, training_c; ref_eval=training_ref, CAEC=true)

calculations/3body_HO_B2R_ACCC.jl

@@ -0,0 +1,68 @@
+using Roots, LinearAlgebra, Plots
+
+include("../EC.jl")
+include("../common.jl")
+include("../ho_basis.jl")
+
+V_of_r(r) = 2 * exp(-(r-3)^2 / (1.5)^2)
+Λ = 0
+m = 1.0
+
+ϕ = 0.1
+μω_global = 0.5 * exp(-2im * ϕ)
+E_max = 40
+
+H0 = get_3b_H_matrix(jacobi, V_of_r, μω_global, E_max, Λ, m, true, true)
+
+# Vp = perturbation to make the state artificially bound
+Vp_of_r(r) = -exp(-(r/3)^2)
+@time "Vp" Vp = get_3b_H_matrix(jacobi, Vp_of_r, μω_global, E_max, Λ, m, false, true)
+
+training_ref = -2.22
+extrapolating_ref = [4.076662025307587-0.012709842443350328im,
+    3.613318119833891-0.007335804709990623im,
+    3.1453431847006783-0.004030580410326795im,
+    2.672967129943755-0.00211498327461944im,
+    2.196542557810288-0.0010719835443437104im,
+    1.7164583929199813-0.0005455212208182736im,
+    1.233088227541505-0.0003070320106485624im]
+
+training_c = range(2.8, 1.8, 5)
+extrapolating_c = 0.0 : 0.2 : 1.2
+
+EC = affine_EC(H0, Vp)
+train!(EC, training_c; ref_eval=training_ref, CAEC=true)
+extrapolate!(EC, extrapolating_c; ref_eval=extrapolating_ref)
+
+# determining c0 with EC
+approx_c0 = 1.5
+quick_extrapolate(c) = minimum(abs2, get_extrapolated_evals(EC.H0_EC, EC.H1_EC, EC.N_EC, c, 1e-14))
+c0 = find_zero(quick_extrapolate, approx_c0)
+
+order::Int = ceil((length(training_c) - 1) / 2) # order of the Pade approximant
+
+# Solve coefficients as a linear system
+training_k = alt_sqrt.(EC.training_E)
+M_left_element(c, i) = alt_sqrt(c - c0)^i
+M_left = M_left_element.(training_c, (0:order)')
+M_right = -training_k .* M_left[:, 2:end] # remove the first column
+M = hcat(M_left, M_right) # M = [M_left | M_right]
+sol = M \ training_k
+a = sol[1:order+1]
+b = [1; sol[order+2:end]]
+
+# Pade approximant
+polynomial(a, c) = sum(i -> a[i+1] * alt_sqrt(c - c0)^i, 0:order)
+pade_approx(c) = polynomial(a, c) / polynomial(b, c)
+
+# Extrapolate
+extrapolated_k = pade_approx.([extrapolating_c; training_c])
+extrapolated_E = extrapolated_k .^ 2
+
+# Plotting
+scatter(real.(EC.training_E), imag.(EC.training_E), label="training")
+scatter!(real.(EC.exact_E), imag.(EC.exact_E), label="exact")
+scatter!(real.(EC.extrapolated_E), imag.(EC.extrapolated_E), label="CAEC", m=:x)
+scatter!(real.(extrapolated_E), imag.(extrapolated_E), label="ACCC", m=:+)
+title!("3-body extrapolation with $(length(training_c)) training points")
+savefig("temp/3body_HO_B2R_ACCC-$(length(training_c)).pdf")

calculations/3body_HO_B2R_EC.jl

@@ -5,19 +5,27 @@ V_of_r(r) = 2 * exp(-(r-3)^2 / (1.5)^2)
 Λ = 0
 m = 1.0
 
-μω_global = 0.5 * exp(-2im * pi / 9)
-E_max = 30
+ϕ = 0.1
+μω_global = 0.5 * exp(-2im * ϕ)
+E_max = 40
 
-H0 = get_3b_H_matrix(src, V_of_r, μω_global, E_max, Λ, m, true, true)
+H0 = get_3b_H_matrix(jacobi, V_of_r, μω_global, E_max, Λ, m, true, true)
 
 # Vp = perturbation to make the state artificially bound
 Vp_of_r(r) = -exp(-(r/3)^2)
-@time "Vp" Vp = get_3b_H_matrix(src, Vp_of_r, μω_global, E_max, Λ, m, false, true)
+@time "Vp" Vp = get_3b_H_matrix(jacobi, Vp_of_r, μω_global, E_max, Λ, m, false, true)
 
-training_ref = -0.72763
-extrapolating_ref = 4.0766890719636635 - 0.01275892774109674im
+training_ref = -2.22
 
-training_c = [2.0, 1.9, 1.8]
+extrapolating_ref = [4.076662025307587-0.012709842443350328im,
+    3.613318119833891-0.007335804709990623im,
+    3.1453431847006783-0.004030580410326795im,
+    2.672967129943755-0.00211498327461944im,
+    2.196542557810288-0.0010719835443437104im,
+    1.7164583929199813-0.0005455212208182736im,
+    1.233088227541505-0.0003070320106485624im]
+
+training_c = [2.6, 2.4, 2.2, 2.0, 1.8]
 extrapolating_c = 0.0 : 0.2 : 1.2
 
 EC = affine_EC(H0, Vp)

calculations/3body_HO_R2R_EC.jl

@@ -6,7 +6,7 @@ V_of_r(r) = 2 * exp(-(r-3)^2 / (1.5)^2)
 m = 1.0
 
 μω_global = 0.5 * exp(-2im * pi / 9)
-E_max = 30
+E_max = 40
 
 T = get_3b_H_matrix(src, V_of_r, μω_global, E_max, Λ, m, true, false)
 V = get_3b_H_matrix(src, V_of_r, μω_global, E_max, Λ, m, false, true)

calculations/3body_dis_CSM_EC.jl

@@ -0,0 +1,77 @@
+include("../p_space.jl")
+include("../EC.jl")
+
+Λ = 0
+m = 1.0
+
+# Distinguishable particles: V12 = bound and V13 & V23 = resonant
+Vsubsystem_of_r(r) = -2 * exp(-r^2/4)
+Vdecay_of_r(r) = -exp(-r^2 / 3) + exp(-r^2 / 10)
+
+ϕ = 0.1
+vertices = [0, 6 * exp(-1im * ϕ)]
+subdivisions = [50]
+jmax = 4
+
+E_max = 40
+μω_global = 0.4
+
+# Jacobi coordinates
+μ1ω1 = μω_global * 1/2
+μ2ω2 = μω_global * 2
+μ1 = m * 1/2
+μ2 = m * 2/3
+
+atol=10^-5; maxevals=10^5; R_cutoff=16; verbose=true;
+
+###########
+
+verbose && println("No of threads = ", Threads.nthreads())
+
+V_l(j, k, kp) = Vl_mat_elem(Vsubsystem_of_r, j, k, kp; atol=atol, maxevals=maxevals, R_cutoff=R_cutoff)
+
+ks, ws = get_mesh(vertices, subdivisions)
+weights = repeat(kron(ws, ws), jmax + 1)
+block_size = length(ks)
+tri((j1, j2)) = triangle_ineq(j1, j2, Λ)
+js = collect(Iterators.filter(tri, iter_prod(0:jmax, 0:jmax)))
+basis = iter_prod(js, zip(ks, ws), zip(ks, ws)) # basis = ((j1, j2), (k1, w1), (k2, w2))
+basis_size = length(js) * length(ks)^2
+@assert length(basis) == basis_size "Something wrong with the basis"
+verbose && println("Basis size = $basis_size")
+
+@time "Block diagonal part" begin
+    blocks = [kron_sum(get_H_matrix((k, kp) -> V_l(j1, k, kp), ks, ws, μ1), get_T_matrix(ks, μ2)) for (j1, _) in js]
+    Ha = blockdiag(sparse.(blocks)...)
+end
+
+basis_ho = ho_basis_2B(E_max, Λ)
+verbose && println("HO basis size = ", basis_ho.dim)
+
+@time "V2_HO" V2_HO =  get_jacobi_V2_matrix(Vdecay_of_r, basis_ho, μω_global)
+
+@time "W_right" W_right = get_W_matrix(basis, basis_ho, μ1ω1, μ2ω2; weights=true)
+@time "W_left" W_left = get_W_matrix(basis, basis_ho, μ1ω1, μ2ω2; weights=false)
+
+@time "V2" Vb = W_left * V2_HO * transpose(W_right)
+
+###########
+
+training_c = [-0.55, -0.7, -0.85, -1, -1.2]
+extrapolating_c = [0.2, 0.1, 0.0, -0.1, -0.2, -0.3]
+
+ref_E = -0.5173809356244544
+
+exact_ref = [-0.31360746615280954-0.07689284936870341im
+    -0.3233372403877718-0.06011323914565665im
+    -0.339615582074795-0.04239442037174759im
+    -0.36333816534241997-0.02648721825958402im
+    -0.39376650561322885-0.014382935339817332im	
+    -0.4299479825535172-0.006510710745123606im]
+
+EC = affine_EC(Ha, Vb)
+train!(EC, training_c; ref_eval=ref_E, CAEC=true)
+extrapolate!(EC, extrapolating_c; precalculated_exact_E=exact_ref)
+
+exportCSV(EC, "temp/dis_CSM_B2R.csv")
+plot(EC, "temp/dis_CSM_B2R.pdf")

calculations/3body_dis_HO_EC.jl

@@ -32,14 +32,21 @@ println("Basis size = ", basis.dim)
 @time "Eigenvalues" target_evals, _ = eigs(Ha, nev=5, ncv=50, which=:SR, maxiter=5000, tol=1e-5, ritzvec=false, check=1)
 display(target_evals)
 
-training_c = [-0.5, -0.65, -0.8, -1, -1.2]
-extrapolating_c = [0.8, 0.6, 0.4, 0.2, 0.1, 0.0, -0.1, -0.2, -0.3]
+training_c = [-0.55, -0.7, -0.85, -1, -1.2]
+extrapolating_c = [0.2, 0.1, 0.0, -0.1, -0.2, -0.3]
 
 ref_E = -0.5173809356244544
 
+exact_ref = [-0.3136074661528041-0.07689284936868852im
+    -0.323337240387771-0.06011323914564878im
+    -0.33961558207479553-0.04239442037174764im
+    -0.3633381653424224-0.026487218259589693im
+    -0.393766505613234-0.014382935339825854im
+    -0.42994798255352606-0.006510710745131777im]
+
 EC = affine_EC(Ha, Vb)
 train!(EC, training_c; ref_eval=ref_E, CAEC=true) # try CAEC=false !!!
-extrapolate!(EC, extrapolating_c)
+extrapolate!(EC, extrapolating_c, precalculated_exact_E = exact_ref)
 
 exportCSV(EC, "temp/dis_HO_B2R.csv")
 plot(EC, "temp/dis_HO_B2R.pdf")

calculations/ACCC.jl

@@ -0,0 +1,50 @@
+using Roots, LinearAlgebra, Plots
+
+include("../EC.jl")
+include("../common.jl")
+include("../p_space.jl")
+
+μ = 0.5
+V_system(c) = (p, q) -> c*(-5*g0(sqrt(3), p, q) + 2*g0(sqrt(10), p, q)) # ResonanceEC: Eq. (20)
+
+# determining c0 with EC
+temp_c = range(1.1, 0.9, 3)
+p, w = get_mesh([0, 8], [256])
+H0 = get_T_matrix(p, μ)
+V = get_V_matrix(V_system(1), p, w)
+EC = affine_EC(H0, V, w)
+train!(EC, temp_c; ref_eval=-0.2, CAEC=false, verbose=false)
+quick_extrapolate(c) = minimum(abs2, get_extrapolated_evals(EC.H0_EC, EC.H1_EC, EC.N_EC, c, 0))
+c0 = find_zero(quick_extrapolate, 0.85)
+
+# Calculation of training and extrapolating E
+training_c = range(1.2, 0.9, 9) # original: range(1.35, 0.9, 5)
+training_E = [quick_pole_E(V_system(c)) for c in training_c]
+training_k = alt_sqrt.(2μ .* training_E)
+
+extrapolating_c = range(0.78, 0.45, 7) # original: range(0.75, 0.40, 8)
+exact_E = [quick_pole_E(V_system(c)) for c in extrapolating_c]
+
+order::Int = ceil((length(training_c) - 1) / 2) # order of the Pade approximant
+
+# Solve coefficients as a linear system
+M_left_element(c, i) = alt_sqrt(c - c0)^i
+M_left = M_left_element.(training_c, (0:order)')
+M_right = -training_k .* M_left[:, 2:end] # remove the first column
+M = hcat(M_left, M_right) # M = [M_left | M_right]
+sol = M \ training_k
+a = sol[1:order+1]
+b = [1; sol[order+2:end]]
+
+# Pade approximant
+polynomial(a, c) = sum(i -> a[i+1] * alt_sqrt(c - c0)^i, 0:order)
+pade_approx(c) = polynomial(a, c) / polynomial(b, c)
+
+# Extrapolate
+extrapolated_k = pade_approx.([training_c; extrapolating_c])
+extrapolated_E = (extrapolated_k .^ 2) / (2μ)
+
+# Plotting
+scatter(real.(training_E), imag.(training_E), label="training")
+scatter!(real.(exact_E), imag.(exact_E), label="exact")
+scatter!(real.(extrapolated_E), imag.(extrapolated_E), label="extrapolated", m=:star5)

calculations/PMM.py

@@ -0,0 +1,81 @@
+#%%
+import pandas as pd
+import torch
+import numpy as np
+
+#%%
+df = pd.read_csv('../temp/2body_data.csv').sort_values(by='c')
+df['E'] = df['re_E'] + 1j * df['im_E']
+train_data = df[df['re_E'] < 0]
+target_data = df[df['re_E'] > 0]
+
+train_cs = train_data['c'].to_numpy()
+train_Es = torch.tensor(train_data['E'].to_numpy(), dtype=torch.complex128)
+
+#%%
+# hyperparameters
+N = 9
+
+# initialize random Hamiltonians
+H0 = torch.randn(N, N, dtype=torch.complex128)
+H0 = (H0 + torch.transpose(H0, 0, 1)).requires_grad_() # symmetric
+H1 = torch.randn(N, N, dtype=torch.complex128)
+H1 = (H1 + torch.transpose(H1, 0, 1)).requires_grad_() # symmetric
+
+#%%
+# training
+
+# generate a set of c values to follow by subdividing the training cs
+subdivisions = 3
+c_steps = np.concatenate([np.linspace(start, stop, subdivisions, endpoint=False) for (start, stop) in zip(train_cs, train_cs[1:])])
+c_steps = np.append(c_steps, train_cs[-1])
+
+lr = 0.05
+epochs = 100000
+for epoch in range(epochs):
+    Es = torch.empty(len(train_data), dtype=torch.complex128)
+    current_E = 0.0 # start at the threshold
+    for c in c_steps:
+        H = H0 + c * H1
+        evals = torch.linalg.eigvals(H)
+        current_E = evals[torch.argmin(torch.abs(evals - current_E))]
+        if np.any(c == train_cs):
+            index = np.where(c == train_cs)[0][0]
+            Es[index] = current_E
+
+    loss = ((Es - train_Es).abs() ** 2).sum()
+
+    if epoch % 1000 == 0:
+        print(f"Training {(epoch+1)/epochs:.1%} \t Loss: {loss}")
+
+    if H0.grad is not None:
+        H0.grad.zero_()
+    if H1.grad is not None:
+        H1.grad.zero_()
+    loss.backward()
+
+    with torch.no_grad():
+        H0 -= lr * H0.grad
+        H1 -= lr * H1.grad
+
+# %%
+# evaluate for all points
+all_c = torch.tensor(df['c'].values, dtype=torch.float64)
+exact_E = torch.tensor(df['E'].values, dtype=torch.complex128)
+pred_Es = torch.empty(len(df), dtype=torch.complex128)
+with torch.no_grad():
+    for (index, (c, E)) in enumerate(zip(all_c, exact_E)):
+        H = H0 + c * H1
+        evals = torch.linalg.eigvals(H)
+        i = torch.argmin(torch.abs(evals - E)) # TODO: more robust way to identify the eigenvector
+        pred_Es[index]= evals[i]
+
+# %%
+# plot the results
+import matplotlib.pyplot as plt
+plt.scatter(train_data['re_E'], train_data['im_E'], label='training')
+plt.scatter(target_data['re_E'], target_data['im_E'], label='target')
+plt.scatter(pred_Es.real, pred_Es.imag, marker='x', label='predicted')
+plt.legend()
+
+# %%

calculations/XZ/2body_HO.jl

@@ -0,0 +1,39 @@
+using Plots
+
+include("../../EC.jl")
+include("../../ho_basis.jl")
+include("../../p_space.jl")
+
+angle = 0.25 * pi # DOESN'T WORK WITHOUT ROTATION
+μω_gen = 0.5 * exp(-2im * angle)
+μ = 0.5
+l = 0
+V1 = -5
+R1 = sqrt(3)
+V2 = 2
+R2 = sqrt(10)
+n_max = 15
+
+ns = collect(0:n_max)
+ls = fill(l, n_max + 1)
+
+T = get_sp_T_matrix(ns, ls; μω_gen=μω_gen, μ=μ)
+V = V1 .* V_Gaussian.(R1, l, ns, transpose(ns); μω_gen=μω_gen) + V2 .* V_Gaussian.(R2, l, ns, transpose(ns); μω_gen=μω_gen)
+
+n_EC = 8
+train_cs = (0.7 .+ 0.05 * randn(n_EC)) - 1im * (0.2 .+ 0.05 * randn(n_EC))
+target_cs = [0.5]
+
+near_E = 0.2 + 0.2im
+exact_E = [0.20845136860234303 - 0.07100640993695649im]
+
+EC = affine_EC(T, V; ensemble_size=32)
+train!(EC, train_cs; ref_eval=near_E, CAEC=false)
+extrapolate!(EC, target_cs; precalculated_exact_E=exact_E)
+
+plot(EC; xlims=(0,0.3), ylims=(-0.3,0.3))
+hline!([0], color=:red, label="continuum")
+xlabel!("Re(E)")
+ylabel!("Im(E)")
+plot!(legend=:bottomleft)
+savefig("temp/2b_HO_XZ.pdf")

calculations/XZ/2body_p_space.jl

@@ -0,0 +1,50 @@
+using Plots
+
+include("../../EC.jl")
+include("../../ho_basis.jl")
+include("../../p_space.jl")
+
+# paramters of the system
+
+angle = 0.0
+μ = 0.5
+l = 0
+V1 = -5
+R1 = sqrt(3)
+V2 = 2
+R2 = sqrt(10)
+
+n_EC = 8
+train_cs = (0.7 .+ 0.03 * randn(n_EC)) - 1im * (0.2 .+ 0.03 * randn(n_EC))
+near_E = 0.2 + 0.2im
+
+target_c = 0.5
+exact_E = 0.20845136860234303 - 0.07100640993695649im
+
+vertices = [0, 4 * exp(-1im * angle)]
+subdivisions = [256]
+ks, ws = get_mesh(vertices, subdivisions)
+
+V_of_r(r) = V1 * exp(-r^2 / R1^2) + V2 * exp(-r^2 / R2^2)
+V_mat_elem(k, kp) = Vl_mat_elem(V_of_r, l, k, kp; atol=10^-5, maxevals=10^5, R_cutoff=16)
+V = get_V_matrix(V_mat_elem, ks, ws)
+T = get_T_matrix(ks, μ)
+
+EC_p_space = affine_EC(T, V)
+train!(EC_p_space, train_cs; ref_eval=near_E, CAEC=false)
+extrapolate!(EC_p_space, [target_c]; precalculated_exact_E=[exact_E])
+
+# Plotting
+
+theme(:dark) # Set the global theme to dark
+
+scatter([real(exact_E)], [imag(exact_E)], label="exact", marker=:circle, markercolor=:white, bg = :black) # black background
+scatter!(real.(EC_p_space.training_E), imag.(EC_p_space.training_E), label="training", marker=:circle, color=:blue)
+scatter!(real.(EC_p_space.extrapolated_E), imag.(EC_p_space.extrapolated_E), label="extrapolated", marker=:x, color=:green)
+hline!([0], color=:red, label="continuum")
+plot!(legend=:bottomleft)
+xlabel!("Re(E)")
+ylabel!("Im(E)")
+xlims!(0, 0.3)
+ylims!(-0.3, 0.3)
+savefig("temp/2body_p_space.pdf")

calculations/XZ/3body_HO.jl

@@ -0,0 +1,36 @@
+include("../../ho_basis.jl")
+include("../../EC.jl")
+
+V_of_r(r) = 2 * exp(-(r-3)^2 / (1.5)^2)
+Λ = 0
+m = 1.0
+
+ϕ = 0.1 # DOESN'T WORK WITHOUT ROTATION
+μω_global = 0.5 * exp(-2im * ϕ)
+E_max = 40
+
+H0 = get_3b_H_matrix(jacobi, V_of_r, μω_global, E_max, Λ, m, true, true)
+
+# Vp = perturbation to make the state artificially bound
+Vp_of_r(r) = -exp(-(r/3)^2)
+@time "Vp" Vp = get_3b_H_matrix(jacobi, Vp_of_r, μω_global, E_max, Λ, m, false, true)
+
+training_ref = 2 + 0.5im
+
+exact_E = [4.076642792419057-0.012998408352259658im,
+    3.6129849325287-0.007397677539402868im,
+    3.145212908643357-0.0038660337822150753im,	
+    2.6729225739451596-0.0021090370393881063im,	
+    2.196385760253282-0.0010430088245526555im,	
+    1.7162659936896967-0.0004515351140200029,	
+    1.2329926791785895-0.00017698044022813525im]
+
+training_c = [0.6 - 0.16im] .+ 0.04 .* (randn(8) .+ 0.5im * randn(8))
+extrapolating_c = 0.0 : 0.2 : 1.2
+
+EC = affine_EC(H0, Vp)
+train!(EC, training_c; ref_eval=training_ref, CAEC=false)
+extrapolate!(EC, extrapolating_c; precalculated_exact_E=exact_E)
+
+exportCSV(EC, "temp/3b_HO_XZ.csv")
+plot(EC, "temp/3b_HO_XZ.pdf")

calculations/XZ/3body_p_space.jl

@@ -0,0 +1,45 @@
+include("../../p_space.jl")
+include("../../EC.jl")
+
+using Arpack
+
+# target = 4.0766890719636875 - 0.012758927741074495im
+
+Λ = 0
+m = 1.0
+V_of_r(r) =  2 * exp(-(r-3)^2 / (1.5)^2)
+
+vertices = [0, 2 - 0.2im, 3, 4] # TODO: real contour instead of Berggren basis
+subdivisions = [15, 10, 10]
+jmax = 4
+
+E_max = 40
+μω_global = 0.5
+
+@time "H0" H0, _ = get_3b_H_matrix(jacobi, V_of_r, vertices, subdivisions, jmax, μω_global, E_max, Λ, m, true, true)
+
+# Vp = perturbation to make the state artificially bound
+Vp_of_r(r) = -exp(-(r/3)^2)
+@time "Vp" Vp, _ = get_3b_H_matrix(jacobi, Vp_of_r, vertices, subdivisions, jmax, μω_global, E_max, Λ, m, false, true)
+
+training_ref = 2 + 0.5im
+
+exact_E = [4.076642792419057-0.012998408352259658im,
+    3.6129849325287-0.007397677539402868im,
+    3.145212908643357-0.0038660337822150753im,	
+    2.6729225739451596-0.0021090370393881063im,	
+    2.196385760253282-0.0010430088245526555im,	
+    1.7162659936896967-0.0004515351140200029,	
+    1.2329926791785895-0.00017698044022813525im]
+
+training_c = [-1.5 - 0.5im] .+ (randn(8) .+ 0.05im * randn(8))
+extrapolating_c = 0.0 : 0.2 : 1.2
+
+EC = affine_EC(H0, Vp)
+train!(EC, training_c; ref_eval=training_ref, CAEC=false)
+extrapolate!(EC, extrapolating_c; precalculated_exact_E=exact_E)
+
+exportCSV(EC, "temp/3b_p_space_XZ.csv")
+plot(EC, "temp/3b_p_space_XZ.pdf")
+
+# Results: training points are all over the place, and extrapolated values are garbage.

calculations/XZ/p_space_vs_HO.jl

@@ -0,0 +1,71 @@
+using Plots
+
+include("../../EC.jl")
+include("../../ho_basis.jl")
+include("../../p_space.jl")
+
+# paramters of the system
+
+angle = 0.0 * pi
+μ = 0.5
+l = 0
+V1 = -5
+R1 = sqrt(3)
+V2 = 2
+R2 = sqrt(10)
+
+n_EC = 8
+train_cs = (0.7 .+ 0.03 * randn(n_EC)) - 1im * (0.2 .+ 0.03 * randn(n_EC))
+near_E = 0.2 + 0.2im
+
+target_c = 0.5
+exact_E = 0.20845136860234303 - 0.07100640993695649im
+
+# HO basis
+
+global EC_HO
+begin
+    println("HO basis calculation")
+    μω_gen = 0.5 * exp(-1im * angle)
+    n_max = 40
+    ns = collect(0:n_max)
+    ls = fill(l, n_max + 1)
+
+    T = get_sp_T_matrix(ns, ls; μω_gen=μω_gen, μ=μ)
+    V = V1 .* V_Gaussian.(R1, l, ns, transpose(ns); μω_gen=μω_gen) + V2 .* V_Gaussian.(R2, l, ns, transpose(ns); μω_gen=μω_gen)
+
+    global EC_HO = affine_EC(T, V)
+    train!(EC_HO, train_cs; ref_eval=near_E, CAEC=false)
+    extrapolate!(EC_HO, [target_c]; precalculated_exact_E=[exact_E])
+end
+
+# p-space
+
+global EC_p_space
+begin
+    println("p-space calculation")
+    vertices = [0, 4 * exp(-1im * angle)]
+    subdivisions = [256]
+    ks, ws = get_mesh(vertices, subdivisions)
+
+    V_of_r(r) = V1 * exp(-r^2 / R1^2) + V2 * exp(-r^2 / R2^2)
+    V_mat_elem(k, kp) = Vl_mat_elem(V_of_r, l, k, kp; atol=10^-5, maxevals=10^5, R_cutoff=16)
+    V = get_V_matrix(V_mat_elem, ks, ws)
+    T = get_T_matrix(ks, μ)
+
+    global EC_p_space = affine_EC(T, V)
+    train!(EC_p_space, train_cs; ref_eval=near_E, CAEC=false)
+    extrapolate!(EC_p_space, [target_c]; precalculated_exact_E=[exact_E])
+end
+
+# Plotting
+
+scatter([real(exact_E)], [imag(exact_E)], label="Exact", marker=:circle, markercolor=:white)
+scatter!(real.(EC_HO.training_E), imag.(EC_HO.training_E), label="HO basis training", marker=:circle, color=:blue)
+scatter!(real.(EC_HO.extrapolated_E), imag.(EC_HO.extrapolated_E), label="HO basis extrapolated", marker=:x, color=:blue)
+scatter!(real.(EC_p_space.training_E), imag.(EC_p_space.training_E), label="p-space training", marker=:circle, color=:red)
+scatter!(real.(EC_p_space.extrapolated_E), imag.(EC_p_space.extrapolated_E), label="p-space extrapolated", marker=:x, color=:red)
+plot!(legend=:bottomleft)
+xlabel!("Re(E)")
+ylabel!("Im(E)")
+savefig("temp/2b_p_space_vs_HO.pdf")

calculations/XZ_technique.jl

@@ -1,30 +0,0 @@
-include("../EC.jl")
-include("../ho_basis.jl")
-include("../p_space.jl")
-
-μω_gen = 0.5 * exp(-1im * 0.47 * pi)
-μ = 0.5
-l = 0
-V1 = -5
-R1 = sqrt(3)
-V2 = 2
-R2 = sqrt(10)
-n_max = 15
-
-ns = collect(0:n_max)
-ls = fill(l, n_max + 1)
-
-T = get_sp_T_matrix(ns, ls; μω_gen=μω_gen, μ=μ)
-V = V1 .* V_Gaussian.(R1, l, ns, transpose(ns); μω_gen=μω_gen) + V2 .* V_Gaussian.(R2, l, ns, transpose(ns); μω_gen=μω_gen)
-
-n_EC = 8
-train_cs = (0.7 .+ 0.05 * randn(n_EC)) - 1im * (0.2 .+ 0.05 * randn(n_EC))
-target_cs = range(0.77, 0.22, 6)
-
-near_E = 0.2 + 0.2im
-
-EC = affine_EC(T, V)
-train!(EC, train_cs; ref_eval=near_E, CAEC=false)
-extrapolate!(EC, target_cs)
-
-plot(EC, "temp/XZ.pdf"; xlims=(-0.2,0.3), ylims=(-0.3,0.3))

common.jl

@@ -2,6 +2,9 @@ using LinearAlgebra, DelimitedFiles, SparseArrays
 
 @enum coordinate_system jacobi src
 
+"Square root function with the branch cut along the postive imaginary axis"
+alt_sqrt(x::Number)::ComplexF64 = sqrt(im * x) / sqrt(im)
+
 "Sum over array while minimizing catastrophic cancellation as much as possible"
 function better_sum(arr::Array{T}) where T<:Real
     pos_arr = arr[arr .> 0]

ho_basis.jl

@@ -4,33 +4,7 @@ using LRUCache
 include("common.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"
+"2-body HO basis (used for 3-body systems without the CM DOF)"
 struct ho_basis_2B
     E_max::Int
     Λ::Int

ho_basis_2body.jl

@@ -1,30 +0,0 @@
-using Arpack, SparseArrays
-include("ho_basis.jl")
-
-E_max = 30
-μω_gen = 0.2
-Λ = 0
-m = 1.0
-
-Va = -2
-Ra = 2
-
-μ1 = m * 1/2
-
-println("No of threads = ", Threads.nthreads())
-
-basis = ho_basis_1B(E_max)
-println("Basis size = ", basis.dim)
-
-println("Constructing KE matrices")
-@time "T1" T1 = get_sp_T_matrix(basis.ns, basis.ls; μω_gen=μω_gen, μ=μ1)
-
-println("Constructing PE matrices")
-V1_elem(l, n1, n2) = Va * V_Gaussian(Ra, l, n1, n2; μω_gen=μω_gen)
-@time "V1" V1 = get_sp_V_matrix(V1_elem, basis.ns, basis.ls; E_max=E_max)
-
-println("Calculating spectrum")
-@time "H" H = T1 + V1
-@time "Eigenvalues" evals, _ = eigs(H, nev=3, ncv=30, which=:SR, maxiter=5000, tol=1e-5, ritzvec=false, check=1)
-
-display(evals)

ho_basis_3body_wf.jl

@@ -1,30 +0,0 @@
-using HDF5
-
-include("ho_basis_3body.jl")
-
-@time "Eigenvectors" evals, evecs = eigs(H, nev=1, ncv=10, which=:SR, maxiter=5000, tol=1e-8, ritzvec=true, check=1)
-
-idx = argmin(real.(evals))
-
-E = evals[idx]
-@assert imag(E) ≈ 0 "Energy is not real"
-E = real(E)
-println("Exporting E = $E")
-
-wf = real.(evecs[:, idx])
-@assert all(imag.(wf) .≈ 0) "Wave function is not real"
-wf = real.(wf)
-
-h5open("temp/HO_3b.h5", "w") do fid
-    write_attribute(fid, "E", E)
-    write_attribute(fid, "Λ", Λ)
-    write_attribute(fid, "μ1ω1", μ1ω1)
-    write_attribute(fid, "μ2ω2", μ2ω2)
-    write_attribute(fid, "μ1", μ1)
-    write_attribute(fid, "μ2", μ2)
-    write_dataset(fid, "n1", n1s)
-    write_dataset(fid, "l1", l1s)
-    write_dataset(fid, "n2", n2s)
-    write_dataset(fid, "l2", l2s)
-    write_dataset(fid, "wf", wf)
-end

test/ho_basis_test.jl

@@ -1,5 +1,5 @@
 using LinearAlgebra
-include("ho_basis.jl")
+include("../ho_basis.jl")
 
 l = 0
 V0 = -10