From eaf236fcc03aa24d87ed2fc70b4842936919edb1 Mon Sep 17 00:00:00 2001
From: Jusong Yu <jusong.yeu@gmail.com>
Date: Fri, 6 Sep 2024 17:52:07 +0200
Subject: [PATCH] Example of pseudolize and logder

---
 examples/logder.jl | 118 ++++++++++++++++++++++++++++++++++-----------
 src/ode_poisson.jl |   4 +-
 src/pseudolize.jl  |   6 +--
 3 files changed, 94 insertions(+), 34 deletions(-)

diff --git a/examples/logder.jl b/examples/logder.jl
index d6f0c32..8f816cc 100644
--- a/examples/logder.jl
+++ b/examples/logder.jl
@@ -3,6 +3,9 @@ using PGEN
 using Plots
 gr()
 
+# Define different color for each l
+colors = Dict(0 => :red, 1 => :blue, 2 => :green, 3 => :orange, 4 => :purple)
+
 # %%
 # AE results for different pseudolize methods
 # Only need to run once
@@ -23,60 +26,117 @@ res = self_consistent_field(
 
 ae_info = (ε_lst = res.ε_lst, ϕs = res.ϕs, vae = res.v_tot, occs = res.occs, xc = res.xc);
 
-# %%
-plot(mesh.r, ae_info.vae, label = "AE pot", line = (2, :solid), lc = :black)
-ylims!(-10, 5)
-xlims!(0, 5)
-
 # %%
 
 # rc for every orbital, pass as a dict
 # angular momentum => rc
 rc = Dict{NamedTuple{(:n, :l),Tuple{Int64,Int64}},Float64}(
-    (n = 2, l = 0) => 2.0,
-    (n = 2, l = 1) => 1.9,
+    (n = 2, l = 0) => 1.9,
+    (n = 2, l = 1) => 2.2,
 )
 
 # Check integration of ρ to inf is equal to Z
 # Check the wavefunction is normalized (∑ϕ^2 dr = Z)
 # plot the wavefunction and density
-v_pspot = pseudolize(ae_info, mesh, rc; method = :TM, kbform = false)
-
-# %%
-
-for nl in keys(rc)
-    vline!([rc[nl]], label = "rc: l=$(nl.l)", linestyle = :dash)
-    plot!(mesh.r, v_pspot.v[nl], label = "pspot l=$(nl.l)")
-end
-display(plot!())
-
-# %%
+v_pspot, ϕ_ps = pseudolize(ae_info, mesh, rc; method = :TM, kbform = false)
 
-# save it to file
-savefig("/tmp/pseudolize_TM.png")
 
 # %%
 # Plot arctan logder and compare with AE
 # start a new plot
-plot()
-
-# Define different color for each l
-colors = Dict(0 => :red, 1 => :blue, 2 => :green, 3 => :orange, 4 => :purple)
 
 rcx = 2.0 # the r cut to compute atan logder
+x_ae = Dict()
+y_ae = Dict()
+x_ps = Dict()
+y_ps = Dict()
+
 for nl in keys(rc)
     l = nl.l
-    x_ae, y_ae =
+    x_ae[nl], y_ae[nl] =
         compute_atanld(l, Z, mesh, ae_info.vae, rcx, window = [-10.0, 10.0], δ = 0.05)
-    plot!(x_ae, y_ae, label = "AE: l = $l", linestyle = :dash, color = colors[l])
 
-    x_ps, y_ps =
+    x_ps[nl], y_ps[nl] =
         compute_atanld(l, Z, mesh, v_pspot.v[nl], rcx, window = [-10.0, 10.0], δ = 0.05)
-    plot!(x_ps, y_ps, label = "PS: l = $l", color = colors[l])
+end
+
+# %%
+plot(title = "logrithmic derivative")
+plot!(xlabel = "Energy (Ha)", ylabel = "tan-1 psi' / psi")
+
+for nl in keys(rc)
+    l = nl.l
+    plot!(x_ae[nl], y_ae[nl], label = "AE: l = $l", line = (2, :solid), color = colors[l])
+
+    plot!(x_ps[nl], y_ps[nl], label = "PS: l = $l", line = (2, :dash), color = colors[l])
 end
 
 ylims!(-10, 10)
 display(plot!())
 
 # %%
-savefig("/tmp/atanlogder_TM_SL.png")
+savefig("/tmp/atanlogder_TM_SL.pdf")
+
+# %%
+# Plot ae wavefunctions
+plot_wfc = plot()
+plot!(xlims = (0, 8), ylims = (-2, 2))
+plot!(title = "Z=$Z, s/p channels")
+
+nl = (n = 2, l = 0)
+vline!([rc[nl]], label = "", line = (1, :dot), lc = colors[nl.l])
+
+ic = findfirst(r -> r > rc[nl], mesh.r) - 1
+
+ϕ_ae = ae_info.ϕs[nl];
+ae_rc = ϕ_ae[ic];
+
+plot!(mesh.r, ϕ_ae, label = "AE, n=$(nl.n), l=$(nl.l)", line = (3, :solid), lc = colors[nl.l])
+
+phi_ps = ϕ_ps[nl]
+ps_rc = phi_ps[ic]
+phi_ps = phi_ps * ae_rc / ps_rc
+plot!(mesh.r, phi_ps, label = "PS, n=$(nl.n), l=$(nl.l)", line = (2, :dash), lc = colors[nl.l])
+
+nl = (n = 2, l = 1)
+vline!([rc[nl]], label = "", line = (1, :dot), lc = colors[nl.l])
+
+ic = findfirst(r -> r > rc[nl], mesh.r) - 1
+
+ϕ_ae = ae_info.ϕs[nl];
+ae_rc = ϕ_ae[ic];
+
+plot!(mesh.r, ϕ_ae, label = "AE, n=$(nl.n), l=$(nl.l)", line = (3, :solid), lc = colors[nl.l])
+
+phi_ps = ϕ_ps[nl]
+ps_rc = phi_ps[ic]
+phi_ps = phi_ps * ae_rc / ps_rc
+plot!(mesh.r, phi_ps, label = "PS, n=$(nl.n), l=$(nl.l)", line = (2, :dash), lc = colors[nl.l])
+plot!(ylabel = "wavefunction")
+
+# %%
+# potential
+plot_pot = plot(mesh.r, ae_info.vae, label = "AE pot", line = (3, :solid), lc = :black)
+ylims!(-3, 1)
+xlims!(0, 8)
+
+for nl in keys(rc)
+    vline!([rc[nl]], label = "l=$(nl.l), rc=$(rc[nl])", line = (1, :dot), lc = colors[nl.l])
+    plot!(
+        mesh.r,
+        v_pspot.v[nl],
+        label = "",
+        line = (2, :dash),
+        lc = colors[nl.l],
+    )
+end
+plot!(ylabel = "V(r) (Ha)")
+plot!(xlabel = "r (a.u.)")
+
+# %%
+plot(plot_wfc, plot_pot, layout=(2, 1))
+
+# %%
+
+# save it to file
+savefig("/tmp/pseudolize_TM.pdf")
diff --git a/src/ode_poisson.jl b/src/ode_poisson.jl
index 1e0d189..4f746d7 100644
--- a/src/ode_poisson.jl
+++ b/src/ode_poisson.jl
@@ -1,6 +1,6 @@
 using OrdinaryDiffEq
 
-function poisson_outward(ρ::Vector{Float64}, mesh::Mesh; max_val = 1e+6)::Vector{Float64}
+function poisson_outward(ρ::Vector{Float64}, mesh::Mesh; max_val = 1e+6, algo = VCABM5())::Vector{Float64}
     # Converted 2nd Poisson ODE in rgrid to uniform grid
     # u1p = u2 * Rp
     # u2p = -(4*pi*rho + 2*u2/r) * Rp
@@ -43,7 +43,7 @@ function poisson_outward(ρ::Vector{Float64}, mesh::Mesh; max_val = 1e+6)::Vecto
 
     cb = DiscreteCallback(condition, affect!)
     prob = DiscreteProblem(f!, y0, (1, N))
-    sol = solve(prob, ABM54(), dt = 1, adaptive = false, callback = cb)
+    sol = solve(prob, algo, dt = 1, adaptive = false, callback = cb)
 
     sol[1, :]
 end
diff --git a/src/pseudolize.jl b/src/pseudolize.jl
index 5070d5e..52a8177 100644
--- a/src/pseudolize.jl
+++ b/src/pseudolize.jl
@@ -21,7 +21,7 @@ function pseudolize(
     rc::Dict{NamedTuple{(:n, :l),Tuple{Int64,Int64}},Float64};
     method = :TM,
     kbform = true,
-)::Potential
+)
 
     # To store the results from pseudolization for post processing
     v_pspot_dict = Dict{NamedTuple{(:n, :l),Tuple{Int64,Int64}},Vector{Float64}}()
@@ -65,9 +65,9 @@ function pseudolize(
 
     if kbform
         v_kb = sl2kb(v_sl)
-        return v_kb
+        return v_kb, ϕ_ps_dict
     else
-        return v_sl
+        return v_sl, ϕ_ps_dict
     end
 end