Bayesian Nonparametric Inference of the Interaction Potential in the Kinetic Vlasov–Fokker–Planck Model

Companion notebook to the thesis "Bayesian Nonparametric Inference in the Kinetic Vlasov–Fokker–Planck Model". All numerical experiments of Section 4 are reproduced here. The section numbering below follows that of the thesis.

Luca Martucci — MSc thesis, 2026.

3.1 A Hermite–Fourier spectral discretisation of (VFP)

Setup, the phase-space grid and Maxwellian-adapted Hermite basis, and the coefficient-hierarchy solver with Strang splitting.

import time, math
import numpy as np
import matplotlib.pyplot as plt

from cycler import cycler
plt.rcParams['axes.prop_cycle'] = cycler(color=['#0000CD', '#D62728', '#2CA02C', '#FF7F0E'])

from scipy.linalg import expm
from numpy.polynomial.hermite_e import hermeval

%config InlineBackend.figure_format = 'retina'
plt.rcParams.update({
    "figure.dpi": 150, "savefig.dpi": 300, "savefig.bbox": "tight",
    "axes.grid": False, "font.size": 11,
})
rng_global = np.random.default_rng(0)

def make_grid(Nq, L=2*np.pi):
    q  = np.arange(Nq) * (L / Nq)
    dq = L / Nq
    k  = 2*np.pi * np.fft.fftfreq(Nq, d=dq)
    return q, dq, k

def hermite_eval(NH, p, beta):
    # B[i, n] = e_n(p_i) = sqrt(beta)/sqrt(2 pi n!) He_n(sqrt(beta) p) exp(-beta p^2 / 2)
    xi    = np.sqrt(beta) * np.atleast_1d(p)
    gauss = np.sqrt(beta) / np.sqrt(2*np.pi) * np.exp(-0.5*xi**2)
    B = np.zeros((len(xi), NH))
    for n in range(NH):
        c = np.zeros(n+1); c[n] = 1.0
        B[:, n] = hermeval(xi, c) / np.sqrt(float(math.factorial(n))) * gauss
    return B

# quick checks of the two defining properties of e_n
q, dq, k = make_grid(48)
pp = np.linspace(-12, 12, 4000); beta_chk = 1.0
B = hermite_eval(6, pp, beta_chk)
ints = np.trapz(B, pp, axis=0)                  # should be (1,0,0,0,0,0)
print("int e_n dp      =", np.round(ints, 4))
print("max |e_0 - M|   =", np.max(np.abs(B[:,0] -
      np.sqrt(beta_chk/(2*np.pi))*np.exp(-0.5*beta_chk*pp**2))))

int e_n dp      = [ 1. -0. -0. -0.  0. -0.]
max |e_0 - M|   = 0.0

def build_linear_exponentials(k, NH, gamma, beta, dt):
    Ehalf = np.empty((len(k), NH, NH), dtype=complex)
    Efull = np.empty((len(k), NH, NH), dtype=complex)
    n = np.arange(NH)
    for ik, kk in enumerate(k):
        L = np.zeros((NH, NH), dtype=complex)
        L[n, n] = -gamma * n                                  # friction (diagonal)
        for m in range(NH):                                   # streaming (tridiagonal)
            if m-1 >= 0:  L[m, m-1] += -(1j*kk/np.sqrt(beta)) * np.sqrt(m)
            if m+1 < NH:  L[m, m+1] += -(1j*kk/np.sqrt(beta)) * np.sqrt(m+1)
        Ehalf[ik] = expm(0.5*dt*L)
        Efull[ik] = expm(dt*L)
    return Ehalf, Efull

def force_kick(c, A, tau, beta):
    # exact exponential of the nilpotent force generator (N c)_n = sqrt(beta n) A c_{n-1}
    NH = c.shape[0]
    sb  = np.sqrt(beta * np.arange(NH))
    out, term = c.copy(), c.copy()
    for j in range(1, NH):                                    # nilpotent => finite sum
        new = np.zeros_like(term)
        new[1:] = sb[1:, None] * A[None, :] * term[:-1]
        term = new * (tau / j)
        out += term
        if not np.any(term): break
    return out

class KineticVFP:
    def __init__(self, q, dq, k, NH, gamma, beta, dt, V):
        self.q, self.dq, self.k = q, dq, k
        self.NH, self.gamma, self.beta, self.dt = NH, gamma, beta, dt
        self.dV = np.real(np.fft.ifft(1j*k*np.fft.fft(V)))                # dV/dq
        self.Ehalf, self.Efull = build_linear_exponentials(k, NH, gamma, beta, dt)

    def force_field(self, rho, W):
        Wconv_hat = self.dq * np.fft.fft(W) * np.fft.fft(rho)            # W * rho
        dWconv    = np.real(np.fft.ifft(1j*self.k*Wconv_hat))
        return -self.dV - dWconv                                          # A = -d/dq (V + W*rho)

    def linear_step(self, c, E):
        chat = np.fft.fft(c, axis=1)                                      # (NH, Nq)
        chat = np.einsum('kmn,nk->mk', E, chat)                           # apply E_k per wavenumber
        return np.real(np.fft.ifft(chat, axis=1))

    def step(self, c, W):
        dt = self.dt
        c = force_kick(c, self.force_field(c[0], W), dt/2, self.beta)
        c = self.linear_step(c, self.Efull)
        c = force_kick(c, self.force_field(c[0], W), dt/2, self.beta)
        return c

    def solve(self, W, c_init, Nt, save_every=1):
        c = c_init.copy()
        saved, times = [c.copy()], [0.0]
        for nstep in range(Nt):
            c = self.step(c, W)
            if (nstep+1) % save_every == 0:
                saved.append(c.copy()); times.append((nstep+1)*self.dt)
        return np.array(times), np.array(saved)        # saved: (Ns, NH, Nq)

3.1.4 Validation of the solver

Four independent checks: (1) exact friction spectrum, (2)–(3) mass conservation, equilibrium stationarity and $O(\Delta t^2)$ consistency, and (4) an independent interacting-particle cross-check.

# (1) friction eigenstructure: c_n(t) = exp(-gamma n t) c_n(0)
Nq, NH = 48, 14
gamma, beta = 1.0, 1.0
q, dq, k = make_grid(Nq)
solver0 = KineticVFP(q, dq, k, NH, gamma, beta, dt=0.05, V=np.zeros(Nq))
c0 = np.zeros((NH, Nq))
c0[0], c0[1], c0[2], c0[3] = 1/(2*np.pi), 0.3/(2*np.pi), 0.2/(2*np.pi), 0.1/(2*np.pi)
times, sol = solver0.solve(np.zeros(Nq), c0, Nt=20, save_every=2)
print("Friction relaxation (uniform in q):")
for n in [1, 2, 3]:
    num   = sol[:, n, :].mean(axis=1) * 2*np.pi
    exact = c0[n, 0] * 2*np.pi * np.exp(-gamma*n*times)
    print(f"  mode n={n}: max rel. error vs exp(-{n}t) = "
          f"{np.max(np.abs(num-exact)/np.abs(exact+1e-15)):.2e}")

Friction relaxation (uniform in q):
  mode n=1: max rel. error vs exp(-1t) = 3.75e-16
  mode n=2: max rel. error vs exp(-2t) = 2.44e-15
  mode n=3: max rel. error vs exp(-3t) = 7.08e-16

# (2)+(3) mass conservation, equilibrium stationarity, O(dt^2)
V  = -np.cos(q)
W  = -np.cos(q) + 0.5*np.cos(2*q)

# self-consistent McKean-Vlasov equilibrium density via Picard iteration
rho = np.ones(Nq) / (2*np.pi)
for _ in range(500):
    Wconv = np.real(np.fft.ifft(dq*np.fft.fft(W)*np.fft.fft(rho)))
    new   = np.exp(-beta*(V + Wconv)); new /= new.sum()*dq
    if np.max(np.abs(new-rho)) < 1e-14: break
    rho = new
c_eq = np.zeros((NH, Nq)); c_eq[0] = rho

solver = KineticVFP(q, dq, k, NH, gamma, beta, dt=0.02, V=V)
rho0 = 1 + 0.6*np.sin(q) + 0.3*np.cos(2*q); rho0 = np.clip(rho0, 1e-3, None); rho0 /= rho0.sum()*dq
c_init = np.zeros((NH, Nq)); c_init[0] = rho0
times, sol = solver.solve(W, c_init, Nt=100, save_every=10)
print("Mass over time:", np.round(sol[:, 0, :].sum(axis=1)*dq, 12))

print("\nEquilibrium drift vs dt (Strang 2nd order):")
for dt in [0.08, 0.04, 0.02, 0.01]:
    s = KineticVFP(q, dq, k, NH, gamma, beta, dt, V)
    _, traj = s.solve(W, c_eq, Nt=int(1.6/dt), save_every=10)
    print(f"  dt={dt:.3f}  max|c(t)-c_eq| = {np.max(np.abs(traj-c_eq[None])):.3e}")

Mass over time: [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]

Equilibrium drift vs dt (Strang 2nd order):
  dt=0.080  max|c(t)-c_eq| = 4.772e-04
  dt=0.040  max|c(t)-c_eq| = 1.194e-04
  dt=0.020  max|c(t)-c_eq| = 2.985e-05
  dt=0.010  max|c(t)-c_eq| = 7.463e-06

# (4) independent particle (IPS) cross-check
dt = 0.01
solver = KineticVFP(q, dq, k, NH, gamma, beta, dt, V)
T = 1.5; Nt = int(T/dt)
_, sol = solver.solve(W, c_init, Nt, save_every=Nt)
rho_pde = sol[-1, 0, :]

