From 6538c30fa8d9f4d62c32a93f4637d6a33114717d Mon Sep 17 00:00:00 2001
From: Nuwan Yapa <yapanuwan@gmail.com>
Date: Mon, 1 Jul 2024 19:45:43 -0400
Subject: [PATCH] Calculate and plot wave function

---
 dirac.jl         |  9 +++++++++
 test/Pb208_wf.jl | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)
 create mode 100644 test/Pb208_wf.jl

diff --git a/dirac.jl b/dirac.jl
index 234611f..2317cf0 100644
--- a/dirac.jl
+++ b/dirac.jl
@@ -23,6 +23,15 @@ function dirac!(du, u, (κ, p, E, Φ0, W0, B0, A0), r)
     du[2] =  (κ/(r + r_reg)) * f - (common1 - common2) * g / ħc
 end
 
+"Solve the Dirac equation and return the wave function u(r)=[g(r), f(r)] where
+    divs is the number of mesh divisions if the solution should be discretized as a 2×(1+divs) matrix (keep divs=0 to obtain an interpolating function),
+    the other parameters are the same from dirac!(...)."
+function solveWf(κ, p, E, Φ0, W0, B0, A0, r_max, divs::Int=0)
+    prob = ODEProblem(dirac!, [0, 1], (0, r_max))
+    sol = solve(prob, RK4(), p=(κ, p, E, Φ0, W0, B0, A0), saveat=(divs == 0 ? [] : r_max/divs))
+    return divs == 0 ? sol : hcat(sol.u...)
+end
+
 "Solve the Dirac equation and return g(r=r_max) where
     r_max is the outer boundary in fm,
     the other parameters are the same from dirac!(...)."
diff --git a/test/Pb208_wf.jl b/test/Pb208_wf.jl
new file mode 100644
index 0000000..e87c3c9
--- /dev/null
+++ b/test/Pb208_wf.jl
@@ -0,0 +1,37 @@
+using DelimitedFiles, Interpolations, Plots
+include("../dirac.jl")
+
+# test data generated from Hartree.f
+# format: x S(x) V(x) R(x) A(x)
+test_data = readdlm("test/Pb208FldsFSUGarnet.csv") 
+xs = test_data[:, 1]
+Ss = test_data[:, 2]
+Vs = test_data[:, 3]
+Rs = test_data[:, 4]
+As = test_data[:, 5]
+
+S_interp = linear_interpolation(xs, Ss)
+V_interp = linear_interpolation(xs, Vs)
+R_interp = linear_interpolation(xs, Rs)
+A_interp = linear_interpolation(xs, As)
+
+κ = -1
+p = true
+r_max = maximum(xs)
+E_min = 880
+E_max = 939
+
+boundEs = findEs(κ, p, S_interp, V_interp, R_interp, A_interp, r_max, E_min, E_max)
+groundE = minimum(boundEs)
+println("ground state E = $groundE")
+
+plot_r_max = r_max * 0.75
+divs = 50
+wf = solveWf(κ, p, groundE, S_interp, V_interp, R_interp, A_interp, plot_r_max, divs)
+rs = collect(0: plot_r_max/divs : plot_r_max)
+gs = wf[1, :]
+fs = wf[2, :]
+
+plot(rs, gs, label="g(r)")
+plot!(rs, fs, label="f(r)")
+xlabel!("r (fm)")