2-body XZ in p-space

calculations/XZ/2body_p_space.jl

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