make_figures.py

"""
make_figures.py
===============
Generates all 8 publication-quality figures for:

    Bayesian Inversion for Diffusion Coefficient Recovery
    Using Kernel Surrogates with Posterior Error Quantification

    Akhilesh Yadav — Independent Research, 2026

HOW TO RUN
----------
    # Run this AFTER src/run_experiment.py has finished:
    python make_figures.py

PREREQUISITES
-------------
    results/summary.json       (generated by run_experiment.py)
    results/samps_fem.npy
    results/samps_surr_5.npy
    results/samps_surr_20.npy
    results/samps_surr_80.npy
    results/trace_fem.npy

FIGURES PRODUCED (saved to figures/)
--------------------------------------
    fig1_data.png                   True solution, observations, a_true(x)
    fig2_posteriors_fem_vs_best.png FEM vs n=80 marginal posteriors
    fig3_all_surrogates.png         FEM vs surrogates n=5, 20, 80
    fig4_reconstruction.png         a(x) posterior mean + 95% credible band
    fig5_wasserstein.png            W_2 convergence + per-dim bar chart
    fig6_diagnostics.png            MCMC trace, ACF, running means
    fig7_predictive_check.png       Posterior predictive check
    fig8_results_table.png          Numerical summary table
"""

import sys
import os
import json

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec

# -- ensure src/ is on path ------------------------------------------
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), 'src'))
from fem_solver   import fem_solution_at_points
from mcmc_sampler import make_a_func, wasserstein2_1d

# ====================================================================
# PATHS
# ====================================================================

ROOT    = os.path.dirname(os.path.abspath(__file__))
RES_DIR = os.path.join(ROOT, 'results')
FIG_DIR = os.path.join(ROOT, 'figures')
os.makedirs(FIG_DIR, exist_ok=True)

# ====================================================================
# CHECK PREREQUISITES
# ====================================================================

required = [
    'summary.json',
    'samps_fem.npy',
    'samps_surr_5.npy',
    'samps_surr_20.npy',
    'samps_surr_80.npy',
    'trace_fem.npy',
]
missing = [f for f in required
           if not os.path.exists(os.path.join(RES_DIR, f))]
if missing:
    print("ERROR: The following result files are missing:")
    for f in missing:
        print(f"  results/{f}")
    print("\nPlease run  python src/run_experiment.py  first.")
    sys.exit(1)

# ====================================================================
# LOAD DATA
# ====================================================================

with open(os.path.join(RES_DIR, 'summary.json')) as fh:
    D = json.load(fh)

THETA_TRUE = np.array(D['theta_true'])
K          = len(THETA_TRUE)
SIGMA_N    = D['sigma_n']
x_obs      = np.array(D['x_obs'])
y_obs      = np.array(D['y_obs'])
u_clean    = np.array(D['u_clean'])
SR         = D['surrogate_results']        # list of dicts

samps_fem  = np.load(os.path.join(RES_DIR, 'samps_fem.npy'))
samps_5    = np.load(os.path.join(RES_DIR, 'samps_surr_5.npy'))
samps_20   = np.load(os.path.join(RES_DIR, 'samps_surr_20.npy'))
samps_80   = np.load(os.path.join(RES_DIR, 'samps_surr_80.npy'))
trace_fem  = np.load(os.path.join(RES_DIR, 'trace_fem.npy'))

f_func = lambda x: np.pi ** 2 * np.sin(np.pi * float(x))

# ====================================================================
# PLOT STYLE
# ====================================================================

plt.rcParams.update({
    'font.family'       : 'serif',
    'font.size'         : 11,
    'axes.labelsize'    : 12,
    'axes.titlesize'    : 12,
    'legend.fontsize'   : 9,
    'figure.dpi'        : 140,
    'lines.linewidth'   : 2.0,
    'axes.spines.top'   : False,
    'axes.spines.right' : False,
    'axes.grid'         : True,
    'grid.alpha'        : 0.20,
    'grid.linestyle'    : '--',
})

C_FEM   = '#1f77b4'   # blue   — FEM reference
C_S5    = '#d62728'   # red    — surrogate n=5
C_S20   = '#ff7f0e'   # orange — surrogate n=20
C_S80   = '#2ca02c'   # green  — surrogate n=80
C_TRUE  = 'black'     # ground truth
C_OBS   = '#333333'   # observations

CI_FEM  = '#aec7e8'   # light blue fill
CI_S80  = '#b5e6b5'   # light green fill

