Sqrt fix (still not working)

EC.jl

@@ -156,11 +156,11 @@ 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", markercolor=:white)
+    scatter!(real.(EC.exact_E), imag.(EC.exact_E), label="exact")
     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)
+        scatter!(real.(EC.extrapolated_E), imag.(EC.extrapolated_E), xerror=real.(EC.extrapolated_CI), yerror=imag.(EC.extrapolated_CI), label="extrapolated")
     else
-        scatter!(real.(EC.extrapolated_E), imag.(EC.extrapolated_E), label="extrapolated", m=:x)
+        scatter!(real.(EC.extrapolated_E), imag.(EC.extrapolated_E), label="extrapolated")
     end
 
     isnothing(basis_points) || scatter!(real.(basis_points), imag.(basis_points), m=:x, label="basis")

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"

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

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")

calculations/ECCC.jl

@@ -0,0 +1,72 @@
+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 - 3im], [512])
+H0 = get_T_matrix(p, μ)
+dim = size(H0, 1)
+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 vectors
+training_c = range(1.2, 0.9, 9) # original: range(1.35, 0.9, 5)
+ref_E = -0.3
+order::Int = ceil((length(training_c) - 1) / 2) # order of the Pade approximant
+
+training_E = ComplexF64[]
+training_vecs = Vector{ComplexF64}[]
+
+for c in training_c
+    println("Training for c = $c")
+    H = H0 + c .* V
+    evals, evecs = eigs(H, sigma=ref_E, maxiter=5000, ritzvec=true, check=1)
+    val = nearest(evals, ref_E)
+    push!(training_E, val)
+    vec = evecs[:, nearestIndex(evals, val)]
+    vec ./= sqrt(only(transpose(vec) * vec)) # normalize
+    push!(training_vecs, vec)
+end
+
+training_vecs = hcat(training_vecs...)
+
+# Calculation of target E
+target_c = range(0.78, 0.45, 7) # original: range(0.75, 0.40, 8)
+target_E = [quick_pole_E(V_system(c)) for c in target_c]
+
+# 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)')
+a = zeros(ComplexF64, dim, order+1)
+b = zeros(ComplexF64, dim, order+1)
+for i in 1:dim
+    println("Fitting coefficients $i/$dim")
+    M_right = -training_vecs[i, :] .* M_left[:, 2:end] # remove the first column
+    M = hcat(M_left, M_right) # M = [M_left | M_right]
+    sol = M \ training_vecs[i, :]
+    a[i, :] .= sol[1:order+1]
+    b[i, :] .= [1; sol[order+2:end]]
+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
+extrapolating_c = [training_c; target_c]
+extrapolated_vec = pade_approx.(extrapolating_c)
+extrapolated_E = [only(transpose(vec) * (H0 + c .* V) * vec) for (vec, c) in zip(extrapolated_vec, extrapolating_c)]
+
+# Plotting
+scatter(real.(training_E), imag.(training_E), label="training")
+scatter!(real.(target_E), imag.(target_E), label="target")
+scatter!(real.(extrapolated_E), imag.(extrapolated_E), label="extrapolated", m=:star5)

calculations/PMM.jl

@@ -0,0 +1,71 @@
+using LinearAlgebra, Random, Plots
+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)
+
+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)
+
+# calculate training data
+data_c = training_c
+data_E = [quick_pole_E(V_system(c)) for c in data_c]
+
+# hyperparameters
+N = 9
+
+# initialize random Hamiltonians
+H0 = randn(ComplexF64, N, N)
+H0 = H0 + transpose(H0) # symmetric
+H1 = randn(ComplexF64, N, N)
+H1 = H1 + transpose(H1) # symmetric
+
+# training
+Es = ComplexF64[]
+ψs = Vector{ComplexF64}[]
+
+lr = 0.05
+epochs = 100000
+for epoch in 1:epochs
+    empty!(Es)
+    empty!(ψs)
+    for (c, E) in zip(data_c, data_E)
+        H = H0 + c * H1
+        evals, evecs = eigen(H)
+        i = nearestIndex(evals, E) # TODO: more robust way to identify the eigenvector
+        push!(Es, evals[i])
+        push!(ψs, evecs[:, i])
+    end
+
+    if epoch % 1000 == 0
+        loss = sum(abs2, Es .- data_E)
+        println("Epoch:$epoch/$epochs \t Loss: $loss")
+    end
+
+    # gradient of the loss function
+    function grad(c_order=0)
+        out = zeros(ComplexF64, N, N)
+        for (c, E_target, ψ, E) in zip(data_c, data_E, ψs, Es)
+            out .+= (c^c_order * conj(E - E_target)) .* (ψ * transpose(ψ))
+        end
+        return 2 .* real.(out)
+    end
+    H0 .-= lr .* grad(0) # update H0
+    H1 .-= lr .* grad(1) # update H1    
+end
+
+# evaluate for all points
+all_c = vcat(training_c, extrapolating_c)
+exact_E = [quick_pole_E(V_system(c)) for c in all_c]
+extrapolated_E = ComplexF64[]
+for (c, ref) in zip(all_c, exact_E)
+    H = H0 + c * H1
+    evals, evecs = eigen(H)
+    evals = vcat(evals, conj.(evals)) # include complex conjugates
+    push!(extrapolated_E, nearest(evals, ref))
+end
+
+# plot results
+scatter(real.(data_E), imag.(data_E), label="training", title="PMM", xlabel="Re E", ylabel="Im E")
+scatter!(real.(exact_E), imag.(exact_E), label="exact", m=:+)
+scatter!(real.(extrapolated_E), imag.(extrapolated_E), label="predicted", m=:x)

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()
-
-# %%

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")

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")

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")

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.

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")

calculations/XZ_technique.jl

@@ -0,0 +1,30 @@
+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))