# Build your own nuclear DFT -- step 3: shell structure from one fixed well.
#
# Two upgrades to step 2 and the magic numbers appear on their own.
#
# (1) Angular momentum. An orbital with l > 0 feels the centrifugal barrier
#     B l(l+1)/r^2 on the diagonal -- that is the entire difference between
#     an s-state and a p-state in the radial picture.
# (2) Spin-orbit. The nuclear ls force is a surface effect: W(r) ~ df/dr peaks
#     where the density falls. It enters as  sdl * W(r)/r  with
#     sdl = <sigma.l> = j(j+1) - l(l+1) - 3/4  = { l         (j = l + 1/2)
#                                                { -(l+1)     (j = l - 1/2)
#     so each l level splits into a pair, and the split grows like 2l+1.
#
# One more tool: to get the SECOND state of a block (the 2s after the 1s), run
# the same inverse iteration but PROJECT OUT the states already found after
# every solve -- deflation. The iteration then converges to the lowest state
# in the space orthogonal to them, which is exactly the next one up.

pi = 3.141592653589793
hbar2_2m = 20.7355
h = 0.1
nr = 140
amass = 16.0

r = zeros(nr)
for i to nr { r[i] = (i + 1.0) * h }

v0 = -51.0
rws = 1.27 * pow(amass, 1.0 / 3.0)
aws = 0.67
u_pot = zeros(nr)
wso = zeros(nr)
for i to nr {
    f = 1.0 / (1.0 + exp((r[i] - rws) / aws))
    u_pot[i] = v0 * f
    # surface-peaked spin-orbit form factor: v_so * df/dr (negative at surface)
    wso[i] = 22.0 * (0.0 - f * (1.0 - f) / aws)
}

# ---- the block solver from step 2, generalized: k lowest states, any (l, j) --
diag = zeros(nr)
off = zeros(nr)
u_new = zeros(nr)
u_cur = zeros(nr)
tc = zeros(nr)
td = zeros(nr)
ev = zeros(1024)     # up to 2 states * 512 slots
eps_k = zeros(2)

def build_h(l, sdl) {
    for i to nr {
        diag[i] = 2.0 * hbar2_2m / (h * h) + hbar2_2m * l * (l + 1.0) / (r[i] * r[i]) + u_pot[i] + sdl * wso[i] / r[i]
        if i < nr - 1.0 { off[i] = 0.0 - hbar2_2m / (h * h) }
    }
    return 0.0
}

def lowest_k(n, kk) {
    sigma = diag[0]
    for i to n {
        g = diag[i]
        if i > 0 { g = g - abs(off[i - 1.0]) }
        if i < n - 1.0 { g = g - abs(off[i]) }
        if g < sigma { sigma = g }
    }
    sigma = sigma - 1.0
    for s to kk {
        # start from a sine with s interior nodes: it already looks like the target
        for i to n { u_new[i] = sin((s + 1.0) * pi * (i + 1.0) / (n + 1.0)) }
        lam = 0.0
        for sweep to 300 {
            for i to n { u_cur[i] = u_new[i] }
            tc[0] = off[0] / (diag[0] - sigma)
            td[0] = u_cur[0] / (diag[0] - sigma)
            for i to n {
                if i > 0 {
                    m = diag[i] - sigma - off[i - 1.0] * tc[i - 1.0]
                    if i < n - 1.0 { tc[i] = off[i] / m }
                    td[i] = (u_cur[i] - off[i - 1.0] * td[i - 1.0]) / m
                }
            }
            u_new[n - 1.0] = td[n - 1.0]
            k = n - 2.0
            while k >= 0.0 {
                u_new[k] = td[k] - tc[k] * u_new[k + 1.0]
                k = k - 1.0
            }
            # deflation: subtract the states already found
            for t to s {
                dot = 0.0
                for i to n { dot += ev[t * 512.0 + i] * u_new[i] }
                for i to n { u_new[i] = u_new[i] - dot * ev[t * 512.0 + i] }
            }
            nrm = 0.0
            for i to n { nrm += u_new[i] * u_new[i] }
            nrm = sqrt(nrm)
            for i to n { u_new[i] = u_new[i] / nrm }
            newlam = 0.0
            for i to n {
                newlam += diag[i] * u_new[i] * u_new[i]
                if i < n - 1.0 { newlam += 2.0 * off[i] * u_new[i] * u_new[i + 1.0] }
            }
            dl = abs(newlam - lam)
            lam = newlam
            if sweep > 2.0 {
                if dl < 1.0e-12 { break }
            }
        }
        eps_k[s] = lam
        for i to n { ev[s * 512.0 + i] = u_new[i] }
    }
    return 0.0
}

# ---- walk the shell-model blocks of a light nucleus --------------------------
# levels in filling order: 1s1/2, 1p3/2, 1p1/2, 1d5/2, 2s1/2, 1d3/2
levels = zeros(6)

build_h(0.0, 0.0)
lowest_k(nr, 2.0)            # the l=0 block holds BOTH the 1s and the 2s
levels[0] = eps_k[0]
levels[4] = eps_k[1]

build_h(1.0, 1.0)            # j = 3/2: sdl = +l = 1
lowest_k(nr, 1.0)
levels[1] = eps_k[0]

build_h(1.0, -2.0)           # j = 1/2: sdl = -(l+1) = -2
lowest_k(nr, 1.0)
levels[2] = eps_k[0]

build_h(2.0, 2.0)            # 1d5/2: sdl = +2
lowest_k(nr, 1.0)
levels[3] = eps_k[0]

build_h(2.0, -3.0)           # 1d3/2: sdl = -3
lowest_k(nr, 1.0)
levels[5] = eps_k[0]

show levels
# Read the plot: 1s alone at the bottom (magic 2); the 1p pair, split by the
# spin-orbit term (fill both: magic 8 -- that's 16O); then the sd shell.
# The 1p3/2 sits BELOW the 1p1/2: in nuclei the aligned partner is MORE bound,
# opposite to atoms, and that sign is what makes 28, 50, 82, 126 magic.
p_split = levels[2] - levels[1]
show p_split                 # the 1p spin-orbit splitting, a few MeV