42 changed files with 879 additions and 1235 deletions
--- a/EC.jl
+++ b/EC.jl
@ -1,173 +0,0 @@
 using Statistics, SparseArrays, LinearAlgebra, Arpack, Plots
 include("common.jl")
 "EC model for a Hamiltonian family H(c) = H0 + c * H1"
 mutable struct affine_EC
    H0::AbstractMatrix{ComplexF64}
    H1::AbstractMatrix{ComplexF64}
    weights::Vector{ComplexF64}
    trained::Bool
    H0_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)); 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 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
        verbose && println("Training for c = $c")
        global current_E
        if ref_eval isa Number
            current_E = ref_eval
            ref_eval = nothing
        elseif !isnothing(ref_eval)
            current_E = popfirst!(ref_eval)
        end
        H = EC.H0 + c .* EC.H1
        evals, evecs = eigs(H, sigma=current_E, maxiter=5000, tol=tol, ritzvec=true, check=1)
        current_E = nearest(evals, current_E)
        push!(EC.training_E, current_E)
        push!(training_vecs, evecs[:, nearestIndex(evals, current_E)])
    end
    CAEC && append!(training_vecs, conj.(training_vecs))
    (EC.H0_EC, EC.H1_EC, EC.N_EC) = get_reduced_matrices(EC, training_vecs, gram_schmidt_threshold; verbose=true)
    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.
 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
        global current_E
        if isnothing(precalculated_exact_E)
            if ref_eval isa Number
                current_E = ref_eval
                ref_eval = nothing
            elseif !isnothing(ref_eval)
                current_E = popfirst!(ref_eval)
            end
            verbose && println("Extact solution for c = $c")
            H = EC.H0 + c .* EC.H1
            evals, _ = eigs(H, sigma=current_E, maxiter=5000, tol=tol, ritzvec=false, check=1)
            current_E = nearest(evals, current_E)
        else
            current_E = popfirst!(precalculated_exact_E)
        end
        push!(EC.exact_E, current_E)
        verbose && println("Extrapolating for c = $c")
        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
 "Export EC data as CSV file"
 exportCSV(EC::affine_EC, filename) = exportCSV(filename, (EC.training_E, EC.exact_E, EC.extrapolated_E), ("training", "exact", "extrapolated"))
 "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", 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")
    isnothing(xlims) || xlims!(xlims...)
    isnothing(ylims) || ylims!(ylims...)
    isnothing(save_fig_filename) || savefig(save_fig_filename)
 end
--- a/Project.toml
+++ b/Project.toml
@ -1,7 +1,5 @@
 [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"
@ -9,7 +7,6 @@ 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"
--- a/berggren_3body_resonance.jl
+++ b/berggren_3body_resonance.jl
@ -1,5 +1,5 @@
 using LinearAlgebra, SparseArrays, Arpack
-include("common.jl")
+include("helper.jl")
 include("p_space.jl")
 include("ho_basis.jl")
 include("berggren.jl")
--- a/calculations/2body_data.jl
+++ b/calculations/2body_data.jl
@ -1,33 +0,0 @@
 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
--- a/calculations/2in3body_HO_B2R_EC.jl
+++ b/calculations/2in3body_HO_B2R_EC.jl
@ -0,0 +1,83 @@
 using Arpack, SparseArrays, LRUCache
 using DelimitedFiles, Plots
 include("../ho_basis.jl")
 Λ = 0
 m = 1.0
 V_of_r(r) = -5 * exp(-r^2 / 3) + 2 * exp(-r^2 / 10)
 E_max = 40
 μω_global = 0.2 * exp(-2im * pi / 9)
 # due to Jacobi coordinates
 μ1ω1 = μω_global * 1/2
 μ2ω2 = μω_global * 2
 μ1 = m * 1/2
 μ2 = m * 2/3
 println("No of threads = ", Threads.nthreads())
 basis = ho_basis_2B(E_max, Λ)
 println("Basis size = ", basis.dim)
@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)
@time "V" Vb = get_jacobi_V_matrix(V_of_r, basis, μ1ω1, μω_global)
@time "H0" Ha = T1 + T2
 # free memory
 basis = T1 = T2 = V1_cache = V_relative_cache = V1 = V_relative = U = V2 = nothing
 GC.gc()
 current_E = -0.26141959851000807
 training_c = range(1.1, 0.9, 5)
 extrapolating_c = range(0.78, 0.45, 7)
 exact = ComplexF64[]
 training = ComplexF64[]
 extrapolated = ComplexF64[]
 training_vecs = Vector{ComplexF64}[]
 for c in training_c
    println("Training for c = $c")
    H = Ha + c .* Vb
    evals, evecs = eigs(H, nev=3, ncv=24, which=:LI, maxiter=5000, tol=1e-5, ritzvec=true, check=1)
    global current_E = nearest(evals, current_E)
    push!(training, current_E)
    push!(training_vecs, evecs[:, nearestIndex(evals, current_E)])
 end
 # CA-EC
 training_vecs = vcat(training_vecs, conj(training_vecs))
 EC_basis = hcat(training_vecs...)
 N_EC = transpose(EC_basis) * EC_basis
 Ha_EC = transpose(EC_basis) * Ha * EC_basis
 Vb_EC = transpose(EC_basis) * Vb * EC_basis
 current_E = 0.05387926316280764 - 0.008900278183544552im
 for c in extrapolating_c
    println("Extrapolating for c = $c")
    H = Ha + c .* Vb
    evals, evecs = eigs(H, nev=3, ncv=24, which=:LI, maxiter=5000, tol=1e-5, ritzvec=true, check=1)
    global current_E = nearest(evals, current_E)
    push!(exact, current_E)
    # extrapolation
    H_EC = Ha_EC + c .* Vb_EC
    evals = eigvals(H_EC, N_EC)
    push!(extrapolated, nearest(evals, current_E))
 end
 scatter(real.(training),imag.(training), label="training")
 scatter!(real.(exact),imag.(exact), label="exact")
 scatter!(real.(extrapolated),imag.(extrapolated), label="extrapolated")
 xlims!(-0.3, 0.3)
 ylims!(-0.2, 0)
 savefig("temp/2in3_HO_B2R.pdf")
--- a/calculations/3body_Berggren_B2R_EC.jl
+++ b/calculations/3body_Berggren_B2R_EC.jl
@ -1,40 +1,82 @@
-include("../p_space.jl")
+using Plots
 include("../EC.jl")
 Λ = 0
 m = 1.0
 V_of_r(r) =  2 * exp(-(r-3)^2 / (1.5)^2)
 ϕ = 0.1
 vertices = [0, 4 * exp(-1im * ϕ), 5, 6]
 subdivisions = [40, 12, 15]
 jmax = 6
 E_max = 40
 μω_global = 0.5
 H0, weights = get_3b_H_matrix(jacobi, V_of_r, vertices, subdivisions, jmax, μω_global, E_max, Λ, m)
 # Vp = perturbation to make the state artificially bound
 Vp_of_r(r) = -exp(-(r/3)^2)
 Vp, _ = get_3b_H_matrix(jacobi, Vp_of_r, vertices, subdivisions, jmax, μω_global, E_max, Λ, m, false, true)
 training_c = [2.6, 2.4, 2.2, 2.0, 1.8]
 extrapolating_c = 0.0 : 0.2 : 1.2
 training_ref = -2.22 # complete list not needed because identification is simple
-extrapolating_ref = [4.076662025307587-0.012709842443350328im,
+exact_ref = reverse([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]
+    1.233088227541505-0.0003070320106485624im])
-EC = affine_EC(H0, Vp, weights)
+include("../p_space_3body_resonance.jl")
-train!(EC, training_c; ref_eval=training_ref, CAEC=true)
+H0 = H
 extrapolate!(EC, extrapolating_c; ref_eval=extrapolating_ref)
-exportCSV(EC, "temp/Berggren_B2R.csv")
+# Vp = perturbation to make the state artificially bound
-plot(EC, "temp/Berggren_B2R.pdf")
+Vp_of_r(r) = -exp(-(r/3)^2)
 Vp_l(j, k, kp) = Vl_mat_elem(Vp_of_r, j, k, kp; atol=atol, maxevals=maxevals, R_cutoff=R_cutoff)
