Step 6 — ²⁴Mg at scale

"""Build your own shell model -- step 6: 24Mg at scale.

Same machinery, bigger nucleus: 4 valence protons + 4 valence neutrons.
Three things change, none of them physics:

  - The basis is 28,503 determinants (out of 735,471 ways to place 8
    particles in 24 seats) -- found with pruned enumeration, an idea from
    cosmo: while building a determinant, bound the total Jz still reachable
    and abandon branches that cannot land on Jz = 0.
  - The Hamiltonian is stored sparse: most determinant pairs differ in more
    than two seats, and a two-body force cannot connect them.
  - Dense diagonalization is replaced by Lanczos (scipy's eigsh) -- the
    iterative eigensolver that finds the lowest states of a large sparse
    matrix from matrix-vector products alone; the companion mean-field
    tutorial builds it from scratch in its step 2.

The printed check is the strongest on this site: TWO independent production
codes, run on this machine with the same interaction, against this code's
ten lowest states.
"""
import math
import time
import numpy as np
from math import factorial, sqrt
from itertools import combinations
from scipy.sparse import coo_matrix
from scipy.sparse.linalg import eigsh

# ── sd-shell single-particle space ──────────────────────────────────────────
# .int labels: 1 = d3/2, 2 = d5/2, 3 = s1/2   (l, 2j per label)
ORBITS = {1: (2, 3), 2: (2, 5), 3: (0, 1)}   # label -> (l, 2j)

def build_sp_states():
    """All (label, 2m, 2tz) single-particle m-states. tz=+1/2 proton, -1/2 neutron
    (only relative signs matter here). Returns list and index lookup."""
    states = []
    for tz2 in (+1, -1):
        for label, (l, j2) in ORBITS.items():
            for m2 in range(-j2, j2 + 1, 2):
                states.append((label, m2, tz2))
    return states, {s: i for i, s in enumerate(states)}

SP, SP_IDX = build_sp_states()
NSP = len(SP)   # 24