ALL_SAMPS  = [samps_fem, samps_5,  samps_20,  samps_80 ]
ALL_COLS   = [C_FEM,     C_S5,     C_S20,     C_S80    ]
ALL_LABS   = ['FEM (ref)', 'Surr n=5', 'Surr n=20', 'Surr n=80']


# ====================================================================
# FIGURE HELPERS
# ====================================================================

def savefig(fig, name):
    path = os.path.join(FIG_DIR, name)
    fig.savefig(path, bbox_inches='tight')
    plt.close(fig)
    print(f"  Saved: figures/{name}")


def posterior_a_bands(samples, x_grid, n_draw=400):
    """Sample diffusion coefficient reconstructions from posterior."""
    rng = np.random.default_rng(1)
    idx = rng.choice(len(samples), size=min(n_draw, len(samples)), replace=False)
    A   = np.array([[make_a_func(samples[j])(xi) for xi in x_grid]
                    for j in idx])
    return A


# ====================================================================
# FIGURE 1 — Data and true coefficient
# ====================================================================

def fig1_data():
    x_fine  = np.linspace(0, 1, 400)
    u_fine  = fem_solution_at_points(make_a_func(THETA_TRUE), f_func, x_fine, N=400)
    a_fine  = np.array([make_a_func(THETA_TRUE)(xi) for xi in x_fine])

    fig, axes = plt.subplots(1, 2, figsize=(12, 4.8))

    ax = axes[0]
    ax.plot(x_fine, u_fine, color=C_FEM, lw=2.5,
            label=r'$u_{\mathrm{true}}(x)$  (FEM, $N=400$)')
    ax.scatter(x_obs, y_obs, s=50, color=C_OBS, zorder=5,
               label=r'Noisy observations ($\sigma_n = $' + str(SIGMA_N) + ')')
    ax.scatter(x_obs, u_clean, s=20, color='tomato',
               marker='x', lw=2, zorder=6,
               label=r'Noiseless $u(x_i)$')
    ax.set_xlabel(r'$x$');  ax.set_ylabel(r'$u(x)$')
    ax.set_title('True PDE Solution and Synthetic Observations')
    ax.legend()

    ax = axes[1]
    ax.step(x_fine, a_fine, color=C_TRUE, lw=2.5, where='mid')
    colors_int = ['steelblue', 'tomato', 'seagreen', 'goldenrod']
    for k in range(K):
        ax.axvspan(k / K, (k + 1) / K, alpha=0.10, color=colors_int[k])
        ax.text((k + 0.5) / K, THETA_TRUE[k] + 0.12,
                r'$\theta_{' + str(k + 1) + r'} = $' + str(THETA_TRUE[k]),
                ha='center', fontsize=10, color='navy')
    ax.set_xlabel(r'$x$');  ax.set_ylabel(r'$a(x)$');  ax.set_ylim(0, 2.7)
    ax.set_title(r'True Piecewise-Constant Diffusion Coefficient $a_{\mathrm{true}}(x)$')

    plt.tight_layout()
    savefig(fig, 'fig1_data.png')


# ====================================================================
# FIGURE 2 — FEM vs best surrogate (n=80) marginal posteriors
# ====================================================================

def fig2_posteriors_fem_vs_best():
    fig, axes = plt.subplots(1, K, figsize=(4.5 * K, 4.8))
    for k in range(K):
        ax = axes[k]
        ax.hist(samps_fem[:, k], bins=45, density=True, alpha=0.65,
                color=C_FEM, label='FEM (reference)',
                edgecolor='white', lw=0.3)
        ax.hist(samps_80[:, k],  bins=45, density=True, alpha=0.65,
                color=C_S80, label='Surrogate $n = 80$',
                edgecolor='white', lw=0.3)
        ax.axvline(THETA_TRUE[k], color='black', ls='--', lw=2.2,
                   label=r'True $\theta_{' + str(k + 1) + r'} = $'
                         + str(THETA_TRUE[k]))
        ax.axvline(np.mean(samps_fem[:, k]), color=C_FEM, ls=':', lw=1.8)
        ax.axvline(np.mean(samps_80[:, k]),  color=C_S80, ls=':', lw=1.8)
        ax.set_xlabel(r'$\theta_{' + str(k + 1) + r'}$', fontsize=13)
        ax.set_ylabel('Density' if k == 0 else '')
        ax.set_title(r'Marginal: $\theta_{' + str(k + 1) + r'}$')
        if k == 0:
            ax.legend(fontsize=8)

    plt.suptitle('Posterior Marginals: FEM (reference) vs Best Surrogate ($n = 80$)',
                 y=1.02, fontsize=13)
    plt.tight_layout()
    savefig(fig, 'fig2_posteriors_fem_vs_best.png')