@time "Vp block diagonal part" begin
    Vpb_blocks = [kron_sum(get_V_matrix((k, kp) -> Vp_l(j1, k, kp), ks, ws), spzeros(length(ks), length(ks))) for (j1, _) in js]
    Vpb = blockdiag(sparse.(Vpb_blocks)...)
 end
@time "Vp2_HO" Vp2_HO =  get_jacobi_V2_matrix(Vp_of_r, basis_ho, μω_global; atol=atol, maxevals=maxevals)
@time "Vp2" Vp2 = W_left * Vp2_HO * transpose(W_right)
@time "Vp" Vp = Vpb + Vp2
 # free memory
 basis = Hb_blocks = Hb = basis_ho = V2_HO = W_right = W_left = V2 = nothing
 GC.gc()
 exact = ComplexF64[]
 training = ComplexF64[]
 extrapolated = ComplexF64[]
 training_vecs = Vector{ComplexF64}[]
 current_E = training_ref
 for c in training_c
    println("Training for c = $c")
    local H = H0 + c .* Vp
    local evals, evecs = eigs(H, sigma=current_E, maxiter=5000, tol=1e-5, ritzvec=true, check=1)
    global current_E = nearest(evals, current_E)
    push!(training, current_E)
    push!(training_vecs, evecs[:, nearestIndex(evals, current_E)])
 end
 training_vecs = vcat(training_vecs, conj(training_vecs)) # CA-EC
 EC_basis = hcat(training_vecs...)
 weights_mat = spdiagm(repeat(kron(ws, ws), jmax + 1))
 N_EC = transpose(EC_basis) * weights_mat * EC_basis
 H0_EC = transpose(EC_basis) * weights_mat * H0 * EC_basis
 Vp_EC = transpose(EC_basis) * weights_mat * Vp * EC_basis
 for c in extrapolating_c
    println("Extrapolating for c = $c")
    global current_E = pop!(exact_ref)
    local H = H0 + c .* Vp
    local evals, _ = eigs(H, sigma=current_E, maxiter=5000, tol=1e-5, ritzvec=false, check=1)
    global current_E = nearest(evals, current_E)
    push!(exact, current_E)
    # extrapolation
    H_EC = H0_EC + c .* Vp_EC
    evals = eigvals(H_EC, N_EC)
    push!(extrapolated, nearest(evals, current_E))
 end
 exportCSV("temp/Berggren_B2R.csv", (training, exact, extrapolated), ("training", "exact", "extrapolated"))
 scatter(real.(training),imag.(training), label="training")
 scatter!(real.(exact),imag.(exact), label="exact")
 scatter!(real.(extrapolated),imag.(extrapolated), label="extrapolated")
 savefig("temp/Berggren_B2R.pdf")
--- a/calculations/3body_Berggren_R2R_EC.jl
+++ b/calculations/3body_Berggren_R2R_EC.jl
@ -1,42 +1,81 @@
-include("../p_space.jl")
+using Plots
 include("../EC.jl")
 Λ = 0
 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
 E_max = 40
 μω_global = 0.5
 H0, weights = get_3b_H_matrix(jacobi, V_of_r, vertices, subdivisions, jmax, μω_global, E_max, Λ, m)
 # Vp = perturbation to make the state artificially bound
 Vp_of_r(r) = -exp(-(r/3)^2)
 Vp, _ = get_3b_H_matrix(jacobi, Vp_of_r, vertices, subdivisions, jmax, μω_global, E_max, Λ, m, false, true)
 training_c = [1.1, 0.9, 0.7, 0.5]
 extrapolating_c = 0.0 : 0.2 : 1.2
-training_ref = [1.4750633616275919 - 0.0003021770706749637im
+training_ref = reverse([1.4750633616275919 - 0.0003021770706749637im
    1.9567078295375822 - 0.0007646829108872369im
    2.4351117758403076 - 0.001281037843108658im
-    2.9096543462392357 - 0.002962488527470604im]
+    2.9096543462392357 - 0.002962488527470604im])
-extrapolating_ref = [4.076662025307587-0.012709842443350328im,
+exact_ref = reverse([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]
+    1.233088227541505-0.0003070320106485624im])
-EC = affine_EC(H0, Vp, weights)
+include("../p_space_3body_resonance.jl")
-train!(EC, training_c; ref_eval=training_ref, CAEC=false)
+H0 = H
 extrapolate!(EC, extrapolating_c; ref_eval=extrapolating_ref)
-exportCSV(EC, "temp/Berggren_R2R.csv")
+# Vp = perturbation to make the state artificially bound
-plot(EC, "temp/Berggren_R2R.pdf")
+Vp_of_r(r) = -exp(-(r/3)^2)
 Vp_l(j, k, kp) = Vl_mat_elem(Vp_of_r, j, k, kp; atol=atol, maxevals=maxevals, R_cutoff=R_cutoff)
@time "Vp block diagonal part" begin
    Vpb_blocks = [kron_sum(get_V_matrix((k, kp) -> Vp_l(j1, k, kp), ks, ws), spzeros(length(ks), length(ks))) for (j1, _) in js]
    Vpb = blockdiag(sparse.(Vpb_blocks)...)
 end
@time "Vp2_HO" Vp2_HO =  get_jacobi_V2_matrix(Vp_of_r, basis_ho, μω_global; atol=atol, maxevals=maxevals)
@time "Vp2" Vp2 = W_left * Vp2_HO * transpose(W_right)
@time "Vp" Vp = Vpb + Vp2
 # free memory
 basis = Hb_blocks = Hb = basis_ho = V2_HO = W_right = W_left = V2 = nothing
 GC.gc()
 exact = ComplexF64[]
 training = ComplexF64[]
 extrapolated = ComplexF64[]
 training_vecs = Vector{ComplexF64}[]
 for c in training_c
    println("Training for c = $c")
    global current_E = pop!(training_ref)
    local H = H0 + c .* Vp
    local evals, evecs = eigs(H, sigma=current_E, maxiter=5000, tol=1e-5, ritzvec=true, check=1)
    global current_E = nearest(evals, current_E)
    push!(training, current_E)
    push!(training_vecs, evecs[:, nearestIndex(evals, current_E)])
 end
 EC_basis = hcat(training_vecs...)
 weights_mat = spdiagm(repeat(kron(ws, ws), jmax + 1))
 N_EC = transpose(EC_basis) * weights_mat * EC_basis
 H0_EC = transpose(EC_basis) * weights_mat * H0 * EC_basis
 Vp_EC = transpose(EC_basis) * weights_mat * Vp * EC_basis
 for c in extrapolating_c
    println("Extrapolating for c = $c")
    global current_E = pop!(exact_ref)
    local H = H0 + c .* Vp
    local evals, _ = eigs(H, sigma=current_E, maxiter=5000, tol=1e-5, ritzvec=false, check=1)
    global current_E = nearest(evals, current_E)
    push!(exact, current_E)
    # extrapolation
    H_EC = H0_EC + c .* Vp_EC
    evals = eigvals(H_EC, N_EC)
    push!(extrapolated, nearest(evals, current_E))
 end
 scatter(real.(training),imag.(training), label="training")
 scatter!(real.(exact),imag.(exact), label="exact")
 scatter!(real.(extrapolated),imag.(extrapolated), label="extrapolated")
 savefig("temp/Berggren_R2R.pdf")
--- a/calculations/3body_CSM_B2R_EC.jl
+++ b/calculations/3body_CSM_B2R_EC.jl
@ -1,40 +0,0 @@
 include("../p_space.jl")
 include("../EC.jl")
 Λ = 0
 m = 1.0
 V_of_r(r) =  2 * exp(-(r-3)^2 / (1.5)^2)
 ϕ = 0.1
 vertices = [0, 6 * exp(-1im * ϕ)]
 subdivisions = [50]
 jmax = 4
 E_max = 40
 μω_global = 0.5
 H0, weights = get_3b_H_matrix(jacobi, V_of_r, vertices, subdivisions, jmax, μω_global, E_max, Λ, m)
 # Vp = perturbation to make the state artificially bound
 Vp_of_r(r) = -exp(-(r/3)^2)
 Vp, _ = get_3b_H_matrix(jacobi, Vp_of_r, vertices, subdivisions, jmax, μω_global, E_max, Λ, m, false, true)
 training_c = [2.6, 2.4, 2.2, 2.0, 1.8]
 extrapolating_c = 0.0 : 0.2 : 1.2
 training_ref = -2.22 # complete list not needed because identification is simple
 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)
 extrapolate!(EC, extrapolating_c; precalculated_exact_E=exact_E)
 exportCSV(EC, "temp/CSM_B2R.csv")
 plot(EC, "temp/CSM_B2R.pdf")