# ── Clebsch-Gordan (Condon–Shortley), all arguments doubled ─────────────────
def cg(j1, m1, j2, m2, J, M):
    if m1 + m2 != M or J > j1 + j2 or J < abs(j1 - j2):
        return 0.0
    if abs(m1) > j1 or abs(m2) > j2 or abs(M) > J:
        return 0.0
    def f(x2):  # factorial of a doubled-integer/2, must be integral
        assert x2 % 2 == 0
        return factorial(x2 // 2)
    pre = (J + 1) * f(J + j1 - j2) * f(J - j1 + j2) * f(j1 + j2 - J) / f(j1 + j2 + J + 2)
    pre *= f(J + M) * f(J - M) * f(j1 - m1) * f(j1 + m1) * f(j2 - m2) * f(j2 + m2)
    s = 0.0
    for k2 in range(0, j1 + j2 + J + 2, 2):
        d = [j1 + j2 - J - k2, j1 - m1 - k2, j2 + m2 - k2,
             J - j2 + m1 + k2, J - j1 - m2 + k2]
        if any(x < 0 for x in d):
            continue
        s += (-1) ** (k2 // 2) / (f(k2) * f(d[0]) * f(d[1]) * f(d[2]) * f(d[3]) * f(d[4]))
    return sqrt(pre) * s

# ── read the .int file ───────────────────────────────────────────────────────
def read_int(path):
    spe = {}
    tbme = {}   # (a,b,c,d,J,T) -> V, canonical a<=b, c<=d, (ab)<=(cd)
    header_tail = None
    n_lines = None
    with open(path) as fh:
        for raw in fh:
            line = raw.strip()
            if not line or line.startswith(('!', '#')):
                continue
            w = line.split()
            if n_lines is None:
                n_lines = int(w[0])
                spe = {1: float(w[1]), 2: float(w[2]), 3: float(w[3])}
                header_tail = [float(x) for x in w[4:]]
                continue
            a, b, c, d, J, T = (int(x) for x in w[:6])
            V = float(w[6])
            # canonicalize pair order inside bra/ket with the standard phase
            def canon(x, y):
                if x <= y:
                    return x, y, 1.0
                jx, jy = ORBITS[x][1], ORBITS[y][1]
                # |ba;JT> = (-1)^{(jx+jy)/2 - J - T} |ab;JT>  (j's doubled):
                # (-1)^{ja+jb-J} from the angular CG swap, (-1)^{1-T} from the
                # isospin CG swap, and one more (-1) from anticommuting a†a†.
                ph = (-1.0) ** ((jx + jy) // 2 - J - T)
                return y, x, ph
            a, b, pab = canon(a, b)
            c, d, pcd = canon(c, d)
            key = (a, b, c, d, J, T) if (a, b) <= (c, d) else (c, d, a, b, J, T)
            tbme[key] = pab * pcd * V
    return spe, tbme, header_tail

# ── numerically built coupled pair states ───────────────────────────────────
def pair_index_maps():
    pairs = list(combinations(range(NSP), 2))
    return pairs, {p: i for i, p in enumerate(pairs)}

PAIRS, PAIR_IDX = pair_index_maps()

def coupled_pair_vector(a, b, J2, M2, T2, Tz2):
    """|ab; J M T Tz> as a vector over ordered pairs (alpha<beta) of m-states.
    Built by brute expansion; normalized numerically (None if it vanishes)."""
    ja, jb = ORBITS[a][1], ORBITS[b][1]
    v = np.zeros(len(PAIRS))
    for ma in range(-ja, ja + 1, 2):
        mb = M2 - ma
        if abs(mb) > jb:
            continue
        cj = cg(ja, ma, jb, mb, J2, M2)
        if cj == 0.0:
            continue
        for ta in (+1, -1):
            tb = Tz2 - ta
            if abs(tb) > 1:
                continue
            ct = cg(1, ta, 1, tb, T2, Tz2)
            if ct == 0.0:
                continue
            ia, ib = SP_IDX[(a, ma, ta)], SP_IDX[(b, mb, tb)]
            if ia == ib:
                continue
            # a†_ia a†_ib |0> = |ia ib> ordered: sign if ia > ib
            if ia < ib:
                v[PAIR_IDX[(ia, ib)]] += cj * ct
            else:
                v[PAIR_IDX[(ib, ia)]] -= cj * ct
    n = np.linalg.norm(v)
    return (v / n, n) if n > 1e-12 else (None, 0.0)

def build_pair_hamiltonian(tbme, scale):
    """V2[p, q]: two-body matrix in the ordered-pair basis, from the TBME file.
    V2 = sum_JT V_JT sum_{M Tz} |ab;JMTTz><cd;JMTTz| with numerically
    normalized projectors — conventions cannot drift."""
    V2 = np.zeros((len(PAIRS), len(PAIRS)))
    for (a, b, c, d, J, T), V in tbme.items():
        J2, T2 = 2 * J, 2 * T
        for M2 in range(-J2, J2 + 1, 2):
            for Tz2 in range(-T2, T2 + 1, 2):
                va, na = coupled_pair_vector(a, b, J2, M2, T2, Tz2)
                if va is None:
                    continue
                vc, nc = (va, na) if (a, b, J, T) == (c, d, J, T) else \
                    coupled_pair_vector(c, d, J2, M2, T2, Tz2)
                if (a, b) == (c, d):
                    vc = va
                elif vc is None:
                    continue
                outer = np.outer(va, vc) * (V * scale)
                V2 += outer
                if (a, b) != (c, d):
                    V2 += outer.T
    return V2

# ── many-body basis: bitstrings with pruned enumeration ─────────────────────
def build_basis(n_protons, n_neutrons, M2_target):
    """All determinants (ints, bit i = m-state i occupied) with the right
    particle numbers per species and total 2*Jz = M2_target. Enumerated per
    species with cosmo-style pruning on the reachable Jz range."""
    def species_states(tz2):
        return [i for i, (lab, m2, t) in enumerate(SP) if t == tz2]

    def enum(states, n, m2_needed_min, m2_needed_max):
        # recursive with bound pruning: sort by m2 so prefix sums bound reach
        out = []
        ms = [SP[i][1] for i in states]
        # suffix min/max attainable sums for k picks from tail
        def rec(start, left, acc_bits, acc_m):
            if left == 0:
                out.append((acc_bits, acc_m))
                return
            for k in range(start, len(states) - left + 1):
                rem = left - 1
                tail = ms[k + 1:]
                lo = acc_m + ms[k] + sum(sorted(tail)[:rem])
                hi = acc_m + ms[k] + sum(sorted(tail)[-rem:] if rem else [])
                if lo > m2_needed_max or hi < m2_needed_min:
                    continue
                rec(k + 1, rem, acc_bits | (1 << states[k]), acc_m + ms[k])
        rec(0, n, 0, 0)
        return out

    prot = species_states(+1)
    neut = species_states(-1)
    # protons can carry any m2p; neutrons must supply M2_target - m2p
    m2_all = [SP[i][1] for i in prot]
    span = sum(sorted(m2_all)[-n_protons:]) if n_protons else 0
    plist = enum(prot, n_protons, -span, span) if n_protons else [(0, 0)]
    from collections import defaultdict
    by_need = defaultdict(list)
    nlist = enum(neut, n_neutrons, -span, span) if n_neutrons else [(0, 0)]
    for bits, m in nlist:
        by_need[m].append(bits)
    basis = []
    for pbits, pm in plist:
        for nbits in by_need.get(M2_target - pm, ()):
            basis.append(pbits | nbits)
    basis.sort()
    return basis

# ── Hamiltonian assembly ─────────────────────────────────────────────────────
def popcount_between(bits, i, j):
    """number of set bits strictly between positions i<j"""
    mask = ((1 << j) - 1) & ~((1 << (i + 1)) - 1)
    return bin(bits & mask).count('1')

def assemble(basis, spe_vec, V2):
    """Sparse H in the determinant basis.

    Pair convention: for p = (i<j), P_p = a_j a_i so P†_p|0> = a†_i a†_j|0> = |p>,
    and H2 = sum_{q,p} V2[q,p] P†_q P_p. Signs, measured with bit counts:
      P_p|D>   = (-1)^{n_<i(D) + n_<j(D) - 1} |D \ {i,j}>      (i<j, both occupied)
      P†_q|B>  = (-1)^{n_<r(B) + n_<s(B)}     |B + {r,s}>       (r<s, both empty)
    (OpenFCI's MatrixMachine does exactly this bookkeeping, in 512-bit registers.)
    """
    idx = {d: k for k, d in enumerate(basis)}
    dim = len(basis)
    conn = [np.nonzero(np.abs(V2[p]) > 1e-12)[0] for p in range(len(PAIRS))]
    rows, cols, vals = [], [], []
    for k, det in enumerate(basis):
        occ = [i for i in range(NSP) if det >> i & 1]
        acc = {k: sum(spe_vec[i] for i in occ)}
        for (i, j) in combinations(occ, 2):
            p = PAIR_IDX[(i, j)]
            s_ket = (-1) ** (bin(det & ((1 << i) - 1)).count('1')
                             + bin(det & ((1 << j) - 1)).count('1') - 1)
            base = det & ~(1 << i) & ~(1 << j)
            for q in conn[p]:
                r, s = PAIRS[q]
                if (base >> r & 1) or (base >> s & 1):
                    continue
                newdet = base | (1 << r) | (1 << s)
                kk = idx.get(newdet)
                if kk is None or kk < k:
                    continue
                s_bra = (-1) ** (bin(base & ((1 << r) - 1)).count('1')
                                 + bin(base & ((1 << s) - 1)).count('1'))
                v = V2[q, p] * s_ket * s_bra
                if abs(v) > 1e-14:
                    acc[kk] = acc.get(kk, 0.0) + v
        for kk, v in acc.items():
            rows.append(k); cols.append(kk); vals.append(v)
            if kk != k:
                rows.append(kk); cols.append(k); vals.append(v)
    return coo_matrix((vals, (rows, cols)), shape=(dim, dim)).tocsr()


COSMO = [-87.1044, -85.6021, -82.9883, -82.7320, -82.0302,
         -81.2219, -79.7662, -79.6227, -79.3076, -79.2863]
BIGSTICK_GS = -87.10445
BIGSTICK_EX = [1.502, 4.116, 4.372]          # 2+(1), 2+(2), 4+(1)

spe, tbme, tail = read_int('usdb.int')
a_core, a_ref, power = tail[-3], tail[-2], tail[-1]
scale = (a_ref / (a_core + 8)) ** power
print(f"24Mg: 4p + 4n; TBME scale = ({a_ref:.0f}/{a_core + 8:.0f})^{power} = {scale:.6f}")

t0 = time.time()
spe_vec = np.array([spe[SP[i][0]] for i in range(NSP)])
V2 = build_pair_hamiltonian(tbme, scale)
basis = build_basis(4, 4, 0)
print(f"m-scheme basis dimension (Jz = 0): {len(basis)}   (BIGSTICK: 28,503)")
H = assemble(basis, spe_vec, V2)
E = np.sort(eigsh(H, k=10, which='SA', return_eigenvectors=False))
dt = time.time() - t0

print(f"\n{'this code':>12}  {'cosmo':>10}")
for e, ref in zip(E, COSMO):
    d = abs(e - ref)
    mark = "ok" if d < 5e-4 else f"delta {d * 1000:.0f} meV (their solver tolerance)"
    print(f"{e:12.4f}  {ref:10.4f}   {mark}")
print(f"\nBIGSTICK ground state: {BIGSTICK_GS}   (this code: {E[0]:.4f})")
for ex, (i, lab) in zip(BIGSTICK_EX, [(1, '2+(1)'), (2, '2+(2)'), (3, '4+(1)')]):
    print(f"BIGSTICK {lab} excitation: {ex:.3f} MeV   (this code: {E[i] - E[0]:.3f})")
print(f"\nruntime: {dt:.1f} s single-core")