rng = np.random.default_rng(0)
n_part = 200_000
cdf = np.cumsum(rho0*dq); cdf /= cdf[-1]
qj = np.interp(rng.random(n_part), cdf, q)        # q0 ~ rho0
pj = rng.standard_normal(n_part) / np.sqrt(beta)  # p0 ~ Maxwellian
W_hat = np.fft.fft(W)
for _ in range(Nt):
    hist, _ = np.histogram(qj % (2*np.pi), bins=Nq, range=(0, 2*np.pi), density=True)
    dWconv  = np.real(np.fft.ifft(1j*k*dq*W_hat*np.fft.fft(hist)))
    idx     = np.minimum((qj % (2*np.pi))/dq, Nq-1).astype(int)
    force   = -np.sin(qj) - dWconv[idx]           # -dV - d(W*rho),  V=-cos
    pj += dt*force - gamma*pj*dt + np.sqrt(2*gamma/beta*dt)*rng.standard_normal(n_part)
    qj  = (qj + pj*dt) % (2*np.pi)
hist_T, _ = np.histogram(qj, bins=Nq, range=(0, 2*np.pi), density=True)
errL2 = np.sqrt(np.trapz((hist_T - rho_pde)**2, q))
print(f"||rho_particle(T) - rho_PDE(T)||_L2 = {errL2:.4f}   "
      f"(MC histogram noise level, rho scale ~ {rho_pde.mean():.3f})")

fig, ax = plt.subplots(figsize=(7, 3.4))
ax.plot(q, rho0,    'k:',   lw=1.5, label=r'$\rho_0$')
ax.plot(q, rho_pde, 'C0',   lw=2,   label='PDE (Hermite-Fourier)')
ax.plot(q, hist_T,  'C3o',  ms=3,   label=f'{n_part:,} particles (IPS)')
ax.set_title("Independent particle cross-check at $t=T$"); ax.set_xlabel("q"); ax.legend()
plt.tight_layout(); plt.show()

||rho_particle(T) - rho_PDE(T)||_L2 = 0.0187   (MC histogram noise level, rho scale ~ 0.159)

3.1.5 Hypocoercivity in the discretised generator

# Spectrum of the linear streaming+friction operator L_k, wavenumber by wavenumber.
NH_spec, gamma_spec, beta_spec = 14, 1.0, 1.0

def Lk_matrix(kk, NH, gamma, beta):
    L = np.zeros((NH, NH), dtype=complex)
    for m in range(NH):
        L[m, m] = -gamma * m                                       # friction (diagonal)
        if m-1 >= 0:  L[m, m-1] += -(1j*kk/np.sqrt(beta)) * np.sqrt(m)     # streaming
        if m+1 < NH:  L[m, m+1] += -(1j*kk/np.sqrt(beta)) * np.sqrt(m+1)
    return L

k_scan = np.arange(0, 9)
slowest = np.array([np.linalg.eigvals(Lk_matrix(kk, NH_spec, gamma_spec, beta_spec)).real.max()
                    for kk in k_scan])
print("k                :", k_scan)
print("slowest Re(lambda):", np.round(slowest, 3))
print(f"\nhypocoercive gap opened by streaming: {0.0:.2f} at k=0  ->  {-slowest[1]:.2f} at k=1")

# (top) least-damped rate vs k
fig, ax = plt.subplots(figsize=(6, 3.8))
ax.plot(k_scan, -slowest, 'o-', color='C0')
ax.axhline(0, color='k', lw=0.5)
ax.set_xlabel("spatial wavenumber $k$"); ax.set_ylabel(r"relaxation rate $-\max_j\,\mathrm{Re}\,\lambda_j$")
ax.set_title("Hypocoercive gap: density relaxes only for $k\\neq0$")
plt.tight_layout(); plt.show()

# (bottom) full complex spectrum at a few wavenumbers
fig, ax = plt.subplots(figsize=(6, 3.8))
for kk, col in zip([0, 1, 3, 6], ['C3', 'C0', 'C2', 'C1']):
    ev = np.linalg.eigvals(Lk_matrix(kk, NH_spec, gamma_spec, beta_spec))
    ax.scatter(ev.real, ev.imag, s=28, color=col, label=f"$k={kk}$")
ax.axvline(0, color='k', lw=0.5)
ax.set_xlabel(r"$\mathrm{Re}\,\lambda$ (decay rate)"); ax.set_ylabel(r"$\mathrm{Im}\,\lambda$ (oscillation)")
ax.set_title("Spectrum of $L_k$ in the complex plane"); ax.legend(fontsize=8)
plt.tight_layout(); plt.show()

k                : [0 1 2 3 4 5 6 7 8]
slowest Re(lambda): [-0.    -1.    -3.674 -4.112 -4.228 -4.278 -4.304 -4.32  -4.33 ]

hypocoercive gap opened by streaming: 0.00 at k=0  ->  1.00 at k=1

3.2 Synthetic data and the discrete measurement models

Ground-truth interaction $W_0$, the initial state $\phi=\rho_0\otimes M$, and the discrete observation designs (generation on a finer resolution to avoid the inverse crime).

def generate_data(seed=0):
    rng = np.random.default_rng(seed)
    Nq, NH_hi, gamma, beta = 48, 20, 1.0, 1.0
    T, dt_hi = 2.0, 0.01
    q, dq, k = make_grid(Nq)
    V = -np.cos(q)

    W_true = -np.cos(q) + 0.5*np.cos(2*q) - 0.25*np.cos(3*q)
    rho0   = 1 + 0.6*np.sin(q) + 0.3*np.cos(2*q)
    rho0   = np.clip(rho0, 1e-3, None); rho0 /= rho0.sum()*dq
    c_init = np.zeros((NH_hi, Nq)); c_init[0] = rho0

    solver_hi = KineticVFP(q, dq, k, NH_hi, gamma, beta, dt_hi, V)
    times, sol = solver_hi.solve(W_true, c_init, Nt=int(T/dt_hi), save_every=2)  # snapshots @ 0.02

    N_obs, sigma = 3000, 0.004
    pmax = 3.5 / np.sqrt(beta)
    t_idx = rng.integers(1, len(times), N_obs)       # exclude t=0
    q_idx = rng.integers(0, Nq, N_obs)
    p_obs = rng.uniform(-pmax, pmax, N_obs)
    B = hermite_eval(NH_hi, p_obs, beta)
    f_clean = np.einsum('in,in->i', sol[t_idx, :, q_idx], B)
    Y = f_clean + sigma*rng.standard_normal(N_obs)

    return dict(q=q, dq=dq, k=k, T=T, gamma=gamma, beta=beta, V=V, rho0=rho0,
                W_true=W_true, times=times, sol_hi=sol, c_init=c_init,
                t_idx=t_idx, q_idx=q_idx, p_obs=p_obs, Y=Y, sigma=sigma, f_clean=f_clean)

data = generate_data(0)
print(f"snapshots={len(data['times'])} (dt=0.02 to T={data['T']}),  "
      f"N_obs={len(data['Y'])},  sigma={data['sigma']}")
print(f"f range = [{data['f_clean'].min():.3f}, {data['f_clean'].max():.3f}],  "
      f"noise/signal ~ {data['sigma']/data['f_clean'].std():.2f}")
print(f"mass t=0: {data['sol_hi'][0,0].sum()*data['dq']:.10f}   "
      f"mass t=T: {data['sol_hi'][-1,0].sum()*data['dq']:.10f}")

snapshots=101 (dt=0.02 to T=2.0),  N_obs=3000,  sigma=0.004
f range = [0.000, 0.113],  noise/signal ~ 0.15
mass t=0: 1.0000000000   mass t=T: 1.0000000000

3.2.1 Visualising the ground-truth trajectory

q = data['q']; ts = data['times']
P = np.linspace(-4, 4, 120)
Bp = hermite_eval(data['sol_hi'].shape[1], P, data['beta'])

def phase_density(coeff_qt):                       # f(q,p) from coeffs c_n(q)
    return (coeff_qt.T @ Bp.T)                      # (Nq, NP)

fig, ax = plt.subplots(figsize=(5, 3.6))
ax.plot(q, data['W_true'], lw=2)
ax.set_title("Symmetric ground-truth $W_0(q)$"); ax.set_xlabel("q")
plt.tight_layout(); plt.show()

f0 = phase_density(data['sol_hi'][0])
fT = phase_density(data['sol_hi'][-1])
fig, ax = plt.subplots(figsize=(5, 3.6))
im = ax.imshow(f0.T, origin='lower', aspect='auto', extent=[0, 2*np.pi, -4, 4], cmap='viridis')
ax.set_title("$f(0,q,p)=\\rho_0(q)M(p)$"); ax.set_xlabel("q"); ax.set_ylabel("p")
plt.colorbar(im, ax=ax)
plt.tight_layout(); plt.show()

fig, ax = plt.subplots(figsize=(5, 3.6))
im = ax.imshow(fT.T, origin='lower', aspect='auto', extent=[0, 2*np.pi, -4, 4], cmap='viridis')
ax.set_title(f"$f(T,q,p)$  (T={data['T']})"); ax.set_xlabel("q"); ax.set_ylabel("p")
plt.colorbar(im, ax=ax)
plt.tight_layout(); plt.show()

# density rho(t,q) and momentum-skew (c_1) heatmaps
fig, ax = plt.subplots(figsize=(5.5, 3.4))
im = ax.imshow(data['sol_hi'][:, 0, :], origin='lower', aspect='auto',
               extent=[0, 2*np.pi, 0, data['T']], cmap='magma')
ax.set_title(r"$\rho(t,q)=c_0$"); ax.set_xlabel("q"); ax.set_ylabel("t"); plt.colorbar(im, ax=ax)
plt.tight_layout(); plt.show()