--- a/calculations/3body_HO_B2R_ACCC.jl
+++ b/calculations/3body_HO_B2R_ACCC.jl
@ -1,68 +0,0 @@
 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")
--- a/calculations/3body_HO_B2R_EC.jl
+++ b/calculations/3body_HO_B2R_EC.jl
@ -1,36 +1,62 @@
-include("../ho_basis.jl")
+using Plots
 include("../EC.jl")
-V_of_r(r) = 2 * exp(-(r-3)^2 / (1.5)^2)
+training_ref = -0.72763
-Λ = 0
+exact_ref = 4.0766890719636635 - 0.01275892774109674im
 m = 1.0
-ϕ = 0.1
+training_c = [2.0, 1.9, 1.8]
-μω_global = 0.5 * exp(-2im * ϕ)
+extrapolating_c = 0.0 : 0.2 : 1.2
 E_max = 40
-H0 = get_3b_H_matrix(jacobi, V_of_r, μω_global, E_max, Λ, m, true, true)
+include("../ho_basis_3body_resonance.jl")
 H0 = H
 # 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)
+@time "Vp" Vp =  get_src_V_matrix(Vp_of_r, basis, μω, μω_global)
-training_ref = -2.22
+exact = ComplexF64[]
 training = ComplexF64[]
 extrapolated = ComplexF64[]
 training_vecs = Vector{ComplexF64}[]
-extrapolating_ref = [4.076662025307587-0.012709842443350328im,
+current_E = training_ref
    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]
+for c in training_c
-extrapolating_c = 0.0 : 0.2 : 1.2
+    println("Training for c = $c")
    local H = H0 + c .* Vp
    local evals, evecs = eigs(H, nev=3, ncv=24, which=:LI, maxiter=5000, tol=1e-5, ritzvec=true, check=1)
-EC = affine_EC(H0, Vp)
+    global current_E = nearest(evals, current_E)
-train!(EC, training_c; ref_eval=training_ref, CAEC=true)
+    push!(training, current_E)
-extrapolate!(EC, extrapolating_c; ref_eval=extrapolating_ref)
+    push!(training_vecs, evecs[:, nearestIndex(evals, current_E)])
 end
-exportCSV(EC, "temp/HO_B2R.csv")
+# CA-EC
-plot(EC, "temp/HO_B2R.pdf")
+training_vecs = vcat(training_vecs, conj(training_vecs))
 EC_basis = hcat(training_vecs...)
 N_EC = transpose(EC_basis) * EC_basis
 H0_EC = transpose(EC_basis) * H0 * EC_basis
 Vp_EC = transpose(EC_basis) * Vp * EC_basis
 current_E = exact_ref
 for c in extrapolating_c
    println("Extrapolating for c = $c")
    local H = H0 + c .* Vp
    local evals, _ = eigs(H, nev=3, ncv=24, which=:LI, maxiter=5000, tol=1e-5, ritzvec=false, check=1)
    global current_E = nearest(evals, current_E)
    push!(exact, current_E)
    # extrapolation
    H_EC = H0_EC + c .* Vp_EC
    evals = eigvals(H_EC, N_EC)
    push!(extrapolated, nearest(evals, current_E))
 end
 exportCSV("temp/HO_B2R.csv", (training, exact, extrapolated), ("training", "exact", "extrapolated"))
 scatter(real.(training),imag.(training), label="training")
 scatter!(real.(exact),imag.(exact), label="exact")
 scatter!(real.(extrapolated),imag.(extrapolated), label="extrapolated")
 savefig("temp/HO_B2R.pdf")
--- a/calculations/3body_HO_R2R_EC.jl
+++ b/calculations/3body_HO_R2R_EC.jl
@ -1,24 +1,54 @@
-include("../ho_basis.jl")
+using DelimitedFiles, Plots
-include("../EC.jl")
+include("../ho_basis_3body_resonance.jl")
-V_of_r(r) = 2 * exp(-(r-3)^2 / (1.5)^2)
+current_E = 5.9673 - 0.0006im
 Λ = 0
 m = 1.0
 μω_global = 0.5 * exp(-2im * pi / 9)
 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)
 ref_E = 5.9673 - 0.0006im
 training_c = 2.0 : -0.2 : 1.2
 extrapolating_c = 1.05 .- [0.0 : 0.1 : 0.4; 0.45 : 0.05 : 0.60]
-EC = affine_EC(T, V)
+@time "H0" H0 = T1 + T2
 train!(EC, training_c; ref_eval=ref_E, CAEC=false)
 extrapolate!(EC, extrapolating_c)
-exportCSV(EC, "temp/HO_R2R.csv")
+# free memory
-plot(EC, "temp/HO_R2R.pdf")
+basis = T1 = T2 = V1_cache = V_relative_cache = V1 = V_relative = U = V2 = nothing
 GC.gc()
 exact = ComplexF64[]
 training = ComplexF64[]
 extrapolated = ComplexF64[]
 training_vecs = Vector{ComplexF64}[]
 for c in training_c
    println("Training for c = $c")
    local H = H0 + c .* V
    local evals, evecs = eigs(H, nev=3, ncv=24, which=:LI, maxiter=5000, tol=1e-5, ritzvec=true, check=1)
    global current_E = nearest(evals, current_E)
    push!(training, current_E)
    push!(training_vecs, evecs[:, nearestIndex(evals, current_E)])
 end
 EC_basis = hcat(training_vecs...)
 N_EC = transpose(EC_basis) * EC_basis
 H0_EC = transpose(EC_basis) * H0 * EC_basis
 V_EC = transpose(EC_basis) * V * EC_basis
 for c in extrapolating_c
    println("Extrapolating for c = $c")
    local H = H0 + c .* V
    local evals, evecs = eigs(H, nev=3, ncv=24, which=:LI, maxiter=5000, tol=1e-5, ritzvec=true, check=1)
    global current_E = nearest(evals, current_E)
    push!(exact, current_E)
    # extrapolation
    H_EC = H0_EC + c .* V_EC
    evals = eigvals(H_EC, N_EC)
    push!(extrapolated, nearest(evals, current_E))
 end
 exportCSV("temp/NCSM.csv", (training, exact, extrapolated), ("training", "exact", "extrapolated"))
 scatter(real.(training),imag.(training), label="training")
 scatter!(real.(exact),imag.(exact), label="exact")
 scatter!(real.(extrapolated),imag.(extrapolated), label="extrapolated")
 savefig("temp/NCSM.pdf")
--- a/calculations/3body_dis_CSM_EC.jl
+++ b/calculations/3body_dis_CSM_EC.jl
@ -1,77 +0,0 @@
 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")
--- a/calculations/3body_dis_HO_EC.jl
+++ b/calculations/3body_dis_HO_EC.jl
@ -1,5 +1,6 @@
 using Arpack, SparseArrays, LRUCache
 using DelimitedFiles, Plots
 include("../ho_basis.jl")
 include("../EC.jl")
 Λ = 0
 m = 1.0
@ -30,23 +31,62 @@ println("Basis size = ", basis.dim)
@time "Ha" Ha = T1 + T2 + Va
@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.55, -0.7, -0.85, -1, -1.2]
+# free memory
-extrapolating_c = [0.2, 0.1, 0.0, -0.1, -0.2, -0.3]
+basis = T1 = T2 = V1_cache = V_relative_cache = V1 = V_relative = U = V2 = nothing
 GC.gc()
-ref_E = -0.5173809356244544
+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]
-exact_ref = [-0.3136074661528041-0.07689284936868852im
+current_E = -0.5173809356244544
    -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)
+exact = ComplexF64[]
-train!(EC, training_c; ref_eval=ref_E, CAEC=true) # try CAEC=false !!!
+training = ComplexF64[]
-extrapolate!(EC, extrapolating_c, precalculated_exact_E = exact_ref)
+extrapolated = ComplexF64[]
 training_vecs = Vector{ComplexF64}[]
