NuclearRMF/nucleons.jl

using DifferentialEquations, Roots

const ħc = 197.33 # MeVfm
const M_n = 939.0 # MeV/c2
const M_p = 939.0 # MeV/c2

const r_reg = 1E-8 # fm # regulator for the centrifugal term

"The spherical Dirac equation that returns du=[dg, df] in-place where
    u=[g, f] are the reduced radial components evaluated at r,
    κ 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,
    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)
    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
    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 so solution would be returned as a 2×(1+divs) matrix,
    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, Φ0, W0, B0, A0, s::system; shooting=true, normalize=true)
    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, RK4(), p=(κ, p, E, Φ0, W0, B0, A0), saveat=Δr(s))
        wf_right = reverse(hcat(sol.u...); dims=2)
        next_r_max = s.r_max / 2 # for the next step
    else
        next_r_max = s.r_max
    end
    
    prob = ODEProblem(dirac!, [0, 1], (0, next_r_max))
    sol = solve(prob, RK4(), p=(κ, p, E, Φ0, W0, B0, A0), saveat=Δr(s))
    wf = hcat(sol.u...)

    if shooting # join two segments
        rescale_factor_g = wf[1, end] / wf_right[1, 1]
        rescale_factor_f = wf[2, end] / wf_right[2, 1]
        @assert isfinite(rescale_factor_g) && isfinite(rescale_factor_f) "Cannot rescale the right partition"
        isapprox(rescale_factor_g, rescale_factor_f; rtol=0.03) || @warn "Discontinuity between two partitions"
        wf_right_rescaled = wf_right .* ((rescale_factor_g + rescale_factor_f) / 2)
        wf = hcat(wf[:, 1:(end - 1)], wf_right_rescaled)
    end

    if normalize
        norm = sum(wf .* wf) * Δr(s) # integration by Reimann sum
        wf = wf ./ sqrt(norm)
    end

    return wf
end

"Returns a function that solves the Dirac equation and returns g(r=r_max) where
    r_max is the outer boundary in fm,
    the other parameters are the same from dirac!(...)."
function boundaryValueFunc(κ, p, Φ0, W0, B0, A0, s::system; dtype=Float64, algo=Tsit5())
    prob = ODEProblem(dirac!, convert.(dtype, [0, 1]), (0, s.r_max))
    func(E) = solve(prob, algo, p=(κ, p, E, Φ0, W0, B0, A0), saveat=[s.r_max], save_idxs=[1])[1, 1]
    return func
end

"Find all bound energies between E_min (=0) and E_max (=mass) where
    tol_digits determines the precision for root finding and the threshold for identifying duplicates,
    the other parameters are the same from dirac!(...)."
function findEs(κ, p, Φ0, W0, B0, A0, s::system, E_min=0, E_max=(p ? M_p : M_n), tol_digits=5)
    func = boundaryValueFunc(κ, p, Φ0, W0, B0, A0, s)
    Es = find_zeros(func, (E_min, E_max); xatol=1/10^tol_digits)
    return unique(E -> round(E; digits=tol_digits), Es)
end

"Find all orbitals and return two lists containing κ values and corresponding energies for a single species where
    the other parameters are defined above"
function findAllOrbitals(p, Φ0, W0, B0, A0, s::system, E_min=0, E_max=(p ? M_p : M_n))
    κs = Int[]
    Es = Float64[]
    # start from κ=-1 and go both up and down
    for direction in [-1, 1]
        for κ in direction * (1:100) # cutoff is 100
            new_Es = findEs(κ, p, Φ0, W0, B0, A0, s, E_min, E_max)
            if isempty(new_Es); break; end
            append!(Es, new_Es)
            append!(κs, fill(κ, length(new_Es)))
        end
    end
    return (κs, Es)
end

"For a given list of κ values with corresponding energies, attempt to fill Z_or_N lowest lying orbitals and return occupancy numbers"
function fillNucleons(Z_or_N::Int, κs, Es)
    sort_i = sortperm(Es)

    occ = zeros(Int, length(κs))

    for i in sort_i
        if Z_or_N ≤ 0; break; end;
        max_occ = 2 * j_κ(κs[i]) + 1
        occ[i] = min(max_occ, Z_or_N)
        Z_or_N -= occ[i]
    end

    Z_or_N == 0 || @warn "All orbitals could not be filled"
    return occ
end

"Total angular momentum j for a given κ value"
j_κ(κ::Int) = abs(κ) - 1/2

"Orbital angular momentum l for a given κ value"
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 arrays κs, Es, occs tabulate the energies and occupation numbers corresponding to each κ,
    the other parameters are defined above"
function calculateNucleonDensity(κs, Es, occs, p, Φ0, W0, B0, A0, s::system)
    ρr2 = zeros(2, s.divs + 1) # ρ×r² values
    
    for (κ, E, occ) in zip(κs, Es, occs)
        wf = solveNucleonWf(κ, p, E, Φ0, W0, B0, A0, s; shooting=true, normalize=true)
        wf2 = wf .* wf
        ρr2 += (occ / (4 * pi)) * wf2 # 2j+1 factor is accounted in the occupancy number
    end

    r2s = rs(s).^2
    ρ = ρr2 ./ transpose(r2s)
    ρ[:, 1] .= ρ[:, 2] # dirty fix for NaN at r=0
    
    ρ_s = ρ[1, :] - ρ[2, :] # g^2 - f^2
    ρ_v = ρ[1, :] + ρ[2, :] # g^2 + f^2
    
    return (ρ_s, ρ_v)
end

"Solve the Dirac equation and calculate scalar and vector densities of a nucleon species where
    the other parameters are defined above"
function solveNucleonDensity(p, Φ0, W0, B0, A0, s::system, E_min=800, E_max=939)
    κs, Es = findAllOrbitals(p, Φ0, W0, B0, A0, s, E_min, E_max)
    occs = fillNucleons(Z_or_N(s, p), κs, Es)
    return calculateNucleonDensity(κs, Es, occs, p, Φ0, W0, B0, A0, s)
end