fig, ax = plt.subplots(figsize=(5.5, 3.4))
im = ax.imshow(data['sol_hi'][:, 1, :], origin='lower', aspect='auto',
               extent=[0, 2*np.pi, 0, data['T']], cmap='coolwarm')
ax.set_title(r"momentum flux mode $c_1(t,q)$"); ax.set_xlabel("q"); ax.set_ylabel("t"); plt.colorbar(im, ax=ax)
plt.tight_layout(); plt.show()

3.3 Empirical forward stability: $W\mapsto f_W$ is locally Lipschitz

inf_kw = dict(NH=14, dt=0.02)                      # inference-resolution solver
solver = KineticVFP(data['q'], data['dq'], data['k'], inf_kw['NH'],
                    data['gamma'], data['beta'], inf_kw['dt'], data['V'])
c_init_inf = np.zeros((inf_kw['NH'], len(q))); c_init_inf[0] = data['rho0']
Nt_inf = int(data['T']/inf_kw['dt'])
Bp_inf = hermite_eval(inf_kw['NH'], P, data['beta'])

rng = np.random.default_rng(42)
_, base = solver.solve(data['W_true'], c_init_inf, Nt_inf, save_every=Nt_inf)
fbase = base[-1].T @ Bp_inf.T

eps_list = np.array([1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1]); errs = []
for eps in eps_list:
    dW = sum(eps*rng.standard_normal()/m**2*np.cos(m*q) for m in range(1, 6))
    _, pert = solver.solve(data['W_true']+dW, c_init_inf, Nt_inf, save_every=Nt_inf)
    fp = pert[-1].T @ Bp_inf.T
    errs.append(np.sqrt(np.trapz(np.trapz((fp-fbase)**2, P, axis=1), q)))
errs = np.array(errs)
slope, intercept = np.polyfit(np.log(eps_list), np.log(errs), 1)
print(f"empirical log-log slope = {slope:.2f}  (Lipschitz: slope <= 1)")

plt.figure(figsize=(5.2, 3.8))
plt.loglog(eps_list, errs, 'o-', label=f"data, slope $\\approx$ {slope:.2f}")
plt.loglog(eps_list, np.exp(intercept)*eps_list, 'k--', label="reference slope 1")
plt.xlabel(r"perturbation amplitude $\epsilon$"); plt.ylabel(r"$\|f_{W+\epsilon\delta}-f_W\|_{L^2(q,p)}$")
plt.title("Forward map is locally Lipschitz"); plt.legend(); plt.tight_layout(); plt.show()

empirical log-log slope = 0.87  (Lipschitz: slope <= 1)

3.4 The prior on $W$

def build_basis_cos(q, M):
    return np.stack([np.cos(m*q) for m in range(1, M+1)], axis=1)   # (Nq, M)

def make_scaling(M, tau=1.0, s=1.5):
    return np.array([tau/m**s for m in range(1, M+1)])

def W_from_xi(xi, Phi, scale):
    return Phi @ (scale*xi)

M, tau, s = 6, 1.0, 1.5
Phi   = build_basis_cos(data['q'], M)
scale = make_scaling(M, tau, s)
rng = np.random.default_rng(1)
plt.figure(figsize=(7, 3.4))
for _ in range(6):
    plt.plot(data['q'], W_from_xi(rng.standard_normal(M), Phi, scale), alpha=0.8)
plt.plot(data['q'], data['W_true'], 'k', lw=2.5, label='$W_0$')
plt.title(f"Samples from the even GP prior (M={M}, $\\tau$={tau}, s={s})")
plt.xlabel("q"); plt.ylabel("W(q)"); plt.legend(); plt.tight_layout(); plt.show()

3.5 Likelihood and posterior sampling by pCN

One likelihood evaluation costs exactly one forward PDE solve; the posterior is sampled by the preconditioned Crank–Nicolson algorithm in whitened coordinates.

B_obs = hermite_eval(inf_kw['NH'], data['p_obs'], data['beta'])   # (N_obs, NH)

def log_likelihood(W):
    _, sol = solver.solve(W, c_init_inf, Nt_inf, save_every=1)     # snapshots @ dt=0.02
    pred = np.einsum('in,in->i', sol[data['t_idx'], :, data['q_idx']], B_obs)
    return -0.5*np.sum((data['Y']-pred)**2)/data['sigma']**2

W0 = W_from_xi(np.zeros(M), Phi, scale)
t0 = time.time(); ll0 = log_likelihood(W0)
print(f"one likelihood eval: {time.time()-t0:.3f}s,   ll(W=0) = {ll0:.1f}")
print(f"target plateau ~ -N_obs/2 = {-len(data['Y'])/2:.0f}")

one likelihood eval: 0.026s,   ll(W=0) = -2502.1
target plateau ~ -N_obs/2 = -1500

def pcn(n_iter=2500, beta_pcn=0.12, seed=1, log_every=400):
    rng = np.random.default_rng(seed)
    xi = np.zeros(M); ll = log_likelihood(W_from_xi(xi, Phi, scale))
    samples = np.empty((n_iter, M)); ll_chain = np.empty(n_iter); n_acc = 0; t0 = time.time()
    for it in range(n_iter):
        eta  = rng.standard_normal(M)
        xi_p = np.sqrt(1-beta_pcn**2)*xi + beta_pcn*eta
        ll_p = log_likelihood(W_from_xi(xi_p, Phi, scale))
        if np.log(rng.random()) < ll_p - ll:
            xi, ll = xi_p, ll_p; n_acc += 1
        samples[it], ll_chain[it] = xi, ll
        if (it+1) % log_every == 0:
            print(f"  iter {it+1:5d}/{n_iter}  acc={n_acc/(it+1):6.2%}  "
                  f"ll={ll:9.1f}  elapsed={time.time()-t0:6.1f}s")
    return samples, ll_chain, n_acc/n_iter

n_iter = 2500
samples, ll_chain, acc = pcn(n_iter=n_iter, beta_pcn=0.12, seed=1)
print(f"\nFinal acceptance rate: {acc:.2%}  (target 20-50%)")

  iter   400/2500  acc=20.50%  ll=  -1534.9  elapsed=   8.8s
  iter   800/2500  acc=19.25%  ll=  -1538.9  elapsed=  17.7s
  iter  1200/2500  acc=19.42%  ll=  -1534.3  elapsed=  26.2s
  iter  1600/2500  acc=18.75%  ll=  -1536.2  elapsed=  34.8s
  iter  2000/2500  acc=19.90%  ll=  -1535.7  elapsed=  43.4s
  iter  2400/2500  acc=20.08%  ll=  -1537.1  elapsed=  52.1s

Final acceptance rate: 20.00%  (target 20-50%)

3.5.1 Posterior inference from phase-space observations

burn = n_iter // 5
xi_post = samples[burn:]
W_mean  = W_from_xi(xi_post.mean(axis=0), Phi, scale)

_, sol_hat = solver.solve(W_mean, c_init_inf, Nt_inf, save_every=1)
err_W   = np.sqrt(np.trapz((W_mean - data['W_true'])**2, data['q']))
err_rho = np.sqrt(np.trapz((sol_hat[-1,0] - data['sol_hi'][-1,0])**2, data['q']))
print(f"||W_bar - W_true||_L2      = {err_W:.4f}")
print(f"||rho_bar(T) - rho_0(T)||_L2 = {err_rho:.4f}   (<< sigma = {data['sigma']})")

||W_bar - W_true||_L2      = 0.3019
||rho_bar(T) - rho_0(T)||_L2 = 0.0011   (<< sigma = 0.004)

rng_p = np.random.default_rng(7)

# (a) potential reconstruction with credible band
fig, ax = plt.subplots(figsize=(6.5, 4.2))
for j in rng_p.choice(len(xi_post), 80, replace=False):
    ax.plot(data['q'], W_from_xi(xi_post[j], Phi, scale), color='gray', alpha=0.12)
ax.plot(data['q'], data['W_true'], lw=2.4, label='True $W_0$')
ax.plot(data['q'], W_mean, '--', lw=2.4, color='C3', label=r'Posterior mean $\bar W$')
ax.set_title("Interaction potential — reconstruction"); ax.set_xlabel("q"); ax.legend()
plt.tight_layout(); plt.show()

# (b) phase-space marginals: rho(T,q) and momentum profile
fig, ax = plt.subplots(figsize=(6.5, 4.2))
rho_true_T = data['sol_hi'][-1, 0]; rho_hat_T = sol_hat[-1, 0]
ax.plot(data['q'], rho_true_T, 'k', lw=2, label=r'true $\rho(T)$')
ax.plot(data['q'], rho_hat_T, 'C3--', lw=1.6, label=r'reconstructed $\rho(T)$')
ax.set_title("Spatial density at $T$: truth vs reconstruction"); ax.set_xlabel("q"); ax.legend(fontsize=9)
plt.tight_layout(); plt.show()

# (c) log-likelihood trace
fig, ax = plt.subplots(figsize=(6.5, 4.2))
ax.plot(ll_chain, lw=0.7); ax.axvline(burn, color='C3', ls=':', label='burn-in')
ax.axhline(-len(data['Y'])/2, color='C2', ls='--', lw=1, label=r'$-N/2$')
ax.set_title("Log-likelihood trace"); ax.set_xlabel("iteration"); ax.legend(fontsize=9)
plt.tight_layout(); plt.show()

# (d) cosine coefficients vs truth
fig, ax = plt.subplots(figsize=(6.5, 4.2))
coef_chain = xi_post * scale
coef_mean, coef_std = coef_chain.mean(axis=0), coef_chain.std(axis=0)
coef_true = np.array([np.trapz(np.cos(m*data['q'])*data['W_true'], data['q'])/np.pi
                      for m in range(1, M+1)])