-exportCSV(EC, "temp/dis_HO_B2R.csv")
+for c in training_c
-plot(EC, "temp/dis_HO_B2R.pdf")
+    print("Training for c = $c: ")
    H = Ha + c .* Vb
    evals, evecs = eigs(H, nev=3, ncv=24, which=:SR, maxiter=5000, tol=1e-5, ritzvec=true, check=1)
    global current_E = nearest(evals, current_E)
    println(current_E)
    push!(training, current_E)
    push!(training_vecs, evecs[:, nearestIndex(evals, current_E)])
 end
 # CA-EC
 training_vecs = vcat(training_vecs, conj(training_vecs))
 EC_basis = hcat(training_vecs...)
 N_EC = transpose(EC_basis) * EC_basis
 Ha_EC = transpose(EC_basis) * Ha * EC_basis
 Vb_EC = transpose(EC_basis) * Vb * EC_basis
 current_E = -0.3005521915662689 - 0.13612069020686351im
 for c in extrapolating_c
    print("Extrapolating for c = $c: ")
    H = Ha + c .* Vb
    evals, evecs = eigs(H, nev=3, ncv=24, which=:SR, maxiter=5000, tol=1e-5, ritzvec=true, check=1)
    global current_E = nearest(evals, current_E)
    println(current_E)
    push!(exact, current_E)
    # extrapolation
    H_EC = Ha_EC + c .* Vb_EC
    evals = eigvals(H_EC, N_EC)
    push!(extrapolated, nearest(evals, current_E))
 end
 # exportCSV("temp/dis_HO_B2R.csv", (training, exact, extrapolated), ("training", "exact", "extrapolated"))
 scatter(real.(training),imag.(training), label="training")
 scatter!(real.(exact),imag.(exact), label="exact")
 scatter!(real.(extrapolated),imag.(extrapolated), label="extrapolated")
 savefig("temp/dis_HO_B2R.pdf")
--- a/calculations/ACCC.jl
+++ b/calculations/ACCC.jl
@ -1,50 +0,0 @@
 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)
--- a/calculations/B2R_Berggren_poles.jl
+++ b/calculations/B2R_Berggren_poles.jl
@ -1,12 +1,12 @@
-include("../EC.jl")
+using Plots
-include("../common.jl")
+include("../helper.jl")
 include("../p_space.jl")
 # contour
 p, w = get_mesh([0, 0.4 - 0.15im, 0.8, 6], [128, 128, 128])
-μ = 0.5
+# ResonanceEC: Eq. (20)
-V_system(c) = (p, q) -> c*(-5*g0(sqrt(3), p, q) + 2*g0(sqrt(10), p, q)) # ResonanceEC: Eq. (20)
+V_system(c) = (p, q) -> c*(-5*g0(sqrt(3), p, q) + 2*g0(sqrt(10), p, q))
 # generating a Berggren basis with a pole using the same system
 basis_c = 0.6
@ -16,18 +16,48 @@ N_berg = sqrt.(diag(transpose(berg_basis .* w) * berg_basis))
 berg_basis = berg_basis ./ transpose(N_berg)
 berg_basis_w = berg_basis .* w
-H0 = transpose(berg_basis_w) * get_T_matrix(p, μ) * berg_basis
+training_points = range(1.1, 0.9, 5) # original: range(1.35, 0.9, 5)
 V = transpose(berg_basis_w) * get_V_matrix(V_system(1), p, w) * berg_basis
-training_c = range(1.1, 0.9, 5) # original: range(1.35, 0.9, 5)
+training_E = Vector{ComplexF64}(undef, length(training_points))
-extrapolating_c = range(0.78, 0.45, 7) # original: range(0.75, 0.40, 8)
+EC_basis = Matrix{ComplexF64}(undef, length(p), length(training_points))
-training_ref = -0.26
+# training
-extrapolating_ref = [quick_pole_E(V_system(c)) for c in extrapolating_c]
+for (j, c) in enumerate(training_points)
    H = get_H_matrix(V_system(c), p, w)
    H_berg = transpose(berg_basis_w) * H * berg_basis
    evals, evecs = eigen(H_berg)
    i = argmin(real.(evals))
 #    i = identify_pole_i(basis_p, evals)
    training_E[j] = evals[i]
    EC_basis[:, j] = evecs[:, i]
 end
-EC = affine_EC(H0, V)
+EC_basis = hcat(EC_basis, conj.(EC_basis)) # CA-EC
-train!(EC, training_c; ref_eval=training_ref, CAEC=true)
+N_EC = transpose(EC_basis) * EC_basis
 extrapolate!(EC, extrapolating_c; ref_eval=extrapolating_ref)
-exportCSV(EC, "temp/2b_GSM_B2R.csv")
+extrapolate_points = range(0.78, 0.45, 7) # original: range(0.75, 0.40, 8)
-plot(EC, "temp/2b_GSM_B2R.pdf"; basis_points=basis_E, xlims=(0, 0.3), ylims=(-0.120, 0.020))
+
 exact_E = Vector{ComplexF64}(undef, length(extrapolate_points))
 extrapolate_E = Vector{ComplexF64}(undef, length(extrapolate_points))
 # extrapolating
 for (j, c) in enumerate(extrapolate_points)
    exact_E[j] = quick_pole_E(V_system(c))
    H = get_H_matrix(V_system(c), p, w)
    H_berg = transpose(berg_basis_w) * H * berg_basis
    H_EC = transpose(EC_basis) * H_berg * EC_basis
    evals = eigvals(H_EC, N_EC)
    i = argmin(abs.(evals .- exact_E[j]))
    extrapolate_E[j] = evals[i]
 end
 exportCSV("temp/2b_GSM_B2R.csv", (training_E, exact_E, extrapolate_E), ("training", "exact", "extrapolated"))
 scatter(real.(training_E), imag.(training_E), label="training")
 scatter!(real.(exact_E), imag.(exact_E), label="exact")
 scatter!(real.(extrapolate_E), imag.(extrapolate_E), label="extrapolated")
 scatter!(real.(basis_E), imag.(basis_E), m=:x, label="Berggren basis")
 xlims!(0,0.3)
 ylims!(-0.120,0.020)
 savefig("temp/2b_GSM_B2R.pdf")
--- a/calculations/B2R_comparison.jl
+++ b/calculations/B2R_comparison.jl
@ -1,30 +1,51 @@
-include("../EC.jl")
+using Plots
 include("../common.jl")
 include("../p_space.jl")
 berggren_mesh = get_mesh([0, 0.4 - 0.15im, 0.8, 6], [128, 128, 128])
 csm_mesh = get_mesh([0, 8 - 3im], [512])
-for (mesh, name) in zip((berggren_mesh, csm_mesh), ("beggren", "csm"))
+for mesh in (berggren_mesh, csm_mesh)
    p, w = mesh
    mesh_E = p.*p ./ (2*0.5)
-    μ = 0.5
+    # ResonanceEC: Eq. (20)
-    V_system(c) = (p, q) -> c*(-5*g0(sqrt(3), p, q) + 2*g0(sqrt(10), p, q)) # ResonanceEC: Eq. (20)
+    V_system(c) = (p, q) -> c*(-5*g0(sqrt(3), p, q) + 2*g0(sqrt(10), p, q))
    H0 = get_T_matrix(p, μ)
    V = get_V_matrix(V_system(1), p, w)
    training_c = range(1.1, 0.9, 5) # original: range(1.35, 0.9, 5)
    extrapolating_c = range(0.78, 0.45, 7) # original: range(0.75, 0.40, 8)
-    training_ref = [quick_pole_E(V_system(c)) for c in training_c]
+    training_points = range(1.1, 0.9, 5) # original: range(1.35, 0.9, 5)
-    exact_E = [quick_pole_E(V_system(c)) for c in extrapolating_c]
+    training_E = Vector{ComplexF64}(undef, length(training_points))
    EC_basis = Matrix{ComplexF64}(undef, length(p), length(training_points))