# ====================================================================
# FIGURE 3 — All four posteriors side by side
# ====================================================================

def fig3_all_surrogates():
    fig, axes = plt.subplots(1, K, figsize=(4.5 * K, 4.8))
    for k in range(K):
        ax = axes[k]
        for samp, cc, lb in zip(ALL_SAMPS, ALL_COLS, ALL_LABS):
            ax.hist(samp[:, k], bins=40, density=True, alpha=0.55,
                    color=cc, label=lb, edgecolor='white', lw=0.3)
        ax.axvline(THETA_TRUE[k], color='black', ls='--', lw=2.2,
                   label=r'True $\theta_{' + str(k + 1) + r'}$')
        ax.set_xlabel(r'$\theta_{' + str(k + 1) + r'}$', fontsize=12)
        ax.set_ylabel('Density' if k == 0 else '')
        ax.set_title(r'Marginal: $\theta_{' + str(k + 1) + r'}$')
        if k == 0:
            ax.legend(fontsize=8)

    plt.suptitle('Posterior Marginals: FEM vs Surrogates ($n = 5, 20, 80$)',
                 y=1.02, fontsize=13)
    plt.tight_layout()
    savefig(fig, 'fig3_all_surrogates.png')


# ====================================================================
# FIGURE 4 — Diffusion coefficient reconstruction
# ====================================================================

def fig4_reconstruction():
    x_grid     = np.linspace(0, 1, 400)
    a_true_arr = np.array([make_a_func(THETA_TRUE)(xi) for xi in x_grid])

    fig, axes = plt.subplots(1, 2, figsize=(13, 5.2))
    for ax, samp, label, cc, cs in [
        (axes[0], samps_fem, 'FEM posterior',    C_FEM, CI_FEM),
        (axes[1], samps_80,  'Surrogate $n=80$', C_S80, CI_S80),
    ]:
        A      = posterior_a_bands(samp, x_grid)
        mean   = np.mean(A, axis=0)
        ci_lo  = np.percentile(A,  2.5, axis=0)
        ci_hi  = np.percentile(A, 97.5, axis=0)

        ax.step(x_grid, a_true_arr, color='black', lw=2.5, where='mid',
                label=r'$a_{\mathrm{true}}(x)$')
        ax.plot(x_grid, mean, color=cc, lw=2.5, label=label + ' mean')
        ax.fill_between(x_grid, ci_lo, ci_hi,
                        alpha=0.35, color=cs, label='95% credible band')
        ax.set_xlabel(r'$x$');  ax.set_ylabel(r'$a(x)$')
        ax.set_title('Diffusion Coefficient Reconstruction\n(' + label + ')')
        ax.legend(fontsize=9);  ax.set_ylim(0, 2.8)

    plt.tight_layout()
    savefig(fig, 'fig4_reconstruction.png')


# ====================================================================
# FIGURE 5 — W_2 convergence
# ====================================================================

def fig5_wasserstein():
    n_vals    = [r['n_train']  for r in SR]
    surr_errs = [r['surr_err'] for r in SR]
    w2_vals   = [r['w2']       for r in SR]
    w2_per_k  = np.array([r['w2_dims'] for r in SR])   # shape (3, K)

    fig, axes = plt.subplots(1, 2, figsize=(12, 5.2))

    # Left: W_2 vs surrogate L^2 error (log-log)
    ax = axes[0]
    ax.loglog(surr_errs, w2_vals, 'o-', color='purple',
              markersize=10, lw=2.5, zorder=5)
    for ne, se, wv in zip(n_vals, surr_errs, w2_vals):
        ax.annotate(f'$n={ne}$', xy=(se, wv),
                    xytext=(se * 1.15, wv * 1.10),
                    fontsize=9, color='purple')
    # Reference slope-1 line
    xs = np.array([min(surr_errs) * 0.7, max(surr_errs) * 1.4])
    ax.loglog(xs, (w2_vals[0] / surr_errs[0]) * xs,
              'k--', lw=1.5, label='Slope 1')
    ax.legend()
    ax.set_xlabel(r'Surrogate $L^2$ relative error')
    ax.set_ylabel(r'$\bar{W}_2(\pi^{\mathrm{FEM}}, \pi^{\mathrm{surr}})$')
    ax.set_title(r'Posterior Error: $W_2$ vs Surrogate Accuracy')

    # Right: per-dimension bar chart
    ax   = axes[1]
    xbar = np.arange(1, K + 1)
    w    = 0.22
    for i, (ne, offs, bc) in enumerate(
            zip(n_vals, [-1, 0, 1], [C_S5, C_S20, C_S80])):
        bars = ax.bar(xbar + offs * w, w2_per_k[i], w,
                      color=bc, alpha=0.82, label=f'$n={ne}$',
                      edgecolor='white')
        for bar, v in zip(bars, w2_per_k[i]):
            ax.text(bar.get_x() + bar.get_width() / 2,
                    bar.get_height() + 0.002,
                    f'{v:.3f}', ha='center', fontsize=8, rotation=40)
    ax.set_xlabel('Parameter index $k$')
    ax.set_ylabel(r'$W_2(\pi^{\mathrm{FEM}}_k,\, \pi^{\mathrm{surr}}_k)$')
    ax.set_title(r'$W_2$ Distance per Marginal')
    ax.set_xticks(xbar)
    ax.set_xticklabels([r'$\theta_{' + str(k) + r'}$'
                        for k in range(1, K + 1)])
    ax.legend()

    plt.tight_layout()
    savefig(fig, 'fig5_wasserstein.png')


