Major optimization to the Dirac equation

nucleons.jl

@@ -11,16 +11,32 @@ const r_reg = 1E-8 # fm # regulator for the centrifugal term
     κ is the generalized angular momentum,
     p is true for proton and false for neutron,
     E in the energy in MeV,
-    Φ0, W0, B0, A0 are the mean-field potentials (couplings included) in MeV as functions of r in fm,
+    f1(r) =  M-Φ0(r)-W0(r)-(p-0.5)*2B0(r)-p*A0(r) is a function of r in MeV (see optimized_dirac_potentials()),
+    f2(r) = -M+Φ0(r)-W0(r)-(p-0.5)*2B0(r)-p*A0(r) is a function of r in MeV (see optimized_dirac_potentials()),
     r is the radius in fm.
 Reference: P. Giuliani, K. Godbey, E. Bonilla, F. Viens, and J. Piekarewicz, Frontiers in Physics 10, (2023)"
-function dirac!(du, u, (κ, p, E, Φ0, W0, B0, A0), r)
+function dirac!(du, u, (κ, E, f1, f2), r)
-    M = p ? M_p : M_n
-    common1 = E - W0(r) - (p - 0.5) * 2B0(r) - p * A0(r)
-    common2 = M - Φ0(r)
     (g, f) = u
-    du[1] = -(κ/(r + r_reg)) * g + (common1 + common2) * f / ħc
+    @inbounds du[1] = -(κ/(r + r_reg)) * g + (E + f1(r)) * f / ħc
-    du[2] =  (κ/(r + r_reg)) * f - (common1 - common2) * g / ħc
+    @inbounds du[2] =  (κ/(r + r_reg)) * f - (E + f2(r)) * g / ħc
+end
+
+"Get the potentials f1 and f2 that goes into the Dirac equation, defined as
+    f1(r) =  M-Φ0(r)-W0(r)-(p-0.5)*2B0(r)-p*A0(r),
+    f2(r) = -M+Φ0(r)-W0(r)-(p-0.5)*2B0(r)-p*A0(r)."
+function optimized_dirac_potentials(p, s::system)
+    M = p ? M_p : M_n
+
+    f1s = zero_array(s)
+    f2s = zero_array(s)
+
+    @. f1s =  M - s.Φ0 - s.W0 - (p - 0.5) * 2 * s.B0 - p * s.A0
+    @. f2s = -M + s.Φ0 - s.W0 - (p - 0.5) * 2 * s.B0 - p * s.A0
+
+    f1 = linear_interpolation(rs(s), f1s)
+    f2 = linear_interpolation(rs(s), f2s)
+
+    return (f1, f2)
 end
 
 "Solve the Dirac equation and return the wave function u(r)=[g(r), f(r)] where
@@ -28,12 +44,12 @@ end
     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 solveNucleonWf(κ, p, E, s::system; shooting=true, normalize=true, algo=Tsit5())
-    (Φ0, W0, B0, A0) = fields_interp(s)
+    (f1, f2) = optimized_dirac_potentials(p, s)
 
     if shooting
         @assert s.divs % 2 == 0 "divs must be an even number when shooting=true"
         prob = ODEProblem(dirac!, [0, 1], (s.r_max, s.r_max / 2))
-        sol = solve(prob, algo, p=(κ, p, E, Φ0, W0, B0, A0), saveat=Δr(s))
+        sol = solve(prob, algo, p=(κ, E, f1, f2), saveat=Δr(s))
         wf_right = reverse(hcat(sol.u...); dims=2)
         next_r_max = s.r_max / 2 # for the next step
     else
@@ -41,7 +57,7 @@ function solveNucleonWf(κ, p, E, s::system; shooting=true, normalize=true, algo
     end
     
     prob = ODEProblem(dirac!, [0, 1], (0, next_r_max))
-    sol = solve(prob, algo, p=(κ, p, E, Φ0, W0, B0, A0), saveat=Δr(s))
+    sol = solve(prob, algo, p=(κ, E, f1, f2), saveat=Δr(s))
     wf = hcat(sol.u...)
 
     if shooting # join two segments
@@ -66,8 +82,8 @@ end
     the other parameters are the same from dirac!(...)."
 function boundaryValueFunc(κ, p, s::system; dtype=Float64, algo=Tsit5())
     prob = ODEProblem(dirac!, convert.(dtype, [0, 1]), (0, s.r_max))
-    (Φ0, W0, B0, A0) = fields_interp(s)
+    (f1, f2) = optimized_dirac_potentials(p, s)
-    func(E) = solve(prob, algo, p=(κ, p, E, Φ0, W0, B0, A0), saveat=[s.r_max], save_idxs=[1])[1, 1]
+    func(E) = solve(prob, algo, p=(κ, E, f1, f2), saveat=[s.r_max], save_idxs=[1])[1, 1]
     return func
 end