-    EC = affine_EC(H0, V, w)
+    for (j, c) in enumerate(training_points)
-    train!(EC, training_c; ref_eval=training_ref, CAEC=true)
+        evals, evecs = eigen(get_H_matrix(V_system(c), p, w))
-    extrapolate!(EC, extrapolating_c; precalculated_exact_E=exact_E)
+        i = identify_pole_i(p, evals)
-    
+        training_E[j] = evals[i]
-    #exportCSV(EC, "temp/2b_comparison_$name.csv")
+        EC_basis[:, j] = evecs[:, i]
-    plot(EC, "temp/2b_comparison_$name.pdf"; basis_contour=mesh_E, xlims=(-0.3,0.3), ylims=(-0.120,0.020))
+    end
    EC_basis = hcat(EC_basis, conj.(EC_basis)) # CA-EC
    EC_basis_w = EC_basis .* w
    N_EC = transpose(EC_basis_w) * EC_basis
    extrapolate_points = range(0.78, 0.45, 7) # original: range(0.75, 0.40, 8)
    exact_E = Vector{ComplexF64}(undef, length(extrapolate_points))
    extrapolate_E = Vector{ComplexF64}(undef, length(extrapolate_points))
    for (j, c) in enumerate(extrapolate_points)
        exact_E[j] = quick_pole_E(V_system(c))
        H = get_H_matrix(V_system(c), p, w)
        H_EC = transpose(EC_basis_w) * H * EC_basis
        evals = eigvals(H_EC, N_EC)
        i = argmin(abs.(evals .- exact_E[j]))
        extrapolate_E[j] = evals[i]
    end
    scatter(real.(training_E), imag.(training_E), label="training")
    scatter!(real.(exact_E), imag.(exact_E), label="exact")
    scatter!(real.(extrapolate_E), imag.(extrapolate_E), label="extrapolated")
    plot!(real.(mesh_E), imag.(mesh_E), label="contour")
    xlims!(-0.3,0.3)
    ylims!(-0.120,0.020)
    savefig("temp/" * string(rand(UInt16)) * ".pdf")
 end
--- a/calculations/PMM.py
+++ b/calculations/PMM.py
@ -1,81 +0,0 @@
 #%%
 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()
 # %%
--- a/calculations/R2R_Berggren.jl
+++ b/calculations/R2R_Berggren.jl
@ -1,5 +1,4 @@
-include("../EC.jl")
+using Plots
 include("../common.jl")
 include("../p_space.jl")
 # contour
@ -8,21 +7,42 @@ subdivisions = [128, 128, 128]
 p, w = get_mesh(vertices, subdivisions)
 mesh_E = p.*p ./ (2*0.5)
-μ = 0.5
+# ResonanceEC: Eq. (20)
-V_system(c) = (p, q) -> c*(-5*g0(sqrt(3), p, q) + 2*g0(sqrt(10), p, q)) # ResonanceEC: Eq. (20)
+V_system(c) = (p, q) -> c*(-5*g0(sqrt(3), p, q) + 2*g0(sqrt(10), p, q))
-H0 = get_T_matrix(p, μ)
+training_points = range(0.75, 0.45, 5)
-V = get_V_matrix(V_system(1), p, w)
+training_E = Vector{ComplexF64}(undef, length(training_points))
 EC_basis = Matrix{ComplexF64}(undef, length(p), length(training_points))
-training_c = range(0.75, 0.45, 5)
+for (j, c) in enumerate(training_points)
-extrapolating_c = range(0.40, 0.25, 5)
+    evals, evecs = eigen(get_H_matrix(V_system(c), p, w))
    i = identify_pole_i(p, evals)
    training_E[j] = evals[i]
    EC_basis[:, j] = evecs[:, i]
 end
-training_ref = [quick_pole_E(V_system(c)) for c in training_c]
+extrapolate_points = range(0.40, 0.25, 5)
-exact_E = [quick_pole_E(V_system(c)) for c in extrapolating_c]
+ref_E = 0.2 - 0.1im
-EC = affine_EC(H0, V, w)
+exact_E = Vector{ComplexF64}(undef, length(extrapolate_points))
-train!(EC, training_c; ref_eval=training_ref, CAEC=false)
+extrapolate_E = Vector{ComplexF64}(undef, length(extrapolate_points))
 extrapolate!(EC, extrapolating_c; precalculated_exact_E=exact_E)
-#exportCSV(EC, "temp/2b_R2R.csv")
+EC_basis_w = EC_basis .* w
-plot(EC, "temp/2b_R2R.pdf"; basis_contour=mesh_E, xlims=(0, 1))
+N_EC = transpose(EC_basis_w) * EC_basis
 for (j, c) in enumerate(extrapolate_points)
    exact_E[j] = quick_pole_E(V_system(c))
    H = get_H_matrix(V_system(c), p, w)
    H_EC = transpose(EC_basis_w) * H * EC_basis
    evals = eigvals(H_EC, N_EC)
    i = argmin(abs.(evals .- ref_E))
    global ref_E = evals[i]
    extrapolate_E[j] = evals[i]
 end
 scatter(real.(training_E), imag.(training_E), label="training")
 scatter!(real.(exact_E), imag.(exact_E), label="exact")
 scatter!(real.(extrapolate_E), imag.(extrapolate_E), label="extrapolated")
 plot!(real.(mesh_E), imag.(mesh_E), label="contour")
 xlims!(0,1)
--- a/calculations/R2R_Berggren_poles.jl
+++ b/calculations/R2R_Berggren_poles.jl
@ -1,33 +1,66 @@
-include("../EC.jl")
+using Plots
-include("../common.jl")
+include("../helper.jl")
 include("../p_space.jl")
 # contour
-p, w = get_mesh([0, 0.4 - 0.15im, 0.8, 6], [128, 128, 128])
+p, w = get_mesh([0, 0.4 - 0.08im, 0.8, 6], [128, 128, 128])
 contour_E = p.^2
-μ = 0.5
+# ResonanceEC: Eq. (20)
-V_system(c) = (p, q) -> c*(-5*g0(sqrt(3), p, q) + 2*g0(sqrt(10), p, q)) # ResonanceEC: Eq. (20)
+V_system(c) = (p, q) -> c*(-5*g0(sqrt(3), p, q) + 2*g0(sqrt(10), p, q))
 # generating a Berggren basis with a pole using the same system
 basis_c = 0.6
 basis_E, berg_basis = eigen(get_H_matrix(V_system(basis_c), p, w); permute=false, scale=false)
 pole_E = quick_pole_E(V_system(basis_c)) # basis only has 1 pole
 basis_p = sqrt.(basis_E)
 N_berg = sqrt.(diag(transpose(berg_basis .* w) * berg_basis))
 berg_basis = berg_basis ./ transpose(N_berg)
 berg_basis_w = berg_basis .* w
-H0 = transpose(berg_basis_w) * get_T_matrix(p, μ) * berg_basis
+training_points = range(0.79, 0.66, 4) # original: range(1.35, 0.9, 5)
 V = transpose(berg_basis_w) * get_V_matrix(V_system(1), p, w) * berg_basis
-training_c = range(0.79, 0.66, 4) # original: range(1.35, 0.9, 5)
+training_E = Vector{ComplexF64}(undef, length(training_points))
-extrapolating_c = range(0.62, 0.40, 6) # original: range(0.75, 0.40, 8)
+EC_basis = Matrix{ComplexF64}(undef, length(p), length(training_points))
-training_ref = [quick_pole_E(V_system(c)) for c in training_c]
+# training
-extrapolating_ref = [quick_pole_E(V_system(c)) for c in extrapolating_c]
+for (j, c) in enumerate(training_points)
    H = get_H_matrix(V_system(c), p, w)
    H_berg = transpose(berg_basis_w) * H * berg_basis
    evals, evecs = eigen(H_berg)
    i = identify_pole_i(basis_p, evals)
    training_E[j] = evals[i]
    EC_basis[:, j] = evecs[:, i]
 end
-EC = affine_EC(H0, V)
+N_EC = transpose(EC_basis) * EC_basis
 train!(EC, training_c; ref_eval=training_ref, CAEC=false)
 extrapolate!(EC, extrapolating_c; ref_eval=extrapolating_ref)