# ====================================================================
# FIGURE 6 — MCMC diagnostics
# ====================================================================

def fig6_diagnostics():
    fig = plt.figure(figsize=(14, 9))
    gs  = GridSpec(2, 3, figure=fig, hspace=0.48, wspace=0.40)

    n_total = len(trace_fem)

    # Log-posterior trace
    ax1 = fig.add_subplot(gs[0, :2])
    ax1.plot(trace_fem, color=C_FEM, lw=0.7, alpha=0.85,
             label='FEM log-posterior')
    ax1.axvline(600, color='grey', ls=':', lw=1.5, label='Burn-in (600)')
    ax1.set_xlabel('MCMC step');  ax1.set_ylabel('Log posterior')
    ax1.set_title('Log-Posterior Trace (FEM reference chain)')
    ax1.legend()

    # Autocorrelation of theta_1
    ax2 = fig.add_subplot(gs[0, 2])

    def autocorr(x, maxlag=80):
        xn = x - x.mean()
        c  = np.correlate(xn, xn, 'full')[len(xn) - 1:]
        return c[:maxlag] / c[0]

    lags = np.arange(80)
    for samp, cc, lb in zip([samps_fem, samps_80], [C_FEM, C_S80],
                             ['FEM', 'Surr $n=80$']):
        ax2.plot(lags, autocorr(samp[:, 0]), color=cc, lw=1.8, label=lb)
    ax2.axhline(0,      color='black', lw=0.8)
    ax2.axhline(1/np.e, color='grey',  lw=0.8, ls='--', label=r'$e^{-1}$')
    ax2.set_xlabel('Lag');  ax2.set_ylabel('ACF')
    ax2.set_title(r'Autocorrelation: $\theta_1$')
    ax2.legend()

    # Running posterior means
    ax3 = fig.add_subplot(gs[1, :])
    n_post = len(samps_fem)
    idx    = np.arange(1, n_post + 1)
    ls_list = ['-', '--', ':', '-.']
    for k in range(K):
        lbl_fem = 'FEM' if k == 0 else None
        lbl_s80 = 'Surr $n=80$' if k == 0 else None
        ax3.plot(idx,
                 np.cumsum(samps_fem[:, k]) / idx,
                 color=C_FEM, lw=1.8, ls=ls_list[k], alpha=0.9,
                 label=lbl_fem)
        ax3.plot(idx,
                 np.cumsum(samps_80[:, k]) / idx,
                 color=C_S80, lw=1.8, ls=ls_list[k], alpha=0.9,
                 label=lbl_s80)
        ax3.axhline(THETA_TRUE[k], color='black', lw=0.8, ls=':')

    ax3.set_xlabel('Post-burn-in sample index')
    ax3.set_ylabel('Running posterior mean')
    ax3.set_title(r'Running Means  (dotted lines = true $\theta_k$)')
    ax3.legend(fontsize=9)

    savefig(fig, 'fig6_diagnostics.png')


# ====================================================================
# FIGURE 7 — Posterior predictive check
# ====================================================================

