From 5248fdafb3eba9ff8c675351b7d50dc1fd32ea2d Mon Sep 17 00:00:00 2001
From: Nuwan Yapa <yapanuwan@gmail.com>
Date: Tue, 16 Jul 2024 20:18:36 -0400
Subject: [PATCH] Calculate density (without normalization, therefore
 incorrect)

---
 dirac.jl              | 27 +++++++++++++++++++++++++--
 test/Pb208_density.jl | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 2 deletions(-)
 create mode 100644 test/Pb208_density.jl

diff --git a/dirac.jl b/dirac.jl
index be2d8ed..53fe0cb 100644
--- a/dirac.jl
+++ b/dirac.jl
@@ -33,7 +33,7 @@ function solveWf(κ, p, E, Φ0, W0, B0, A0, r_max, divs::Int=0, refine=true)
         algo = Feagin12()
 
         f(E_) = boundaryValue(κ, p, E_, Φ0, W0, B0, A0, r_max; dtype=dtype, algo=algo)
-        E = find_zero(f, E)
+        E = find_zero(f, convert(dtype, E))
     else
         dtype = Float64
         algo = RK()
@@ -98,4 +98,27 @@ end
 j_κ(κ::Int) = abs(κ) - 1/2
 
 "Orbital angular momentum l for a given κ value"
-l_κ(κ::Int) = abs(κ) - (κ < 0) # since true = 1 and false = 0
\ No newline at end of file
+l_κ(κ::Int) = abs(κ) - (κ < 0) # since true = 1 and false = 0
+
+"Calculate scalar and vector densities of a nucleon species on [0,r_max] divided into (divs+1) points and returns them as vectors (ρ_s, ρ_v) where
+    the other parameters are defined above"
+function calculateNucleonDensity(N, p, Φ0, W0, B0, A0, r_max, divs, E_min=0, E_max=(p ? M_p : M_n))
+    κs, Es = findAllOrbitals(p, Φ0, W0, B0, A0, r_max, E_min, E_max)
+    occs = fillNucleons(N, κs, Es)
+    
+    r2s = (collect ∘ range)(0, r_max, length=divs+1).^2 |> transpose
+
+    ρ_s = zeros(divs + 1)
+    ρ_v = zeros(divs + 1)
+    
+    for (κ, E, occ) in zip(κs, Es, occs)
+        wf = solveWf(κ, p, E, Φ0, W0, B0, A0, r_max, divs) # TODO: Needs to be normalized
+        wf2 = wf .* wf
+        ρ = (occ / (4 * pi)) * (wf2 ./ r2s)  # 2j+1 factor is accounted in the occupancy number
+
+        ρ_s += ρ[1, :] - ρ[2, :] # g^2 - f^2
+        ρ_v += ρ[1, :] + ρ[2, :] # g^2 + f^2
+    end
+
+    return (ρ_s, ρ_v)
+end
\ No newline at end of file
diff --git a/test/Pb208_density.jl b/test/Pb208_density.jl
new file mode 100644
index 0000000..bb33597
--- /dev/null
+++ b/test/Pb208_density.jl
@@ -0,0 +1,32 @@
+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
+N = 82
+r_max = maximum(xs)
+E_min = 880
+E_max = 939
+divs = 50
+
+(ρ_s, ρ_v) = calculateNucleonDensity(N, p, S_interp, V_interp, R_interp, A_interp, r_max, divs, E_min, E_max)
+
+rs = range(0, r_max, length=divs+1)
+
+plot(rs, ρ_s, label="ρₛ(r)")
+plot!(rs, ρ_v, label="ρᵥ(r)")
+xlabel!("r (fm)")