-exportCSV(EC, "temp/2b_GSM_R2R.csv")
+extrapolate_points = range(0.62, 0.40, 6) # original: range(0.75, 0.40, 8)
-plot(EC, "temp/2b_GSM_R2R.pdf"; basis_points=basis_E, xlims=(0, 0.3), ylims=(-0.120, 0.020))
+
 exact_E = Vector{ComplexF64}(undef, length(extrapolate_points))
 extrapolate_E = Vector{ComplexF64}(undef, length(extrapolate_points))
 # extrapolating
 for (j, c) in enumerate(extrapolate_points)
    exact_E[j] = quick_pole_E(V_system(c))
    H = get_H_matrix(V_system(c), p, w)
    H_berg = transpose(berg_basis_w) * H * berg_basis
    H_EC = transpose(EC_basis) * H_berg * EC_basis
    evals = eigvals(H_EC, N_EC)
    i = argmin(abs.(evals .- exact_E[j]))
    extrapolate_E[j] = evals[i]
 end
 exportCSV("temp/2b_GSM_R2R.csv", (training_E, exact_E, extrapolate_E, [pole_E]), ("training", "exact", "extrapolated", "basis"))
 contour_E_export = contour_E[real.(contour_E) .< 1] # to trim unnecessary data outside axis limits
 exportCSV("temp/2b_GSM_R2R_contour.csv", contour_E_export)
 scatter(real.(training_E), imag.(training_E), label="training")
 scatter!(real.(exact_E), imag.(exact_E), label="exact")
 scatter!(real.(extrapolate_E), imag.(extrapolate_E), label="extrapolated")
 scatter!(real.(basis_E), imag.(basis_E), m=:x, label="Berggren basis")
 xlims!(0,0.3)
 ylims!(-0.120,0.020)
 savefig("temp/2b_GSM_R2R.pdf")
--- a/calculations/XZ/2body_HO.jl
+++ b/calculations/XZ/2body_HO.jl
@ -1,39 +0,0 @@
 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")
--- a/calculations/XZ/2body_p_space.jl
+++ b/calculations/XZ/2body_p_space.jl
@ -1,50 +0,0 @@
 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")
--- a/calculations/XZ/3body_HO.jl
+++ b/calculations/XZ/3body_HO.jl
@ -1,36 +0,0 @@
 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")
--- a/calculations/XZ/3body_p_space.jl
+++ b/calculations/XZ/3body_p_space.jl
@ -1,45 +0,0 @@
 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.
--- a/calculations/XZ/p_space_vs_HO.jl
+++ b/calculations/XZ/p_space_vs_HO.jl
@ -1,71 +0,0 @@
 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")
--- a/calculations/XZ_technique.jl
+++ b/calculations/XZ_technique.jl
@ -0,0 +1,55 @@
 using LinearAlgebra, Plots
 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)
 train_E = zeros(ComplexF64, n_EC)
 EC_basis = zeros(ComplexF64, (n_max + 1, length(train_cs)))
 exact_E = zeros(ComplexF64, length(target_cs))
 extrapolate_E = similar(exact_E)
 near_E = 0.2 + 0.2im
 for (j, c) in enumerate(train_cs)
    H = T + c .* V
    evals, evecs = eigen(H)
    i = argmin(abs.(evals .- near_E))
    train_E[j] = evals[i]
    EC_basis[:, j] = evecs[:, i]
 end
 N_EC = transpose(EC_basis) * EC_basis
 for (j, c) in enumerate(target_cs)
    exact_E[j] = quick_pole_E((p, q) -> c*(V1*g0(R1, p, q) + V2*g0(R2, p, q)), μ; cs_angle=0.5)
    H = T + c .* V
    H_EC = transpose(EC_basis) * H * EC_basis
    evals = eigvals(H_EC, N_EC)
    i = argmin(abs.(evals .- exact_E[j]))
    extrapolate_E[j] = evals[i]
 end
 scatter(real.(train_E), imag.(train_E), label="training")
 scatter!(real.(exact_E), imag.(exact_E), label="exact")
 scatter!(real.(extrapolate_E), imag.(extrapolate_E), label="extrapolated")
 xlims!(-0.2,0.3)
 ylims!(-0.3,0.3)
--- a/common.jl
+++ b/common.jl
@ -1,106 +0,0 @@
 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]
    neg_arr = arr[arr .< 0]
    sort!(pos_arr)
    sort!(neg_arr, rev=true)
    return sum(pos_arr) + sum(neg_arr)
 end
 better_sum(arr::Array{ComplexF64}) = better_sum(real.(arr)) + 1im * better_sum(imag.(arr))
 "The triangle inequality for angular momenta"
 triangle_ineq(l1, l2, L) = abs(l1 - l2) ≤ L && L ≤ (l1 + l2)
 "Index of the nearest value in a list to a given reference point"
 nearestIndex(list::Array, ref) = argmin(norm.(list .- ref))
 "Nearest value in a list to a given reference point"
 nearest(list::Array, ref) = list[nearestIndex(list, ref)]
 "Simple implementation of the Kronecker sum"
 function kron_sum(A::AbstractMatrix, B::AbstractMatrix)
    @assert size(A, 1) == size(A, 2) && size(B, 1) == size(B, 2) "Matrices should be square"
    return kron(A, I(size(B, 1))) + kron(I(size(A, 1)), B)
 end
 "Flattened vector version of Iterators.product(...) with index hierachy reversed -- leftmost index has the highest hierachy"
 iter_prod(args...) = reverse.(collect(Iterators.product(reverse(args)...))[:])
 "Export CSV data for a scatter plot taking in data as a list of complex vectors (x=Re, y=Im), and a list of corresponding labels, typically used for EC results"
 function exportCSV(file::String, data, labels=nothing)
    columns = ["x" "y" "label"]
    if isnothing(labels)
        columns[end] = ""
        labels = fill("", length(data))
    end
    open(file, "w") do f
        writedlm(f, columns)
        for (d, label) in zip(data, labels)
            l = fill(label, length(d))
            writedlm(f, hcat(reim(d)..., l))
        end
    end
 end
 "In-place c-orthogonalization via (modified) Gram-Schmidt. Only significant vectors are returned (c-normalized).
 The number of significant vectors to return are determined by the original singular values (compared to the threshold)."
 function gram_schmidt!(vecs::Vector{Vector{T}}, ws=ones(length(vecs[1])), threshold=1e-5; verbose=false) where T<: Number
    c_product(i, j) = sum(vecs[i] .* ws .* vecs[j])
    norm(i) = c_product(i, i)
    proj(i, j) = (c_product(i, j) / norm(i)) .* vecs[i] # component of vec[i] in vec[j]
    # initial normalization
    for i in eachindex(vecs)
        vecs[i] ./= sqrt(norm(i))
    end
    verbose && println("Absolute singular values = $(round.(c_singular_values(vecs, ws); sigdigits=1))")
    target_dim = c_rank(vecs, ws, threshold)
    verbose && println("Target dimensionality = $target_dim")
    selected_vecs_i = Integer[]
    while length(selected_vecs_i) < target_dim
        i = argmax(i -> abs(norm(i)), setdiff(eachindex(vecs), selected_vecs_i)) # find the largest vector from the remaining
        push!(selected_vecs_i, i)
        for j in setdiff(eachindex(vecs), selected_vecs_i)
            vecs[j] .-= proj(i, j)
        end
    end
    verbose && println("Absolute norms of selected vectors = $(round.(selected_vecs_i .|> norm .|> abs; sigdigits=1))")
    verbose && println("Absolute norms of dropped vectors = $(round.(setdiff(eachindex(vecs), selected_vecs_i) .|> norm .|> abs; sigdigits=1))")
    # final normalization
    for i in selected_vecs_i
        vecs[i] ./= sqrt(norm(i))
    end
    return vecs[selected_vecs_i]
 end
 "Same as above but the basis is provided as a matrix instead of an array of vectors"
 gram_schmidt!(vecs::Matrix{T}, ws=ones(size(vecs, 1)), threshold=1e-5; verbose=false) where T<: Number = hcat(gram_schmidt!([vecs[:, i] for i in axes(vecs, 2)], ws, threshold; verbose=verbose)...)
 "Rank of a basis set (pre-normalized) under c-product"
 c_rank(vecs, ws, threshold=1e-8) = count(c_singular_values(vecs, ws) .> threshold)
 "Singular values (magnitudes) of a basis set (pre-normalized) under c-product"
 function c_singular_values(vecs, ws)
    basis = hcat(vecs...)
    N = transpose(basis) * spdiagm(ws) * basis
    singular_values = eigvals(N) .|> abs .|> sqrt
    return singular_values
 end