pos = np.arange(1, M+1)
ax.errorbar(pos, coef_mean, yerr=coef_std, fmt='o', capsize=3, label=r'posterior mean $\pm$ std')
ax.plot(pos, coef_true, 'rx', ms=11, mew=2.2, label='true coefficient')
ax.axhline(0, color='k', lw=0.4); ax.set_xticks(pos)
ax.set_xticklabels([f"$a_{m}$" for m in range(1, M+1)])
ax.set_title("Cosine coefficients of $W$"); ax.set_xlabel("mode $m$"); ax.legend()
plt.tight_layout(); plt.show()

print("true coeffs:", np.round(coef_true, 3))
print("post coeffs:", np.round(coef_mean, 3), " +-", np.round(coef_std, 3))

true coeffs: [-0.969  0.531 -0.22   0.029  0.028  0.027]
post coeffs: [-0.987  0.548 -0.291 -0.149  0.021 -0.055]  +- [0.022 0.06  0.082 0.07  0.052 0.055]

3.6 Identifiability: the Jacobian ladder and reversion to the prior

Linearising the forward map at the truth: the SVD of the Jacobian predicts, mode by mode, which Fourier frequencies of $W$ the data identify.

# Linearise the forward map at W_true and read off the identifiability structure.
xi_true = coef_true / scale          # true potential in whitened coordinates

def predict_f(xi):
    _, sol = solver.solve(W_from_xi(xi, Phi, scale), c_init_inf, Nt_inf, save_every=1)
    return np.einsum('in,in->i', sol[data['t_idx'], :, data['q_idx']], B_obs)

h = 1e-3
F0 = predict_f(xi_true)
J = np.zeros((len(F0), M))
for m in range(M):
    e = np.zeros(M); e[m] = h
    J[:, m] = (predict_f(xi_true + e) - predict_f(xi_true - e)) / (2*h)

U, S, Vt = np.linalg.svd(J, full_matrices=False)
print("singular values of J   :", np.array2string(S, formatter={'float_kind':lambda v:f'{v:.2e}'}))
print("ratio to the largest    :", np.array2string(S/S[0], formatter={'float_kind':lambda v:f'{v:.1e}'}))
print("\nEach singular direction is essentially a single cosine mode:")
for j in range(M):
    dom = np.argmax(np.abs(Vt[j]))
    print(f"  sigma_{j} = {S[j]:.2e}   <-->   mode a_{dom+1}  (weight {Vt[j,dom]:+.2f})")

# a mode is 'learnable' when its per-dataset curvature (sigma_j/sigma_eps)^2 beats the prior precision 1/scale^2
info   = (S / data['sigma'])**2                  # ~ Fisher info along each direction (per full data set)
prior  = 1.0 / np.ones(M)**2                     # whitened prior precision is 1 in every direction
fig, ax = plt.subplots(figsize=(6.4, 3.8))
ax.semilogy(np.arange(1, M+1), info, 'o-', label=r'data curvature $(\sigma_j/\sigma_\varepsilon)^2$')
ax.axhline(1.0, color='C3', ls='--', label='whitened prior precision (=1)')
ax.set_xlabel("singular direction $\\approx$ mode $m$"); ax.set_ylabel("information (log scale)")
ax.set_title("Identifiability ladder: data vs prior, mode by mode")
ax.legend(); plt.tight_layout(); plt.show()

singular values of J   : [1.73e-01 2.71e-02 1.96e-03 1.21e-04 1.16e-05 5.60e-07]
ratio to the largest    : [1.0e+00 1.6e-01 1.1e-02 7.0e-04 6.7e-05 3.2e-06]

Each singular direction is essentially a single cosine mode:
  sigma_0 = 1.73e-01   <-->   mode a_1  (weight +1.00)
  sigma_1 = 2.71e-02   <-->   mode a_2  (weight -1.00)
  sigma_2 = 1.96e-03   <-->   mode a_3  (weight +1.00)
  sigma_3 = 1.21e-04   <-->   mode a_4  (weight -1.00)
  sigma_4 = 1.16e-05   <-->   mode a_5  (weight -1.00)
  sigma_5 = 5.60e-07   <-->   mode a_6  (weight -1.00)

3.6.1 A ground truth engineered to defeat the method

# fine-solver "truth" for the harder potential
qh, dqh, kh = data['q'], data['dq'], data['k']
Vh = data['V']; g_, b_ = data['gamma'], data['beta']
W_hard = 0.8*np.cos(qh) - 0.6*np.cos(2*qh) + 0.4*np.cos(3*qh) - 0.25*np.cos(4*qh)

sol_hi_h = KineticVFP(qh, dqh, kh, 20, g_, b_, 0.01, Vh)
th, solh = sol_hi_h.solve(W_hard, data['c_init'][:20] if data['c_init'].shape[0]>=20 else
                          np.vstack([data['rho0'][None], np.zeros((19, len(qh)))]),
                          Nt=int(data['T']/0.01), save_every=2)

# observations: same design, full phase space
rngh = np.random.default_rng(7)
Nh, sig_h = 3000, 0.004
ti = rngh.integers(1, len(th), Nh); qi = rngh.integers(0, len(qh), Nh)
pi = rngh.uniform(-3.5, 3.5, Nh)
Bhi = hermite_eval(20, pi, b_)
fH = np.einsum('in,in->i', solh[ti, :, qi], Bhi)
YH = fH + sig_h*rngh.standard_normal(Nh)

# inference at coarse resolution, wider prior M=8
Mh = 8
Phi_h = build_basis_cos(qh, Mh); scale_h = make_scaling(Mh, 1.0, 1.5)
solver_h = KineticVFP(qh, dqh, kh, 14, g_, b_, 0.02, Vh)
c0h = np.zeros((14, len(qh))); c0h[0] = data['rho0']; Nth = int(data['T']/0.02)
B_h = hermite_eval(14, pi, b_)
coef_true_h = np.array([np.trapz(np.cos(m*qh)*W_hard, qh)/np.pi for m in range(1, Mh+1)])

def loglik_h(xi):
    _, s = solver_h.solve(W_from_xi(xi, Phi_h, scale_h), c0h, Nth, save_every=1)
    return -0.5*np.sum((YH - np.einsum('in,in->i', s[ti, :, qi], B_h))**2)/sig_h**2

rngc = np.random.default_rng(3); beta_h = 0.07; n_h = 3000
xi = np.zeros(Mh); ll = loglik_h(xi); samp_h = np.empty((n_h, Mh)); nacc = 0; t0 = time.time()
for it in range(n_h):
    xp = np.sqrt(1-beta_h**2)*xi + beta_h*rngc.standard_normal(Mh)
    lp = loglik_h(xp)
    if np.log(rngc.random()) < lp - ll: xi, ll = xp, lp; nacc += 1
    samp_h[it] = xi
    if (it+1) % 750 == 0: print(f"  iter {it+1}/{n_h}  acc={nacc/(it+1):.0%}  ll={ll:.0f}  t={time.time()-t0:.0f}s")
