From 00a5f664188198cccb1e523f28b0f2aabd87b057 Mon Sep 17 00:00:00 2001
From: Nuwan Yapa <yapanuwan@gmail.com>
Date: Fri, 22 Nov 2024 18:21:08 -0500
Subject: [PATCH] Shooting method for wave function

---
 dirac.jl | 34 ++++++++++++++++++++--------------
 1 file changed, 20 insertions(+), 14 deletions(-)

diff --git a/dirac.jl b/dirac.jl
index 9a78a3c..d8765df 100644
--- a/dirac.jl
+++ b/dirac.jl
@@ -25,23 +25,29 @@ end
 
 "Solve the Dirac equation and return the wave function u(r)=[g(r), f(r)] where
     divs is the number of mesh divisions so solution would be returned as a 2×(1+divs) matrix,
-    refine determines whether to internally enable high-precision mode and optimize the energy beforehand (assuming a bound state),
+    shooting method divides the interval into two partitions at r_max/2, ensuring convergence at both r=0 and r=r_max,
     the other parameters are the same from dirac!(...)."
-function solveWf(κ, p, E, Φ0, W0, B0, A0, r_max, divs; refine=true, normalize=true)
-    if refine
-        dtype = BigFloat
-        algo = Feagin12()
+function solveWf(κ, p, E, Φ0, W0, B0, A0, r_max, divs; shooting=true, normalize=true)
+    saveat = r_max / divs
 
-        f(E_) = boundaryValue(κ, p, E_, Φ0, W0, B0, A0, r_max; dtype=dtype, algo=algo)
-        E = find_zero(f, convert(dtype, E))
-    else
-        dtype = Float64
-        algo = RK()
+    if shooting
+        @assert divs % 2 == 0 "divs must be an even number when shooting=true"
+        prob = ODEProblem(dirac!, [0, 1], (r_max, r_max / 2))
+        sol = solve(prob, RK4(), p=(κ, p, E, Φ0, W0, B0, A0), saveat=saveat)
+        wf_right = reverse(hcat(sol.u...); dims=2)
+        r_max = r_max / 2 # for the next step
     end
+    
+    prob = ODEProblem(dirac!, [0, 1], (0, r_max))
+    sol = solve(prob, RK4(), p=(κ, p, E, Φ0, W0, B0, A0), saveat=saveat)
+    wf = hcat(sol.u...)
 
-    prob = ODEProblem(dirac!, convert.(dtype, [0, 1]), (0, r_max))
-    sol = solve(prob, algo, p=(κ, p, E, Φ0, W0, B0, A0), saveat=r_max/divs)
-    wf = Float64.(hcat(sol.u...))
+    if shooting # join two segments
+        rescale_factor = wf[1, end] / wf_right[1, 1]
+        rescale_factor ≈ wf[2, end] / wf_right[2, 1] || @warn "Shooting method joint has a discontinuity"
+        rescale_factor = isfinite(rescale_factor) ? rescale_factor : 1
+        wf = hcat(wf[:, 1:(end - 1)], wf_right .* rescale_factor)
+    end
 
     if normalize
         norm = sum(wf .* wf) * r_max / divs # integration by Reimann sum
@@ -116,7 +122,7 @@ function calculateNucleonDensity(N, p, Φ0, W0, B0, A0, r_max, divs, E_min=0, E_
     ρr2 = zeros(2, divs + 1) # ρ×r² values
     
     for (κ, E, occ) in zip(κs, Es, occs)
-        wf = solveWf(κ, p, E, Φ0, W0, B0, A0, r_max, divs; refine=true, normalize=true)
+        wf = solveWf(κ, p, E, Φ0, W0, B0, A0, r_max, divs; shooting=true, normalize=true)
         wf2 = wf .* wf
         ρr2 += (occ / (4 * pi)) * wf2 # 2j+1 factor is accounted in the occupancy number
     end