--- a/helper.jl
+++ b/helper.jl
@ -0,0 +1,50 @@
 using LinearAlgebra, DelimitedFiles
 "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]
    neg_arr = arr[arr .< 0]
    sort!(pos_arr)
    sort!(neg_arr, rev=true)
    return sum(pos_arr) + sum(neg_arr)
 end
 better_sum(arr::Array{ComplexF64}) = better_sum(real.(arr)) + 1im * better_sum(imag.(arr))
 "The triangle inequality for angular momenta"
 triangle_ineq(l1, l2, L) = abs(l1 - l2) ≤ L && L ≤ (l1 + l2)
 "Index of the nearest value in a list to a given reference point"
 nearestIndex(list::Array, ref) = argmin(norm.(list .- ref))
 "Nearest value in a list to a given reference point"
 nearest(list::Array, ref) = list[nearestIndex(list, ref)]
 "Simple implementation of the Kronecker sum"
 function kron_sum(A::AbstractMatrix, B::AbstractMatrix)
    @assert size(A, 1) == size(A, 2) && size(B, 1) == size(B, 2) "Matrices should be square"
    return kron(A, I(size(B, 1))) + kron(I(size(A, 1)), B)
 end
 "Flattened vector version of Iterators.product(...) with index hierachy reversed -- leftmost index has the highest hierachy"
 iter_prod(args...) = reverse.(collect(Iterators.product(reverse(args)...))[:])
 "Export CSV data for a scatter plot taking in data as a list of complex vectors (x=Re, y=Im), and a list of corresponding labels, typically used for EC results"
 function exportCSV(file::String, data, labels=nothing)
    columns = ["x" "y" "label"]
    if isnothing(labels)
        columns[end] = ""
        labels = fill("", length(data))
    end
    open(file, "w") do f
        writedlm(f, columns)
        for (d, label) in zip(data, labels)
            l = fill(label, length(d))
            writedlm(f, hcat(reim(d)..., l))
        end
    end
 end
--- a/ho_basis.jl
+++ b/ho_basis.jl
@ -1,10 +1,36 @@
 using SparseArrays
 using QuadGK
 using LRUCache
-include("common.jl")
+include("helper.jl")
 include("math.jl")