post_h = samp_h[n_h//5:]*scale_h
cmh, csh = post_h.mean(0), post_h.std(0)
print(f"\n  true coeffs: {np.round(coef_true_h,3)}")
print(f"  post coeffs: {np.round(cmh,3)}  +- {np.round(csh,3)}")

fig, ax = plt.subplots(figsize=(6, 3.9))
ax.plot(qh, W_hard, 'k', lw=2.4, label=r'true $W_{\rm hard}$')
ax.plot(qh, Phi_h@cmh, 'C3--', lw=2.2, label=r'posterior mean $\bar W$')
ax.set_title("Harder potential — reconstruction"); ax.set_xlabel("q"); ax.legend()
plt.tight_layout(); plt.show()

pos = np.arange(1, Mh+1)
fig, ax = plt.subplots(figsize=(6, 3.9))
ax.errorbar(pos, cmh, yerr=csh, fmt='o', capsize=3, label=r'posterior $\pm$ std')
ax.plot(pos, coef_true_h, 'rx', ms=11, mew=2.2, label='truth')
ax.axhline(0, color='k', lw=0.4); ax.set_xticks(pos)
ax.set_xticklabels([f"$a_{m}$" for m in range(1, Mh+1)])
ax.set_title("Cosine coefficients — harder potential"); ax.set_xlabel("mode $m$"); ax.legend(fontsize=8)
plt.tight_layout(); plt.show()

  iter 750/3000  acc=36%  ll=-1436  t=17s
  iter 1500/3000  acc=37%  ll=-1437  t=32s
  iter 2250/3000  acc=38%  ll=-1435  t=49s
  iter 3000/3000  acc=39%  ll=-1435  t=65s

  true coeffs: [ 0.785 -0.615  0.386 -0.264 -0.013 -0.013 -0.012 -0.011]
  post coeffs: [ 0.788 -0.582  0.142  0.01  -0.036 -0.106  0.015  0.052]  +- [0.031 0.034 0.101 0.039 0.075 0.052 0.03  0.034]

3.6.2 Growing the parameter dimension: the fully parametric limit

# sweep the prior truncation M from a few modes up to the fully-parametric limit
def run_pcn_M(Mtry, n_iter=2000, beta_pcn=0.09, seed=1, s_prior=1.5, tau_prior=1.0):
    """Self-contained pCN at prior truncation Mtry (leaves the §12 globals untouched)."""
    Phi_M   = build_basis_cos(data['q'], Mtry)
    scale_M = make_scaling(Mtry, tau_prior, s_prior)
    def loglik(xi):
        _, sol = solver.solve(Phi_M @ (scale_M*xi), c_init_inf, Nt_inf, save_every=1)
        pred = np.einsum('in,in->i', sol[data['t_idx'], :, data['q_idx']], B_obs)
        return -0.5*np.sum((data['Y']-pred)**2)/data['sigma']**2
    rng = np.random.default_rng(seed)
    xi = np.zeros(Mtry); ll = loglik(xi)
    samples = np.empty((n_iter, Mtry)); n_acc = 0
    for it in range(n_iter):
        xi_p = np.sqrt(1-beta_pcn**2)*xi + beta_pcn*rng.standard_normal(Mtry)
        ll_p = loglik(xi_p)
        if np.log(rng.random()) < ll_p - ll:
            xi, ll = xi_p, ll_p; n_acc += 1
        samples[it] = xi
    return dict(samples=samples, scale=scale_M, Phi=Phi_M, acc=n_acc/n_iter)

M_full = len(data['q'])//2 - 1            # every cosine mode the Nq-point grid resolves
M_grid = [3, 6, 12, M_full]
print(f"fully-parametric limit for Nq={len(data['q'])}:  M_full = {M_full}\n")

sweep = {}
t0 = time.time()
for Mtry in M_grid:
    res  = run_pcn_M(Mtry, n_iter=2000, beta_pcn=0.09)
    post = res['samples'][2000//5:]                     # discard first fifth as burn-in
    W_hat = res['Phi'] @ (res['scale'] * post.mean(0))
    res['post'], res['W_hat'] = post, W_hat
    res['errW'] = np.sqrt(np.trapz((W_hat - data['W_true'])**2, data['q']))
    sweep[Mtry] = res
    print(f"  M={Mtry:2d}   acc={res['acc']:5.1%}   ||W_bar - W_true||_L2 = {res['errW']:.4f}")
print(f"\nsweep wall-clock: {time.time()-t0:.0f}s")

fig, ax = plt.subplots(1, 2, figsize=(11, 3.9))
ax[0].plot(M_grid, [sweep[M]['errW'] for M in M_grid], 'o-', color='C0')
ax[0].set_xlabel("prior truncation $M$"); ax[0].set_ylabel(r"$\|\bar W - W_0\|_{L^2}$")
ax[0].set_title("Error is bounded as $M\\to$ fully-parametric"); ax[0].set_ylim(bottom=0)
for M in M_grid:
    ax[1].plot(data['q'], sweep[M]['W_hat'], lw=1.8, label=f"$M={M}$")
ax[1].plot(data['q'], data['W_true'], 'k', lw=2.6, label='truth')
ax[1].set_xlabel("q"); ax[1].set_title(r"Posterior mean $\bar W$ for growing $M$")
ax[1].legend(fontsize=8, ncol=2)
plt.tight_layout(); plt.show()

fully-parametric limit for Nq=48:  M_full = 23

  M= 3   acc=25.6%   ||W_bar - W_true||_L2 = 0.1967
  M= 6   acc=26.6%   ||W_bar - W_true||_L2 = 0.1970
  M=12   acc=27.3%   ||W_bar - W_true||_L2 = 0.2224
  M=23   acc=28.2%   ||W_bar - W_true||_L2 = 0.1684

sweep wall-clock: 171s

3.6.3 How the higher modes return to the prior (Laplace shrinkage and marginals)

# linear-Gaussian (Laplace) posterior: which modes revert to the prior
Phi_full   = build_basis_cos(data['q'], M_full)
scale_full = make_scaling(M_full, 1.0, 1.5)
coef_true_full = np.array([np.trapz(np.cos(m*data['q'])*data['W_true'], data['q'])/np.pi
                           for m in range(1, M_full+1)])
xi_true_full = coef_true_full / scale_full            # truth in whitened coordinates

def predict_full(xi):
    _, sol = solver.solve(Phi_full @ (scale_full*xi), c_init_inf, Nt_inf, save_every=1)
    return np.einsum('in,in->i', sol[data['t_idx'], :, data['q_idx']], B_obs)

h  = 1e-3                                              # forward-map Jacobian by finite differences
F0 = predict_full(xi_true_full)
Jf = np.zeros((len(F0), M_full))
for m in range(M_full):
    e = np.zeros(M_full); e[m] = h
    Jf[:, m] = (predict_full(xi_true_full + e) - predict_full(xi_true_full - e)) / (2*h)

Sigma   = np.linalg.inv(np.eye(M_full) + Jf.T @ Jf / data['sigma']**2)   # whitened Laplace covariance
shrink  = 1 - np.diag(Sigma)                          # posterior shrinkage per mode
prior_sd = scale_full                                 # physical prior std of a_m
post_sd  = scale_full * np.sqrt(np.clip(np.diag(Sigma), 0, None))        # physical posterior std
coef_pcn = sweep[M_full]['post'] * scale_full         # empirical pCN marginals from §13.3

gauss = lambda x, sd: np.exp(-0.5*(x/sd)**2)/(sd*np.sqrt(2*np.pi))

# (1) prior vs posterior marginal, mode by mode (first 12 modes)
fig, axes = plt.subplots(3, 4, figsize=(12, 7.4))
for m, ax in enumerate(axes.ravel()):
    xg = np.linspace(-3*prior_sd[m], 3*prior_sd[m], 400)
    ax.fill_between(xg, gauss(xg, prior_sd[m]), color='C7', alpha=0.30, label='prior')
    ax.plot(xg, gauss(xg, post_sd[m]), 'C3', lw=2, label='posterior')
    ax.axvline(coef_true_full[m], color='k', ls='--', lw=1.4, label='truth')
    ax.set_title(f"$a_{{{m+1}}}$   (shrink {shrink[m]:.2f})", fontsize=10); ax.set_yticks([])
    if m == 0: ax.legend(fontsize=7, loc='upper left')
fig.suptitle("Prior vs posterior, mode by mode: the high modes fall back onto the prior", y=1.01)
plt.tight_layout(); plt.show()

# (2) shrinkage ladder over all modes, with the (noisy) empirical pCN check
emp_shrink = 1 - (coef_pcn.std(0)/prior_sd)**2
fig, ax = plt.subplots(figsize=(7.6, 4))
ax.bar(np.arange(1, M_full+1), shrink, color='C0', alpha=0.75, label='Laplace (from Jacobian)')
ax.plot(np.arange(1, M_full+1), emp_shrink, 'C3o', ms=4, label='empirical pCN (noisy)')
ax.axhline(0, color='k', lw=0.5)
ax.set_xlabel("Fourier mode $m$"); ax.set_ylabel(r"posterior shrinkage $s_m$")
ax.set_title(r"$s_m=1-\mathrm{Var}_{\rm post}/\mathrm{Var}_{\rm prior}$:   1 = learned,   0 = prior")
ax.legend(fontsize=8); plt.tight_layout(); plt.show()

print("mode   :", " ".join(f"{m:5d}" for m in range(1, 9)))
print("shrink :", " ".join(f"{v:5.2f}" for v in shrink[:8]))

mode   :     1     2     3     4     5     6     7     8
shrink :  1.00  0.98  0.19  0.00  0.00  0.00  0.00  0.00

3.7 Density-only observations and the comparison of the two models

The physically natural observable is the macroscopic density $\rho_W$. Below, the inference is redone observing $\rho$ alone, and the two measurement models are compared head to head under a common noise-to-signal ratio. Setup (solver, ground truth, prior, likelihood, sampler) is re-defined self-contained for this part.

import time, math
import numpy as np

import matplotlib.pyplot as plt

from cycler import cycler
plt.rcParams['axes.prop_cycle'] = cycler(color=['#0000CD', '#D62728', '#2CA02C', '#FF7F0E'])

from scipy.linalg import expm
from numpy.polynomial.hermite_e import hermeval

%config InlineBackend.figure_format = 'retina'
plt.rcParams.update({
    "figure.dpi": 150, "savefig.dpi": 300, "savefig.bbox": "tight",
    "axes.grid": False, "font.size": 11,
})

def make_grid(Nq, L=2*np.pi):
    q = np.arange(Nq)*(L/Nq); dq = L/Nq
    k = 2*np.pi*np.fft.fftfreq(Nq, d=dq)
    return q, dq, k

def hermite_eval(NH, p, beta):
    # B[i,n] = e_n(p_i): Maxwellian-weighted Hermite functions (orthonormal, int e_n = delta_n0)
    xi = np.sqrt(beta)*np.atleast_1d(p)
    gauss = np.sqrt(beta)/np.sqrt(2*np.pi)*np.exp(-0.5*xi**2)
    B = np.zeros((len(xi), NH))
    for n in range(NH):
        c = np.zeros(n+1); c[n] = 1.0
        B[:, n] = hermeval(xi, c)/np.sqrt(float(math.factorial(n)))*gauss
    return B

def build_linear_exponentials(k, NH, gamma, beta, dt):
    # per-wavenumber exp(dt L_k) for the streaming(tridiagonal)+friction(diagonal) linear part
    Efull = np.empty((len(k), NH, NH), dtype=complex)
    n = np.arange(NH)
    for ik, kk in enumerate(k):
        L = np.zeros((NH, NH), dtype=complex)
        L[n, n] = -gamma*n
        for m in range(NH):
            if m-1 >= 0: L[m, m-1] += -(1j*kk/np.sqrt(beta))*np.sqrt(m)
            if m+1 < NH: L[m, m+1] += -(1j*kk/np.sqrt(beta))*np.sqrt(m+1)
        Efull[ik] = expm(dt*L)
    return Efull

def force_kick(c, A, tau, beta):
    # exact exponential of the nilpotent force generator (N c)_n = sqrt(beta n) A c_{n-1}
    NH = c.shape[0]; sb = np.sqrt(beta*np.arange(NH))
    out, term = c.copy(), c.copy()
    for j in range(1, NH):
        new = np.zeros_like(term)
        new[1:] = sb[1:, None]*A[None, :]*term[:-1]
        term = new*(tau/j); out += term
        if not np.any(term): break
    return out

class KineticVFP:
    def __init__(self, q, dq, k, NH, gamma, beta, dt, V):
        self.q, self.dq, self.k = q, dq, k
        self.NH, self.gamma, self.beta, self.dt = NH, gamma, beta, dt
        self.dV = np.real(np.fft.ifft(1j*k*np.fft.fft(V)))
        self.Efull = build_linear_exponentials(k, NH, gamma, beta, dt)
    def force_field(self, rho, W):
        Wconv_hat = self.dq*np.fft.fft(W)*np.fft.fft(rho)         # W*rho via FFT
        dWconv = np.real(np.fft.ifft(1j*self.k*Wconv_hat))
        return -self.dV - dWconv                                 # A = -d/dq (V + W*rho)
    def linear_step(self, c):
        chat = np.fft.fft(c, axis=1)
        chat = np.einsum('kmn,nk->mk', self.Efull, chat)
        return np.real(np.fft.ifft(chat, axis=1))
    def step(self, c, W):
        dt = self.dt
        c = force_kick(c, self.force_field(c[0], W), dt/2, self.beta)  # Strang: half kick
        c = self.linear_step(c)                                        #         full linear
        c = force_kick(c, self.force_field(c[0], W), dt/2, self.beta)  #         half kick
        return c
    def solve(self, W, c_init, Nt, save_every=1):
        c = c_init.copy(); saved, times = [c.copy()], [0.0]
        for nstep in range(Nt):
            c = self.step(c, W)
            if (nstep+1) % save_every == 0:
                saved.append(c.copy()); times.append((nstep+1)*self.dt)
        return np.array(times), np.array(saved)   # (Ns, NH, Nq)

Nq, NH, gamma, beta = 48, 14, 1.0, 1.0
q, dq, k = make_grid(Nq)

# (i) mass conservation on a non-trivial run
V = -np.cos(q); W = -np.cos(q) + 0.5*np.cos(2*q)
rho0 = 1 + 0.6*np.sin(q) + 0.3*np.cos(2*q); rho0 = np.clip(rho0, 1e-3, None); rho0 /= rho0.sum()*dq
c_init = np.zeros((NH, Nq)); c_init[0] = rho0
s = KineticVFP(q, dq, k, NH, gamma, beta, 0.02, V)
_, sol = s.solve(W, c_init, Nt=100, save_every=25)
mass = sol[:, 0, :].sum(axis=1)*dq
print(f"mass drift over the run   : {np.max(np.abs(mass-mass[0])):.2e}")

# (ii) exact friction spectrum (uniform in q, W=0)
s0 = KineticVFP(q, dq, k, NH, gamma, beta, 0.05, np.zeros(Nq))
c0 = np.zeros((NH, Nq)); c0[0]=1/(2*np.pi); c0[2]=0.2/(2*np.pi)
tt, sl = s0.solve(np.zeros(Nq), c0, Nt=20, save_every=4)
num = sl[:, 2, :].mean(axis=1)*2*np.pi; exact = c0[2,0]*2*np.pi*np.exp(-gamma*2*tt)
print(f"friction mode n=2 rel err : {np.max(np.abs(num-exact)/np.abs(exact)):.2e}")

mass drift over the run   : 4.88e-15
friction mode n=2 rel err : 2.44e-15

def generate_truth(seed=0, W_true=None, Nq=48, NH_hi=20, gamma=1.0, beta=1.0, T=2.0, dt_hi=0.01):
    q, dq, k = make_grid(Nq); V = -np.cos(q)
    if W_true is None:
        W_true = -np.cos(q) + 0.5*np.cos(2*q) - 0.25*np.cos(3*q)
    rho0 = 1 + 0.6*np.sin(q) + 0.3*np.cos(2*q); rho0 = np.clip(rho0, 1e-3, None); rho0 /= rho0.sum()*dq
    c_init = np.zeros((NH_hi, Nq)); c_init[0] = rho0
    solver_hi = KineticVFP(q, dq, k, NH_hi, gamma, beta, dt_hi, V)
    times, sol = solver_hi.solve(W_true, c_init, Nt=int(T/dt_hi), save_every=2)
    return dict(q=q, dq=dq, k=k, T=T, gamma=gamma, beta=beta, V=V, rho0=rho0,
                W_true=W_true, times=times, sol_hi=sol, NH_hi=NH_hi)

truth = generate_truth(0)
q, dq, k = truth['q'], truth['dq'], truth['k']

# shared observation design: same (t_i, q_i) for both models
rng = np.random.default_rng(0)
N_obs = 3000
t_idx = rng.integers(1, len(truth['times']), N_obs)   # exclude t=0
q_idx = rng.integers(0, len(q), N_obs)

# --- f-data: also draw momenta p_i and evaluate the full state ---
p_obs = rng.uniform(-3.5, 3.5, N_obs)
B_hi  = hermite_eval(truth['NH_hi'], p_obs, truth['beta'])
f_clean = np.einsum('in,in->i', truth['sol_hi'][t_idx, :, q_idx], B_hi)

# --- rho-data: just read c_0 at the same (t_i, q_i); no momentum ---
rho_clean = truth['sol_hi'][t_idx, 0, q_idx]

# match SNR (noise / signal-std) across the two observables
SNR = 0.15
sig_f  = SNR*f_clean.std()
sig_rho = SNR*rho_clean.std()
Y_f   = f_clean   + sig_f  *rng.standard_normal(N_obs)
Y_rho = rho_clean + sig_rho*rng.standard_normal(N_obs)

print(f"f  :  signal std = {f_clean.std():.4f},  sigma = {sig_f:.4f}")
print(f"rho:  signal std = {rho_clean.std():.4f},  sigma = {sig_rho:.4f}")
print(f"matched noise/signal ratio = {SNR}")
print(f"mass at t=0 = {truth['sol_hi'][0,0].sum()*dq:.10f},  at t=T = {truth['sol_hi'][-1,0].sum()*dq:.10f}")

f  :  signal std = 0.0269,  sigma = 0.0040
rho:  signal std = 0.0734,  sigma = 0.0110
matched noise/signal ratio = 0.15
mass at t=0 = 1.0000000000,  at t=T = 1.0000000000

fig, ax = plt.subplots(figsize=(5, 3.5))
ax.plot(q, truth['W_true'], lw=2, color='C0')
ax.set_title(r"Ground-truth $W_0(q)$ (to be recovered)"); ax.set_xlabel("q")
plt.tight_layout(); plt.show()

fig, ax = plt.subplots(figsize=(5, 3.5))
im = ax.imshow(truth['sol_hi'][:, 0, :], origin='lower', aspect='auto',
               extent=[0, 2*np.pi, 0, truth['T']], cmap='magma')
ax.set_title(r"OBSERVED: density $\rho(t,q)=c_0$"); ax.set_xlabel("q"); ax.set_ylabel("t")
plt.colorbar(im, ax=ax)
plt.tight_layout(); plt.show()

fig, ax = plt.subplots(figsize=(5, 3.5))
im = ax.imshow(truth['sol_hi'][:, 1, :], origin='lower', aspect='auto',
               extent=[0, 2*np.pi, 0, truth['T']], cmap='coolwarm')
ax.set_title(r"NOT observed: mean momentum $c_1(t,q)$"); ax.set_xlabel("q"); ax.set_ylabel("t")
plt.colorbar(im, ax=ax)
plt.tight_layout(); plt.show()

# a couple of density snapshots with the actual noisy observations overlaid
fig, ax = plt.subplots(figsize=(7, 3.4))
for tt, c in [(truth['T']/2, 'C0'), (truth['T'], 'C3')]:
    it = int(np.argmin(np.abs(truth['times']-tt)))
    ax.plot(q, truth['sol_hi'][it, 0], color=c, lw=2, label=fr"$\rho(t={truth['times'][it]:.1f})$")
    sel = (t_idx == it)
    ax.plot(q[q_idx[sel]], Y_rho[sel], '.', color=c, ms=4, alpha=0.5)
ax.plot(q, truth['rho0'], 'k:', lw=1.3, label=r'$\rho_0$ (start)')
ax.set_title("Density snapshots (lines) with noisy $\\rho$-observations (dots)")
ax.set_xlabel("q"); ax.legend(fontsize=8); plt.tight_layout(); plt.show()

def build_basis_cos(q, M): return np.stack([np.cos(m*q) for m in range(1, M+1)], axis=1)
def make_scaling(M, tau=1.0, s=1.5): return np.array([tau/m**s for m in range(1, M+1)])
def W_from_xi(xi, Phi, scale): return Phi @ (scale*xi)

M, tau, s = 6, 1.0, 1.5
Phi = build_basis_cos(q, M); scale = make_scaling(M, tau, s)
coef_true = np.array([np.trapz(np.cos(m*q)*truth['W_true'], q)/np.pi for m in range(1, M+1)])

# inference-resolution solver (coarser than the truth: NH=14, dt=0.02) and observation basis
NH_inf, dt_inf = 14, 0.02
solver = KineticVFP(q, dq, k, NH_inf, truth['gamma'], truth['beta'], dt_inf, truth['V'])
c_init_inf = np.zeros((NH_inf, len(q))); c_init_inf[0] = truth['rho0']
Nt_inf = int(truth['T']/dt_inf)
B_obs = hermite_eval(NH_inf, p_obs, truth['beta'])   # only needed for the f-likelihood

def loglik_rho(xi):
    _, sol = solver.solve(W_from_xi(xi, Phi, scale), c_init_inf, Nt_inf, save_every=1)
    pred = sol[t_idx, 0, q_idx]                        # rho = c_0 only
    return -0.5*np.sum((Y_rho - pred)**2)/sig_rho**2

def loglik_f(xi):
    _, sol = solver.solve(W_from_xi(xi, Phi, scale), c_init_inf, Nt_inf, save_every=1)
    pred = np.einsum('in,in->i', sol[t_idx, :, q_idx], B_obs)   # full f = sum_n c_n e_n
    return -0.5*np.sum((Y_f - pred)**2)/sig_f**2

def pcn(loglik, beta_pcn, n_iter=2500, seed=1, log_every=500):
    rng = np.random.default_rng(seed)
    xi = np.zeros(M); ll = loglik(xi)
    samples = np.empty((n_iter, M)); ll_chain = np.empty(n_iter); nacc = 0; t0 = time.time()
    for it in range(n_iter):
        xp = np.sqrt(1-beta_pcn**2)*xi + beta_pcn*rng.standard_normal(M)
        lp = loglik(xp)
        if np.log(rng.random()) < lp - ll: xi, ll = xp, lp; nacc += 1
        samples[it], ll_chain[it] = xi, ll
        if (it+1) % log_every == 0:
            print(f"  iter {it+1:5d}/{n_iter}  acc={nacc/(it+1):5.1%}  ll={ll:9.1f}  t={time.time()-t0:5.0f}s")
    return samples, ll_chain, nacc/n_iter

print(f"one rho-likelihood eval costs ~", end=" ")
_t=time.time(); loglik_rho(np.zeros(M)); print(f"{time.time()-_t:.3f}s;  target plateau -N/2 = {-N_obs/2:.0f}")

one rho-likelihood eval costs ~ 0.042s;  target plateau -N/2 = -1500

3.7.1 Recovery from the density alone

samples_rho, ll_rho, acc_rho = pcn(loglik_rho, beta_pcn=0.09, n_iter=2500, seed=1)
print(f"\nrho-observed acceptance: {acc_rho:.1%}")

burn = 2500//5
post_rho = samples_rho[burn:]
W_rho = W_from_xi(post_rho.mean(0), Phi, scale)
_, sol_rho = solver.solve(W_rho, c_init_inf, Nt_inf, save_every=1)
errW_rho   = np.sqrt(np.trapz((W_rho - truth['W_true'])**2, q))
errRho_rho = np.sqrt(np.trapz((sol_rho[-1,0] - truth['sol_hi'][-1,0])**2, q))
print(f"||W_bar - W_true||_L2        = {errW_rho:.4f}")
print(f"||rho_bar(T) - rho_true(T)|| = {errRho_rho:.4f}   (<< sigma_rho = {sig_rho:.4f})")

  iter   500/2500  acc=19.4%  ll=  -1489.6  t=   11s
  iter  1000/2500  acc=15.2%  ll=  -1488.3  t=   21s
  iter  1500/2500  acc=15.5%  ll=  -1489.6  t=   32s
  iter  2000/2500  acc=16.0%  ll=  -1490.0  t=   42s
  iter  2500/2500  acc=16.2%  ll=  -1488.8  t=   53s

rho-observed acceptance: 16.2%
||W_bar - W_true||_L2        = 0.2477
||rho_bar(T) - rho_true(T)|| = 0.0004   (<< sigma_rho = 0.0110)

rp = np.random.default_rng(7)

fig, ax = plt.subplots(figsize=(6.5, 4.2))
for j in rp.choice(len(post_rho), 80, replace=False):
    ax.plot(q, W_from_xi(post_rho[j], Phi, scale), color='gray', alpha=0.12)
ax.plot(q, truth['W_true'], lw=2.4, label='True $W_0$')
ax.plot(q, W_rho, '--', lw=2.4, color='C3', label=r'Posterior mean $\bar W$ ($\rho$-observed)')
ax.set_title("Interaction potential from density data"); ax.set_xlabel("q"); ax.legend()
plt.tight_layout(); plt.show()

fig, ax = plt.subplots(figsize=(6.5, 4.2))
ax.plot(q, truth['sol_hi'][-1,0], 'k', lw=2, label=r'true $\rho(T)$')
ax.plot(q, sol_rho[-1,0], 'C3--', lw=1.6, label=r'reconstructed $\rho(T)$')
ax.set_title("Observable at $T$: truth vs fit"); ax.set_xlabel("q"); ax.legend(fontsize=9)
plt.tight_layout(); plt.show()

fig, ax = plt.subplots(figsize=(6.5, 4.2))
ax.plot(ll_rho, lw=0.7); ax.axvline(burn, color='C3', ls=':', label='burn-in')
ax.axhline(-N_obs/2, color='C2', ls='--', lw=1, label=r'$-N/2$')
ax.set_title("Log-likelihood trace ($\\rho$-observed)"); ax.set_xlabel("iteration"); ax.legend(fontsize=9)
plt.tight_layout(); plt.show()

fig, ax = plt.subplots(figsize=(6.5, 4.2))
cchain = post_rho*scale; cmean, cstd = cchain.mean(0), cchain.std(0)
pos = np.arange(1, M+1)
ax.errorbar(pos, cmean, yerr=cstd, fmt='o', capsize=3, label=r'posterior $\pm$ std')
ax.plot(pos, coef_true, 'rx', ms=11, mew=2.2, label='truth')
ax.axhline(0, color='k', lw=0.4); ax.set_xticks(pos)
ax.set_xticklabels([f"$a_{m}$" for m in range(1, M+1)])
ax.set_title("Cosine coefficients of $W$ ($\\rho$-observed)"); ax.set_xlabel("mode $m$"); ax.legend()
plt.tight_layout(); plt.show()
print("true:", np.round(coef_true,3)); print("post:", np.round(cmean,3), "+-", np.round(cstd,3))

true: [-0.969  0.531 -0.22   0.029  0.028  0.027]
post: [-0.998  0.521 -0.258 -0.104  0.068 -0.061] +- [0.014 0.04  0.067 0.088 0.023 0.046]

3.7.2 Head to head at matched noise

samples_f, ll_f, acc_f = pcn(loglik_f, beta_pcn=0.11, n_iter=2500, seed=1)
print(f"\nf-observed acceptance: {acc_f:.1%}")
post_f = samples_f[burn:]
W_f = W_from_xi(post_f.mean(0), Phi, scale)
errW_f = np.sqrt(np.trapz((W_f - truth['W_true'])**2, q))
cmean_f, cstd_f = (post_f*scale).mean(0), (post_f*scale).std(0)

print("\n            " + "  ".join(f"a{m}" for m in range(1,M+1)))
print("true      :", " ".join(f"{v:+.3f}" for v in coef_true))
print("rho posterior:", " ".join(f"{v:+.3f}" for v in cmean))
print("f   posterior:", " ".join(f"{v:+.3f}" for v in cmean_f))
print(f"\n||W_bar - W_true||_L2   rho: {errW_rho:.4f}    f: {errW_f:.4f}")

fig, ax = plt.subplots(figsize=(6.5, 4.2))
ax.plot(q, truth['W_true'], 'k', lw=2.6, label='True $W_0$')
ax.plot(q, W_rho, '--', lw=2.2, color='C3', label=r'$\bar W$ from $\rho$ only')
ax.plot(q, W_f, ':', lw=2.4, color='C0', label=r'$\bar W$ from full $f$')
ax.set_title("Reconstructed potential: $\\rho$ vs full-$f$ observation"); ax.set_xlabel("q"); ax.legend()
plt.tight_layout(); plt.show()

w = 0.36; pos = np.arange(1, M+1)
fig, ax = plt.subplots(figsize=(6.5, 4.2))
ax.bar(pos - w/2, np.abs(cmean - coef_true),   w, color='C3', label=r'$\rho$-observed')
ax.bar(pos + w/2, np.abs(cmean_f - coef_true), w, color='C0', label=r'$f$-observed')
ax.set_xticks(pos); ax.set_xticklabels([f"$a_{m}$" for m in range(1, M+1)])
ax.set_title("|posterior mean $-$ truth| per mode"); ax.set_xlabel("mode $m$"); ax.legend()
plt.tight_layout(); plt.show()

  iter   500/2500  acc=22.6%  ll=  -1535.5  t=   11s
  iter  1000/2500  acc=21.5%  ll=  -1534.4  t=   21s
  iter  1500/2500  acc=22.1%  ll=  -1533.8  t=   32s
  iter  2000/2500  acc=23.0%  ll=  -1535.2  t=   43s
  iter  2500/2500  acc=23.2%  ll=  -1534.2  t=   54s

f-observed acceptance: 23.2%

            a1  a2  a3  a4  a5  a6
true      : -0.969 +0.531 -0.220 +0.029 +0.028 +0.027
rho posterior: -0.998 +0.521 -0.258 -0.104 +0.068 -0.061
f   posterior: -0.986 +0.546 -0.335 -0.192 +0.042 -0.017

||W_bar - W_true||_L2   rho: 0.2477    f: 0.3845

3.7.3 Shared identifiability structure of the two observation models

def jacobian(predict, xi0, h=1e-3):
    base = predict(xi0); J = np.zeros((len(base), M))
    for m in range(M):
        e = np.zeros(M); e[m] = h
        J[:, m] = (predict(xi0 + e) - predict(xi0 - e))/(2*h)
    return J

def predict_rho(xi):
    _, sol = solver.solve(W_from_xi(xi, Phi, scale), c_init_inf, Nt_inf, save_every=1)
    return sol[t_idx, 0, q_idx]
def predict_f(xi):
    _, sol = solver.solve(W_from_xi(xi, Phi, scale), c_init_inf, Nt_inf, save_every=1)
    return np.einsum('in,in->i', sol[t_idx, :, q_idx], B_obs)

xi_true = coef_true/scale
S_rho = np.linalg.svd(jacobian(predict_rho, xi_true), compute_uv=False)
S_f   = np.linalg.svd(jacobian(predict_f,   xi_true), compute_uv=False)

print("singular values of the forward-map Jacobian (whitened):")
print("  rho:", np.array2string(S_rho, formatter={'float_kind':lambda v:f'{v:.2e}'}))
print("  f  :", np.array2string(S_f,   formatter={'float_kind':lambda v:f'{v:.2e}'}))
print("\nshape (ratio to the largest) -- this is the identifiability structure:")
print("  rho:", np.array2string(S_rho/S_rho[0], formatter={'float_kind':lambda v:f'{v:.1e}'}))
print("  f  :", np.array2string(S_f/S_f[0],   formatter={'float_kind':lambda v:f'{v:.1e}'}))

fig, ax = plt.subplots(figsize=(6.5, 4))
ax.semilogy(np.arange(1, M+1), S_rho, 'o-', color='C3', label=r'$\rho$-observed')
ax.semilogy(np.arange(1, M+1), S_f,   's--', color='C0', label=r'$f$-observed')
ax.set_xlabel("singular direction $\\approx$ mode $m$"); ax.set_ylabel(r"singular value $\sigma_j$ (log)")
ax.set_title("Absolute sensitivity (f sits below: values of $f$ are smaller)"); ax.legend()
plt.tight_layout(); plt.show()

fig, ax = plt.subplots(figsize=(6.5, 4))
ax.semilogy(np.arange(1, M+1), S_rho/S_rho[0], 'o-', color='C3', label=r'$\rho$-observed')
ax.semilogy(np.arange(1, M+1), S_f/S_f[0],   's--', color='C0', label=r'$f$-observed')
ax.set_xlabel("singular direction $\\approx$ mode $m$"); ax.set_ylabel(r"$\sigma_j/\sigma_1$ (log)")
ax.set_title("Same SHAPE = same identifiability structure"); ax.legend()
plt.tight_layout(); plt.show()

singular values of the forward-map Jacobian (whitened):
  rho: [7.90e-01 1.12e-01 8.30e-03 5.26e-04 4.84e-05 2.43e-06]
  f  : [1.73e-01 2.71e-02 1.96e-03 1.21e-04 1.16e-05 5.60e-07]

shape (ratio to the largest) -- this is the identifiability structure:
  rho: [1.0e+00 1.4e-01 1.1e-02 6.7e-04 6.1e-05 3.1e-06]
  f  : [1.0e+00 1.6e-01 1.1e-02 7.0e-04 6.7e-05 3.2e-06]

3.7.4 Robustness of the density model in the fully parametric limit

# sweep prior truncation M up to the fully-parametric limit, from rho data
def run_pcn_M_rho(Mtry, n_iter=2000, beta_pcn=0.06, seed=1, s_prior=1.5, tau_prior=1.0):
    """Self-contained pCN on the rho-likelihood at prior truncation Mtry."""
    Phi_M   = build_basis_cos(q, Mtry)
    scale_M = make_scaling(Mtry, tau_prior, s_prior)
    def loglik(xi):
        _, sol = solver.solve(Phi_M @ (scale_M*xi), c_init_inf, Nt_inf, save_every=1)
        return -0.5*np.sum((Y_rho - sol[t_idx, 0, q_idx])**2)/sig_rho**2   # rho = c_0 only
    rng = np.random.default_rng(seed)
    xi = np.zeros(Mtry); ll = loglik(xi)
    samples = np.empty((n_iter, Mtry)); n_acc = 0
    for it in range(n_iter):
        xi_p = np.sqrt(1-beta_pcn**2)*xi + beta_pcn*rng.standard_normal(Mtry)
        ll_p = loglik(xi_p)
        if np.log(rng.random()) < ll_p - ll:
            xi, ll = xi_p, ll_p; n_acc += 1
        samples[it] = xi
    return dict(samples=samples, scale=scale_M, Phi=Phi_M, acc=n_acc/n_iter)

M_full = len(q)//2 - 1
M_grid = [3, 6, 12, M_full]
print(f"fully-parametric limit for Nq={len(q)}:  M_full = {M_full}\n")

sweep = {}
t0 = time.time()
for Mtry in M_grid:
    res  = run_pcn_M_rho(Mtry, n_iter=2000, beta_pcn=0.06)
    post = res['samples'][2000//5:]
    W_hat = res['Phi'] @ (res['scale'] * post.mean(0))
    res['post'], res['W_hat'] = post, W_hat
    res['errW'] = np.sqrt(np.trapz((W_hat - truth['W_true'])**2, q))
    sweep[Mtry] = res
    print(f"  M={Mtry:2d}   acc={res['acc']:5.1%}   ||W_bar - W_true||_L2 = {res['errW']:.4f}")
print(f"\nsweep wall-clock: {time.time()-t0:.0f}s")

fig, ax = plt.subplots(1, 2, figsize=(11, 3.9))
ax[0].plot(M_grid, [sweep[M]['errW'] for M in M_grid], 'o-', color='C3')
ax[0].set_xlabel("prior truncation $M$"); ax[0].set_ylabel(r"$\|\bar W - W_0\|_{L^2}$")
ax[0].set_title(r"Density data: error bounded as $M\to$ fully-parametric"); ax[0].set_ylim(bottom=0)
for M in M_grid:
    ax[1].plot(q, sweep[M]['W_hat'], lw=1.8, label=f"$M={M}$")
ax[1].plot(q, truth['W_true'], 'k', lw=2.6, label='truth')
ax[1].set_xlabel("q"); ax[1].set_title(r"$\bar W$ from $\rho$ for growing $M$"); ax[1].legend(fontsize=8, ncol=2)
plt.tight_layout(); plt.show()

fully-parametric limit for Nq=48:  M_full = 23

  M= 3   acc=24.9%   ||W_bar - W_true||_L2 = 0.2246
  M= 6   acc=24.1%   ||W_bar - W_true||_L2 = 0.3021
  M=12   acc=23.8%   ||W_bar - W_true||_L2 = 0.3822
  M=23   acc=26.8%   ||W_bar - W_true||_L2 = 0.3070

sweep wall-clock: 176s

3.7.5 Higher modes return to the prior — and to the same place for $\rho$ and $f$

# Laplace shrinkage per mode, for rho and f, at the fully-parametric prior
Phi_full   = build_basis_cos(q, M_full)
scale_full = make_scaling(M_full, 1.0, 1.5)
coef_true_full = np.array([np.trapz(np.cos(m*q)*truth['W_true'], q)/np.pi for m in range(1, M_full+1)])
xi_true_full = coef_true_full / scale_full

def predict_rho_full(xi):
    _, sol = solver.solve(Phi_full @ (scale_full*xi), c_init_inf, Nt_inf, save_every=1)
    return sol[t_idx, 0, q_idx]
def predict_f_full(xi):
    _, sol = solver.solve(Phi_full @ (scale_full*xi), c_init_inf, Nt_inf, save_every=1)
    return np.einsum('in,in->i', sol[t_idx, :, q_idx], B_obs)

def laplace_shrink(predict, sig):
    h = 1e-3; F0 = predict(xi_true_full); J = np.zeros((len(F0), M_full))
    for m in range(M_full):
        e = np.zeros(M_full); e[m] = h
        J[:, m] = (predict(xi_true_full + e) - predict(xi_true_full - e)) / (2*h)
    Sig = np.linalg.inv(np.eye(M_full) + J.T @ J / sig**2)
    return 1 - np.diag(Sig), np.sqrt(np.clip(np.diag(Sig), 0, None))

shrink_rho, sd_rho_w = laplace_shrink(predict_rho_full, sig_rho)
shrink_f,   sd_f_w   = laplace_shrink(predict_f_full,   sig_f)
prior_sd    = scale_full
post_sd_rho = scale_full * sd_rho_w

gauss = lambda x, sd: np.exp(-0.5*(x/sd)**2)/(sd*np.sqrt(2*np.pi))

# (1) density-only posterior vs prior, mode by mode (first 12 modes)
fig, axes = plt.subplots(3, 4, figsize=(12, 7.4))
for m, ax in enumerate(axes.ravel()):
    xg = np.linspace(-3*prior_sd[m], 3*prior_sd[m], 400)
    ax.fill_between(xg, gauss(xg, prior_sd[m]), color='C7', alpha=0.30, label='prior')
    ax.plot(xg, gauss(xg, post_sd_rho[m]), 'C3', lw=2, label=r'posterior ($\rho$)')
    ax.axvline(coef_true_full[m], color='k', ls='--', lw=1.4, label='truth')
    ax.set_title(f"$a_{{{m+1}}}$   (shrink {shrink_rho[m]:.2f})", fontsize=10); ax.set_yticks([])
    if m == 0: ax.legend(fontsize=7, loc='upper left')
fig.suptitle("Density-only posterior vs prior, mode by mode: high modes revert to the prior", y=1.01)
plt.tight_layout(); plt.show()

# (2) shrinkage ladders: rho vs f fall off together
mm = np.arange(1, M_full+1)
fig, ax = plt.subplots(figsize=(7.6, 4))
ax.plot(mm, shrink_rho, 'o-',  color='C3', label=r'$\rho$-observed')
ax.plot(mm, shrink_f,   's--', color='C0', label='$f$-observed')
ax.axhline(0, color='k', lw=0.5)
ax.set_xlabel("Fourier mode $m$"); ax.set_ylabel(r"posterior shrinkage $s_m$")
ax.set_title(r"$\rho$ and $f$ send the same modes back to the prior")
ax.legend(); plt.tight_layout(); plt.show()

print("mode      :", " ".join(f"{m:5d}" for m in range(1, 8)))
print("shrink rho:", " ".join(f"{v:5.2f}" for v in shrink_rho[:7]))
print("shrink f  :", " ".join(f"{v:5.2f}" for v in shrink_f[:7]))

mode      :     1     2     3     4     5     6     7
shrink rho:  1.00  0.99  0.36  0.00  0.00  0.00  0.00
shrink f  :  1.00  0.98  0.19  0.00  0.00  0.00  0.00