def fig7_predictive_check():
    x_pp      = np.linspace(0, 1, 250)
    u_true_pp = fem_solution_at_points(make_a_func(THETA_TRUE), f_func,
                                        x_pp, N=300)
    rng = np.random.default_rng(0)

    fig, axes = plt.subplots(1, 2, figsize=(13, 5.2))
    for ax, samp, label, cc, cs in [
        (axes[0], samps_fem, 'FEM posterior',    C_FEM, CI_FEM),
        (axes[1], samps_80,  'Surrogate $n=80$', C_S80, CI_S80),
    ]:
        n_draw = 150
        idx    = rng.choice(len(samp), size=n_draw, replace=False)
        U_pp   = np.array([
            fem_solution_at_points(make_a_func(samp[j]), f_func,
                                   x_pp, N=100)
            for j in idx
        ])
        mean   = np.mean(U_pp, axis=0)
        ci_lo  = np.percentile(U_pp,  2.5, axis=0)
        ci_hi  = np.percentile(U_pp, 97.5, axis=0)

        ax.plot(x_pp, u_true_pp, 'k--', lw=2.2,
                label=r'$u_{\mathrm{true}}$')
        ax.plot(x_pp, mean, color=cc, lw=2.5,
                label='Posterior predictive mean')
        ax.fill_between(x_pp, ci_lo, ci_hi,
                        alpha=0.30, color=cs, label='95% predictive band')
        ax.scatter(x_obs, y_obs, s=30, color=C_OBS, zorder=6,
                   label='Observations')
        ax.set_xlabel(r'$x$');  ax.set_ylabel(r'$u(x)$')
        ax.set_title('Posterior Predictive Check\n(' + label + ')')
        ax.legend(fontsize=8)

    plt.tight_layout()
    savefig(fig, 'fig7_predictive_check.png')


# ====================================================================
# FIGURE 8 — Numerical results table
# ====================================================================

def fig8_results_table():
    summ_fem = D['fem']
    sr_best  = SR[-1]           # n=80

    col_headers = [
        r'$k$', r'True $\theta_k$',
        'FEM mean', 'FEM std', 'FEM 95% CI',
        'Surr $n=80$ mean', 'Surr $n=80$ std',
        r'$W_2$ ($n=80$)', 'In 95% CI',
    ]
    rows = []
    for k in range(K):
        in_fem  = (summ_fem['ci_lower'][k]
                   <= THETA_TRUE[k]
                   <= summ_fem['ci_upper'][k])
        in_surr = (sr_best['ci_lower'][k]
                   <= THETA_TRUE[k]
                   <= sr_best['ci_upper'][k])
        rows.append([
            str(k + 1),
            str(THETA_TRUE[k]),
            f"{summ_fem['mean'][k]:.3f}",
            f"{summ_fem['std'][k]:.3f}",
            f"[{summ_fem['ci_lower'][k]:.2f}, {summ_fem['ci_upper'][k]:.2f}]",
            f"{sr_best['mean'][k]:.3f}",
            f"{sr_best['std'][k]:.3f}",
            f"{sr_best['w2_dims'][k]:.4f}",
            ('FEM:Y' if in_fem else 'FEM:N') + ' | '
            + ('Surr:Y' if in_surr else 'Surr:N'),
        ])

    fig, ax = plt.subplots(figsize=(14, 4.0))
    ax.axis('off')
    tbl = ax.table(cellText=rows, colLabels=col_headers,
                   cellLoc='center', loc='center')
    tbl.auto_set_font_size(False)
    tbl.set_fontsize(9.5)
    tbl.scale(1.1, 1.8)

    # Style header row
    for j in range(len(col_headers)):
        tbl[(0, j)].set_facecolor('#2c3e50')
        tbl[(0, j)].set_text_props(color='white', fontweight='bold')
    # Alternate row colours
    for i in range(1, K + 1):
        for j in range(len(col_headers)):
            tbl[(i, j)].set_facecolor('#ecf0f1' if i % 2 == 0 else 'white')

    ax.set_title(
        'Numerical Results Summary: '
        'Bayesian Inversion for Diffusion Coefficient Recovery',
        pad=20, fontsize=12, fontweight='bold',
    )
    plt.tight_layout()
    savefig(fig, 'fig8_results_table.png')


# ====================================================================
# MAIN
# ====================================================================

if __name__ == '__main__':
    print()
    print("=" * 60)
    print("  Generating figures ...")
    print("=" * 60)

    fig1_data()
    fig2_posteriors_fem_vs_best()
    fig3_all_surrogates()
    fig4_reconstruction()
    fig5_wasserstein()
    fig6_diagnostics()
    fig7_predictive_check()
    fig8_results_table()

    print()
    print("=" * 60)
    print(f"  All 8 figures saved to:  figures/")
    print("=" * 60)