-"2-body HO basis (used for 3-body systems without the CM DOF)"
+"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
@ -155,7 +181,7 @@ function get_jacobi_V2_matrix(V_of_r, basis::ho_basis_2B, μω_global; atol=10^-
    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
+    V2 =  U' * V_relative * U
    return V2
 end
@ -191,7 +217,7 @@ function get_src_V12_matrix(V_of_r, basis::ho_basis_2B, μω_global; atol=10^-6,
    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
+    V12 =  U' * V_relative * U
    return V12
 end
@ -210,42 +236,3 @@ function get_W_matrix(basis_p, basis::ho_basis_2B, μ1ω1, μ2ω2=μ1ω1; weight
    end
    return sparse(W)
 end
 function get_3b_H_matrix(coord_system::coordinate_system, V_of_r, μω_global, E_max, Λ, m=1.0, kinetic_part=true, potential_part=true; 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)
    out = spzeros(basis.dim, basis.dim)
    if kinetic_part
        verbose && println("Constructing KE matrices")
        @time "T1" out += get_sp_T_matrix(basis.n1s, basis.l1s, [basis.n2s, basis.l2s]; μω_gen=μ1ω1, μ=μ1)
        @time "T2" out += get_sp_T_matrix(basis.n2s, basis.l2s, [basis.n1s, basis.l1s]; μω_gen=μ2ω2, μ=μ2)
        if coord_system == src
            @time "T_cross" out += get_2p_p1p2_matrix(basis, μ1ω1, μ2ω2; dtype=ComplexF64) ./ (2*μ)
        end
    end
    if potential_part
        verbose && println("Constructing PE matrices")
        if coord_system == jacobi
            @time "V" out += get_jacobi_V_matrix(V_of_r, basis, μ1ω1, μω_global; atol=atol, maxevals=maxevals)
        elseif coord_system == src
            @time "V" out += get_src_V_matrix(V_of_r, basis, μω, μω_global; atol=atol, maxevals=maxevals)
        end
    end
    return out
 end
--- a/ho_basis_2body.jl
+++ b/ho_basis_2body.jl
@ -0,0 +1,30 @@
 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)
--- a/ho_basis_3body.jl
+++ b/ho_basis_3body.jl
@ -1,14 +1,40 @@
-using Arpack
+using Arpack, SparseArrays
 include("ho_basis.jl")
 Λ = 0
 m = 1.0
-V_of_r(r) = -2 * exp(-r^2 / 4)
+Va = -2
 Ra = 2
 E_max = 40
 μω_global = 0.3
-H = get_3b_H_matrix(src, V_of_r, μω_global, E_max, Λ, m)
+# 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, μω, μω) ./ (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 =  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)
--- a/ho_basis_3body_resonance.jl
+++ b/ho_basis_3body_resonance.jl
@ -1,15 +1,35 @@
-using Arpack
+using Arpack, SparseArrays, LRUCache
 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)
 μω_global = 0.5 * exp(-2im * pi / 9)
 E_max = 30
 μω_global = 0.5 * exp(-2im * pi / 9)
-H = get_3b_H_matrix(src, V_of_r, μω_global, E_max, Λ, m)
+# 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)
@time "Eigenvalues" evals, _ = eigs(H, nev=5, ncv=50, sigma=target_E, maxiter=5000, tol=1e-5, ritzvec=false, check=1)
 display(evals)
--- a/ho_basis_3body_wf.jl
+++ b/ho_basis_3body_wf.jl
@ -0,0 +1,30 @@
 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
--- a/test/ho_basis_test.jl
+++ b/test/ho_basis_test.jl
@ -1,5 +1,5 @@
 using LinearAlgebra
-include("../ho_basis.jl")
+include("ho_basis.jl")
 l = 0
 V0 = -10
--- a/math.jl
+++ b/math.jl
@ -1,5 +1,5 @@
 using SpecialFunctions, WignerSymbols
-include("common.jl")
+include("helper.jl")
 # Gaussian potentials in HO space
 inv_factorial(n) = Iterators.prod(inv.(1:n))
--- a/p_space.jl
+++ b/p_space.jl
@ -1,6 +1,5 @@
 using LinearAlgebra
 using SpecialFunctions, FastGaussQuadrature, QuadGK
 include("ho_basis.jl")
 function gausslegendre_shifted(a, b, n)
    scale = (b - a) / 2
@ -60,56 +59,3 @@ function Vl_mat_elem(V_of_r, l, p, q; atol=0, maxevals=10^7, R_cutoff=Inf)
    (integral, _) = quadgk(integrand, 0, R_cutoff; atol=atol, maxevals=maxevals)
    return (2 / pi) * integral
 end
 "Return the Hamiltonian matrix (and the array of weights) for the given system."
 function get_3b_H_matrix(coord_system::coordinate_system, V_of_r, vertices, subdivisions, jmax, μω_global, E_max, Λ, m=1.0, kinetic_part=true, potential_part=true; atol=10^-5, maxevals=10^5, R_cutoff=16, verbose=true)
    if coord_system == jacobi
        μ1ω1 = μω_global * 1/2
        μ2ω2 = μω_global * 2
        μ1 = m * 1/2
        μ2 = m * 2/3
    else
        error("Only Jacobi coordinates are implemented")
    end
    verbose && println("No of threads = ", Threads.nthreads())
    V_l(j, k, kp) = Vl_mat_elem(V_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")
    out = spzeros(basis_size, basis_size)
    @time "Block diagonal part" begin
        if kinetic_part & potential_part
            Hb_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]
        elseif kinetic_part
            Hb_blocks = [kron_sum(get_T_matrix(ks, μ1), get_T_matrix(ks, μ2)) for _ in js]
        elseif potential_part
            Hb_blocks = [kron_sum(get_V_matrix((k, kp) -> V_l(j1, k, kp), ks, ws), spzeros(block_size, block_size)) for (j1, _) in js]
        end
        out += blockdiag(sparse.(Hb_blocks)...)
    end
    if potential_part
        basis_ho = ho_basis_2B(E_max, Λ)
        verbose && println("HO basis size = ", basis_ho.dim)
        @time "V2_HO" V2_HO =  get_jacobi_V2_matrix(V_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" out += W_left * V2_HO * transpose(W_right)
    end
    return (out, weights)
 end
--- a/p_space_2body.jl
+++ b/p_space_2body.jl
@ -1,5 +1,5 @@
 using SparseArrays, Arpack
-include("common.jl")
+include("helper.jl")
 include("p_space.jl")
 E_target = -0.3919
--- a/p_space_3body.jl
+++ b/p_space_3body.jl
@ -1,18 +1,56 @@
-using Arpack
+using LinearAlgebra, SparseArrays, Arpack
 include("helper.jl")
 include("p_space.jl")
 include("ho_basis.jl")
 println("No of threads = ", Threads.nthreads())
 atol = 10^-5
 maxevals = 10^5
 R_cutoff = 16
 Λ = 0
 m = 1.0
-V_of_r(r) = -2 * exp(-r^2 / 4)
+μ1 = m * 1/2
 μ2 = m * 2/3
 Va = -2
 Ra = 2
 V_of_r(r) = Va * exp(-r^2 / Ra^2)
 V_l(j, k, kp) = Vl_mat_elem(V_of_r, j, k, kp; atol=atol, maxevals=maxevals, R_cutoff=R_cutoff)
 vertices = [0, 0.5 - 0.3im, 1, 4]
 subdivisions = [10, 10, 10]
 ks, ws = get_mesh(vertices, subdivisions)
 jmax = 4
 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"
 println("Basis size = $basis_size")
@time "Block diagonal part" begin
    Hb_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]
    Hb = blockdiag(sparse.(Hb_blocks)...)
 end
 E_max = 30
 μω_global = 0.5
 μ1ω1 = μω_global * 1/2
 μ2ω2 = μω_global * 2
-H, _ = get_3b_H_matrix(jacobi, V_of_r, vertices, subdivisions, jmax, μω_global, E_max, Λ, m)
+basis_ho = ho_basis_2B(E_max, Λ)
@time "V2_HO" V2_HO =  get_jacobi_V2_matrix(V_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" V2 = W_left * V2_HO * transpose(W_right)
@time "H" H = Hb + V2
@time "Eigenvalues" evals, _ = eigs(H, nev=3, ncv=24, which=:SR, maxiter=5000, tol=1e-5, ritzvec=false, check=1)
 display(evals)
--- a/p_space_3body_resonance.jl
+++ b/p_space_3body_resonance.jl
@ -1,20 +1,55 @@
-using Arpack
+using LinearAlgebra, SparseArrays, Arpack
 include("helper.jl")
 include("p_space.jl")
 include("ho_basis.jl")
-target = 4.0766890719636875 - 0.012758927741074495im
+println("No of threads = ", Threads.nthreads())
 atol = 10^-5
 maxevals = 10^5
 R_cutoff = 16
 Λ = 0
 m = 1.0
 μ1 = m * 1/2
 μ2 = m * 2/3
 target = 4.0766890719636875 - 0.012758927741074495im
 V_of_r(r) =  2 * exp(-(r-3)^2 / (1.5)^2)
 V_l(j, k, kp) = Vl_mat_elem(V_of_r, j, k, kp; atol=atol, maxevals=maxevals, R_cutoff=R_cutoff)
 vertices = [0, 2 - 0.2im, 3, 4]
 subdivisions = [15, 10, 10]
 ks, ws = get_mesh(vertices, subdivisions)
 jmax = 4
 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"
 println("Basis size = $basis_size")
@time "Block diagonal part" begin
    Hb_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]
    Hb = blockdiag(sparse.(Hb_blocks)...)
 end
 E_max = 40
 μω_global = 0.5
 μ1ω1 = μω_global * 1/2
 μ2ω2 = μω_global * 2
-H, _ = get_3b_H_matrix(jacobi, V_of_r, vertices, subdivisions, jmax, μω_global, E_max, Λ, m)
+basis_ho = ho_basis_2B(E_max, Λ)
@time "V2_HO" V2_HO =  get_jacobi_V2_matrix(V_of_r, basis_ho, μω_global; atol=atol, maxevals=maxevals)
@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" V2 = W_left * V2_HO * transpose(W_right)
@time "H" H = Hb + V2
@time "Eigenvalues" evals, _ = eigs(H, sigma=target, maxiter=5000, tol=1e-5, ritzvec=false, check=1)
 display(evals)
--- a/test/misc.jl
+++ b/test/misc.jl
@ -1,65 +0,0 @@
 using SparseArrays, LinearAlgebra
 include("../common.jl")
 #### gram_schmidt ####
 n = 64
 N = 8
 vecs = [rand(n) + 1im .* rand(n) for _ in 1:N]
 ws = rand(n)
 ws_mat = spdiagm(ws)
 basis = hcat(vecs...)
 H = rand(n, n)
 H += transpose(H) # complex symmetric
 H_EC = transpose(basis) * ws_mat * H * ws_mat * basis
 N_EC = transpose(basis) * ws_mat * basis
 evals = eigvals(H_EC, N_EC)
 println("Eigenvalues with GEVP:")
 display(evals)
 ortho_basis = hcat(gram_schmidt!(vecs, ws)...)
 N_ortho = transpose(ortho_basis) * ws_mat * ortho_basis
 println("Norm matrix after Gram-Schmidt:")
 display(round.(N_ortho; digits=2)) # should be ≈I
@assert N_ortho ≈ I(N) "Gram-Schmidt did not yield an orthogonal basis"
 H_EC_ortho = transpose(ortho_basis) * ws_mat * H * ws_mat * ortho_basis
 evals_ortho = eigvals(H_EC_ortho)
 println("Eigenvalues after Gram-Schmidt:")
 display(evals_ortho)
@assert evals ≈ evals_ortho "Gram-Schmidt did not preserve the eigenvalues"
 ############
 println("\nRepeat with a redundant basis\n")
 println("Original dimensionality = $(length(vecs))")
 for pow in 6:9
    noise = rand(n) ./ 10^pow
    new_vec = vecs[1] + noise
    push!(vecs, new_vec)
 end
 println("Dimensionality before Gram-Schmidt = $(length(vecs))")
 ortho_vecs = gram_schmidt!(vecs, ws; verbose=true)
 ortho_basis = hcat(ortho_vecs...)
 println("Dimensionality after Gram-Schmidt = $(length(ortho_vecs))")
 H_EC_ortho = transpose(ortho_basis) * ws_mat * H * ws_mat * ortho_basis
 evals_ortho = eigvals(H_EC_ortho)
 println("Eigenvalues after Gram-Schmidt:")
 display(evals_ortho)
@assert isapprox(evals, evals_ortho; atol=1e-3) "Gram-Schmidt did not approximately preserve the eigenvalues"
 ######################
--- a/test/moshinsky.jl
+++ b/test/moshinsky.jl
@ -1,4 +1,4 @@
-println("### Test: transpose(U) * U == identity")
+println("### Test: U' * U == identity")
 using LinearAlgebra
@ -15,7 +15,7 @@ println("Basis size = ", basis.dim)
@time "Moshinsky brackets" U = Moshinsky_transform(basis)
-check = transpose(U) * U - I
+check = U' * U - I
 println("Maximum = ", maximum(abs.(check)))
 println("Norm = ", sum(check .* conj(check)))
--- a/test/p_space_basic.jl
+++ b/test/p_space_basic.jl
@ -1,5 +1,5 @@
 using Plots
-include("../common.jl")
+include("../helper.jl")
 include("../p_space.jl")
 vertices = [0, 1 - 0.5im, 2, 6]
--- a/test/simple_berggren.jl
+++ b/test/simple_berggren.jl
@ -1,5 +1,5 @@
 using Plots
-include("../common.jl")
+include("../helper.jl")
 include("../p_space.jl")
 vertices = [0, 0.5 - 0.3im, 1, 6]