# Build your own nuclear DFT -- step 2: one bound state in a fixed well.
#
# Before any self-consistency, solve the simplest quantum problem in the file:
# a single nucleon in a FIXED Woods-Saxon potential. For a spherical potential
# the 3-D Schrodinger equation separates, and the s-wave (l = 0) radial part
# for u(r) = r R(r) is one-dimensional:
#
#     -(hbar^2/2m) u'' + U(r) u = eps u,      u(0) = u(rmax) = 0.
#
# On the mesh, u'' becomes the (u[i-1] - 2u[i] + u[i+1])/h^2 stencil, so H is
# a TRIDIAGONAL matrix. We want its lowest eigenpair, and there is an O(N)
# trick for that: pick a shift sigma below the whole spectrum (Gershgorin:
# diagonal minus the reach of the off-diagonals), then apply (H - sigma I)^-1
# over and over. Each application is a Thomas solve -- forward sweep, back
# substitution -- and the iteration converges to the lowest state because
# that's the eigenvalue closest to sigma.

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 }

# a standard fixed Woods-Saxon well: 51 MeV deep, R = 1.27 A^(1/3), a = 0.67
v0 = -51.0
rws = 1.27 * pow(amass, 1.0 / 3.0)
aws = 0.67
u_pot = zeros(nr)
for i to nr { u_pot[i] = v0 / (1.0 + exp((r[i] - rws) / aws)) }

# tridiagonal H: constant mass for now (the effective mass comes in step 5)
diag = zeros(nr)
off = zeros(nr)
for i to nr {
    diag[i] = 2.0 * hbar2_2m / (h * h) + u_pot[i]
    if i < nr - 1.0 { off[i] = 0.0 - hbar2_2m / (h * h) }
}

# Gershgorin lower bound: min over rows of (diagonal - |off-diagonals|).
# The kinetic +2/h^2 on the diagonal cancels against the two -1/h^2 couplings,
# so this lands near the bottom of the potential -- a tight shift.
sigma = diag[0]
for i to nr {
    g = diag[i]
    if i > 0 { g = g - abs(off[i - 1.0]) }
    if i < nr - 1.0 { g = g - abs(off[i]) }
    if g < sigma { sigma = g }
}
sigma = sigma - 1.0
show sigma           # ~ -52: just under the well floor, as promised

# inverse iteration with Thomas solves
u = zeros(nr)
w = zeros(nr)
tc = zeros(nr)
td = zeros(nr)
for i to nr { u[i] = 1.0 }
eps = 0.0
for sweep to 200 {
    # forward sweep of (H - sigma I) w = u
    tc[0] = off[0] / (diag[0] - sigma)
    td[0] = u[0] / (diag[0] - sigma)
    for i to nr {
        if i > 0 {
            m = diag[i] - sigma - off[i - 1.0] * tc[i - 1.0]
            if i < nr - 1.0 { tc[i] = off[i] / m }
            td[i] = (u[i] - off[i - 1.0] * td[i - 1.0]) / m
        }
    }
    # back substitution
    w[nr - 1.0] = td[nr - 1.0]
    k = nr - 2.0
    while k >= 0.0 {
        w[k] = td[k] - tc[k] * w[k + 1.0]
        k = k - 1.0
    }
    # normalize, measure the change, adopt
    nrm = 0.0
    for i to nr { nrm += w[i] * w[i] }
    nrm = sqrt(nrm)
    change = 0.0
    for i to nr {
        d = w[i] / nrm - u[i]
        change += d * d
        u[i] = w[i] / nrm
    }
    if change < 1.0e-24 { break }
}
# Rayleigh quotient: the eigenvalue of the converged vector
for i to nr {
    eps += diag[i] * u[i] * u[i]
    if i < nr - 1.0 { eps += 2.0 * off[i] * u[i] * u[i + 1.0] }
}
show eps             # the 1s level of the well, ~ -38 MeV: deeply bound

show u               # u(r) = r R(r): rises from 0, one bump, decays -- no nodes
show u_pot           # the well it lives in