ergodic-sampling/ergodic_sampling_test.py

708 lines
38 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import numpy as np
import pywt
from scipy.interpolate import griddata
from scipy.stats import pearsonr, spearmanr
import matplotlib.pyplot as plt
rng = np.random.default_rng(0)
NX = NY = 64; N = NX*NY
BETA = 2.5 # potential-field truth: power spectrum ~ k^-BETA (before attenuation)
Z_SRC = 1.5 # source depth (cells): upward-continuation attenuation exp(-2*pi*k*Z_SRC)
N_MODES = 12 # Fourier-sparse truth: number of modes (low-frequency biased)
STRIDE = 3 # sparse (regular-grid) stride -> delta ~ 0.118
WAVELET, LEVELS, MODE = 'sym8', 2, 'periodization' # 2 = pywt.dwt_max_level(64, 'sym8')
SYM_SPINS = 8 # cycle-spin shifts for the symlet reconstruction (our choice; the
# paper never specifies cycle spinning or any solver parameters)
NBR, NBA, RMAX, BLK = 45, 36, 90.0, 4 # BLK=4 ~ paper's patch size A/N_phi (2.9 px, non-integer)
SA_ITERS = 4000
POCS_ITERS, POCS_LMIN = 800, 1e-4
LOO_M = 60 # leave-one-out calibration points
ENS_K = 64 # ensemble validation realizations (our construct, not the paper's)
# ---------------- ground truths (two tiers) ----------------
def powerlaw_field(beta, z_src=Z_SRC):
# random-phase power law with source-depth attenuation: without the exp(-2*pi*k*z) factor the
# field is rough at pixel scale, which real (upward-continued) potential-field data never is
kx = np.fft.fftfreq(NX)[:, None]; ky = np.fft.fftfreq(NY)[None, :]
k = np.hypot(kx, ky); k[0, 0] = 1.0
amp = k**(-beta/2)*np.exp(-2*np.pi*k*z_src); amp[0, 0] = 0.0
f = np.real(np.fft.ifft2(amp*np.exp(2j*np.pi*rng.random((NX, NY)))))
return (f - f.mean())/f.std()
def fourier_sparse_field(n_modes):
# visually structured sparse signal: amplitude ~ |k|^-1/2 so low frequencies dominate the look,
# with half the modes beyond the sparse grid's Nyquist 1/(2*STRIDE) so that grid aliases
k_alias = 1/(2*STRIDE)
C = np.zeros((NX, NY), complex); picked = set()
m = 0
while m < n_modes:
lo, hi = ((k_alias, 0.30) if m < n_modes//2 else (0.03, k_alias))
kmag = np.exp(rng.uniform(np.log(lo), np.log(hi)))
ang = rng.uniform(0, 2*np.pi)
i = int(round(kmag*np.cos(ang)*NX)) % NX
j = int(round(kmag*np.sin(ang)*NY)) % NY
if (i, j) in picked or (i, j) == (0, 0):
continue
picked.add((i, j))
C[i, j] = kmag**-0.5*np.exp(2j*np.pi*rng.random())
m += 1
f = np.fft.ifft2(C).real
return (f - f.mean())/f.std()
truth_pl = powerlaw_field(BETA) # tier 2: hard, not l1-sparse in any basis
truth_sp = fourier_sparse_field(N_MODES) # tier 1: exactly sparse, the paper's Fig. A.1 setting
# ---------------- sampling patterns ----------------
mask_reg = np.zeros((NX, NY), bool); mask_reg[::STRIDE, ::STRIDE] = True
NS = int(mask_reg.sum()); DELTA = NS/N
def random_mask():
m = np.zeros(N, bool); m[rng.choice(N, NS, replace=False)] = True
return m.reshape(NX, NY)
# ---------------- ISA properties (Eq. 4 ingredients) ----------------
def H(p):
p = np.abs(p); nz = p > 0
return -np.sum(p[nz]*np.log(p[nz]))
def dense_pair_hists(nx, ny):
# all-pairs interval/angle histograms of the dense grid via closed-form pair counts
DX, DY = np.meshgrid(np.arange(-(nx-1), nx), np.arange(0, ny), indexing='ij')
keep = (DY > 0) | ((DY == 0) & (DX > 0))
DX, DY = DX[keep], DY[keep]
w = (nx - np.abs(DX))*(ny - np.abs(DY))
hr, _ = np.histogram(np.hypot(DX, DY), NBR, (0, RMAX), weights=w)
# angles folded to [0, pi): the all-pairs rose diagram is 180-degree symmetric
ha, _ = np.histogram(np.arctan2(DY, DX) % np.pi, NBA, (0, np.pi), weights=w)
return hr/hr.max(), ha/ha.max()
def sparse_pair_hists(mask):
c = np.argwhere(mask).astype(float)
d = c[:, None, :] - c[None, :, :]
iu = np.triu_indices(len(c), 1)
dx, dy = d[..., 0][iu], d[..., 1][iu]
hr, _ = np.histogram(np.hypot(dx, dy), NBR, (0, RMAX))
ha, _ = np.histogram(np.arctan2(dy, dx) % np.pi, NBA, (0, np.pi))
return hr/hr.max(), ha/ha.max()
def block_density(mask):
# per-block occupancy fraction gamma_phi; compared against delta*gamma_theta = DELTA (Eq. 2)
return mask.reshape(NX//BLK, BLK, NY//BLK, BLK).sum((1, 3))/BLK**2
def srf(mask):
P = np.abs(np.fft.fft2(mask.astype(float))); P /= P[0, 0]
xi = P.ravel()[1:]
return xi.max(), xi
HR0, HA0 = dense_pair_hists(NX, NY)
def objective(mask):
# Eq. 4 with the paper's 2D weights w_alpha = w_beta = w_gamma = mu
hr, ha = sparse_pair_hists(mask)
mu, xi = srf(mask)
return mu*(H(hr - HR0) + H(ha - HA0) + H(block_density(mask) - DELTA) + H(xi))
def ergodic_mask(iters=SA_ITERS, T0=0.05):
m = random_mask(); f = objective(m); best = (f, m.copy())
for it in range(iters):
T = T0*(1 - it/iters)
on = np.flatnonzero(m.ravel()); off = np.flatnonzero(~m.ravel())
i, j = rng.choice(on), rng.choice(off)
m.flat[i] = False; m.flat[j] = True
f2 = objective(m)
if f2 < f or rng.random() < np.exp(-(f2-f)/max(T, 1e-9)):
f = f2
if f < best[0]: best = (f, m.copy())
else:
m.flat[i] = True; m.flat[j] = False
return best[1], best[0]
mask_erg, f_erg = ergodic_mask()
f_rand = np.mean([objective(random_mask()) for _ in range(20)])
print(f"N_samples={NS} (delta={DELTA:.3f}) obj: sparse={objective(mask_reg):.3f} "
f"random(mean)={f_rand:.3f} ergodic={f_erg:.3f}")
# ---------------- transforms ----------------
_, SLICES = pywt.coeffs_to_array(pywt.wavedec2(np.zeros((NX, NY)), WAVELET, mode=MODE, level=LEVELS))
def Wsym(x): return pywt.coeffs_to_array(pywt.wavedec2(x, WAVELET, mode=MODE, level=LEVELS))[0]
def Wtsym(c): return pywt.waverec2(pywt.array_to_coeffs(c, SLICES, output_format='wavedec2'), WAVELET, mode=MODE)
SYM_APP = np.zeros((NX, NY), bool); SYM_APP[SLICES[0]] = True # approximation subband
def Wfft(x): return np.fft.fft2(x)
def Wtfft(c): return np.fft.ifft2(c).real
# ---------------- CS reconstruction: iterative thresholding + data reinsertion (POCS) ----------------
def cs_reconstruct(y_img, mask, W=Wfft, Wt=Wtfft, n_iter=POCS_ITERS, lam_min_frac=POCS_LMIN):
x = y_img.copy()
lam_max = np.abs(W(y_img)).max()
for lam in np.geomspace(lam_max, lam_max*lam_min_frac, n_iter):
c = W(x)
c *= np.maximum(1 - lam/np.maximum(np.abs(c), 1e-300), 0) # soft threshold on magnitude
x = Wt(c)
x[mask] = y_img[mask]
return x
# symlet variant (the paper's transform), tuned to reproduce the paper's accuracy in its regime:
# never threshold the approximation subband, and average reconstructions over circular shifts
# (cycle spinning) to remove the decimated transform's shift-variance
def cs_reconstruct_sym(y_img, mask, n_iter=POCS_ITERS, lam_min_frac=POCS_LMIN, n_shift=SYM_SPINS):
def single(y2, m2):
x = y2.copy()
lam_max = np.abs(np.where(SYM_APP, 0, Wsym(y2))).max()
for lam in np.geomspace(lam_max, lam_max*lam_min_frac, n_iter):
c = Wsym(x)
t = np.sign(c)*np.maximum(np.abs(c) - lam, 0)
x = Wtsym(np.where(SYM_APP, c, t))
x[m2] = y2[m2]
return x
recs = []
for _ in range(n_shift):
s = (int(rng.integers(NX)), int(rng.integers(NY)))
xs = single(np.roll(y_img, s, (0, 1)), np.roll(mask, s, (0, 1)))
recs.append(np.roll(xs, (-s[0], -s[1]), (0, 1)))
x = np.mean(recs, 0)
x[mask] = y_img[mask]
return x
# ---------------- kriging / GP: Bayes-optimal for the power-law field class ----------------
KXF = np.fft.fftfreq(NX)[:, None]; KYF = np.fft.fftfreq(NY)[None, :]
KRAD = np.hypot(KXF, KYF)
GRID = np.argwhere(np.ones((NX, NY), bool))
def circulant_cov(beta, z_src):
# power spectrum k^-beta * exp(-4*pi*k*z): the amplitude-attenuated generator's power spectrum
S = np.where(KRAD > 0, KRAD, 1.0)**(-beta)*np.exp(-4*np.pi*KRAD*z_src); S[0, 0] = 0.0
c = np.fft.ifft2(S).real
return c/c[0, 0]
def gp_matrices(mask, beta, z_src):
on = np.argwhere(mask)
cov = circulant_cov(beta, z_src)
Kss = cov[(on[:, None, 0]-on[None, :, 0]) % NX, (on[:, None, 1]-on[None, :, 1]) % NY]
Kxs = cov[(GRID[:, None, 0]-on[None, :, 0]) % NX, (GRID[:, None, 1]-on[None, :, 1]) % NY]
return Kss + 1e-6*np.eye(len(on)), Kxs
def gp_fit_params(y_img, mask):
# truth-free ML over spectral slope and source depth, variance profiled out
# (only the sample-sample covariance is needed here, not the full-grid cross-covariance)
yv = y_img[mask]; best = None
on = np.argwhere(mask)
for beta in np.arange(1.5, 4.01, 0.5):
for z_src in (0.0, 0.5, 1.0, 1.5, 2.0, 3.0):
cov = circulant_cov(beta, z_src)
Kss = cov[(on[:, None, 0]-on[None, :, 0]) % NX,
(on[:, None, 1]-on[None, :, 1]) % NY] + 1e-6*np.eye(len(on))
try:
L = np.linalg.cholesky(Kss)
except np.linalg.LinAlgError:
continue
a = np.linalg.solve(L, yv)
ll = -0.5*len(yv)*np.log(a @ a/len(yv)) - np.log(np.diag(L)).sum()
if best is None or ll > best[0]: best = (ll, beta, z_src)
return best[1], best[2]
def gp_predictor(mask, beta, z_src):
# precomputed kriging weights: with the mask and model fixed, reconstruction is a single
# matrix multiplication per field (the second argument exists only for signature parity)
Kss, Kxs = gp_matrices(mask, beta, z_src)
Wm = Kxs @ np.linalg.inv(Kss)
return lambda y_img, m=None: (Wm @ y_img[mask]).reshape(NX, NY)
def gp_reconstruct(y_img, mask, beta, z_src):
yv = y_img[mask]
Kss, Kxs = gp_matrices(mask, beta, z_src)
w = np.linalg.solve(Kss, yv)
xh = (Kxs @ w).reshape(NX, NY)
A = np.linalg.solve(Kss, Kxs.T)
var = np.maximum(1.0 - np.einsum('ij,ji->i', Kxs, A), 0)
s2 = yv @ np.linalg.solve(Kss, yv)/len(yv) # profiled variance scale
return xh, np.sqrt(var*s2).reshape(NX, NY)
def apriori_level_correction(beta, z_src, mask, n_top=400, m_draws=20000):
# The potential-field generator uses fixed-amplitude random phases and divides each
# realization by its sample standard deviation. For a linear reconstruction the resulting
# error variance is E[X/Y] with X = sum_p t_p u_p (per-mode error contributions) and
# Y = sum_p s_p u_p (the realization variance), where u_p = 1 + cos(theta_p) are the random
# mode powers. Both expectations involve only scalar mode powers, so the level bias of the
# posterior sigma is computable a priori from the fitted spectrum and the mask alone.
Sfit = np.where(KRAD > 0, KRAD, 1.0)**(-beta)*np.exp(-4*np.pi*KRAD*z_src); Sfit[0, 0] = 0.0
Sn = (Sfit/Sfit.sum()).ravel()
idx = np.arange(N); ix, iy = np.divmod(idx, NY)
neg = ((-ix) % NX)*NY + ((-iy) % NY)
take = ((idx < neg) | (idx == neg)) & (Sn > 0)
s_pair = np.where(idx == neg, Sn, 2*Sn)[take]
kxp, kyp = ix[take], iy[take]
o = np.argsort(s_pair)[::-1]
s_pair, kxp, kyp = s_pair[o], kxp[o], kyp[o]
rs = np.random.default_rng(123) # local stream, leaves rng untouched
u_top = 1 + np.cos(2*np.pi*rs.random((m_draws, n_top)))
Y = u_top @ s_pair[:n_top]
for j0 in range(n_top, len(s_pair), 512):
blk = s_pair[j0:j0+512]
Y += (1 + np.cos(2*np.pi*rs.random((m_draws, len(blk))))) @ blk
g_top = (u_top/Y[:, None]).mean(0) # E[u_p / Y] for the dominant modes
g_tail = (1/Y).mean() # small modes are independent of Y
Kss_c, Kxs_c = gp_matrices(mask, beta, z_src)
Wc = Kxs_c @ np.linalg.inv(Kss_c)
sig2_u = np.maximum(1.0 - np.einsum('ij,ji->i', Kxs_c, np.linalg.solve(Kss_c, Kxs_c.T)), 0)
ph = np.exp(2j*np.pi*(np.outer(GRID[:, 0], kxp[:n_top])/NX
+ np.outer(GRID[:, 1], kyp[:n_top])/NY))
resp = Wc @ ph[mask.ravel()] - ph # error response of each dominant mode
t_top = (np.abs(resp)**2)*s_pair[:n_top]/2
t_tail = np.maximum(sig2_u - t_top.sum(1), 0)
sig2_c = t_top @ g_top + t_tail*g_tail
return np.sqrt(sig2_c/np.maximum(sig2_u, 1e-30)).reshape(NX, NY)
# ---------------- acquire + reconstruct ----------------
def leakage(x, truth): return np.corrcoef((x - truth).ravel(), truth.ravel())[0, 1]
def rms(x, truth): return np.sqrt(np.mean((x - truth)**2))
def spline_interp(truth, mask):
pts = np.argwhere(mask); X, Y = np.mgrid[0:NX, 0:NY]
xr = griddata(pts, truth[mask], (X, Y), method='cubic')
nn = np.isnan(xr)
xr[nn] = griddata(pts, truth[mask], (X, Y), method='nearest')[nn]
return xr
def fourier_oracle(truth, k=NS):
C = np.fft.fft2(truth).ravel()
C[np.argsort(np.abs(C))[:-k]] = 0
return np.fft.ifft2(C.reshape(NX, NY)).real
# tier 1: Fourier-sparse truth: the CS claim of the paper (Fig. A.1 / Figs. 15, 22)
print("\n--- tier 1: Fourier-sparse truth (paper's CS regime) ---")
x_sp_spl = spline_interp(truth_sp, mask_reg)
x_sp_erg_spl = spline_interp(truth_sp, mask_erg)
x_sp_reg = cs_reconstruct(truth_sp*mask_reg, mask_reg)
x_sp_erg = cs_reconstruct(truth_sp*mask_erg, mask_erg)
x_sp_sym = cs_reconstruct_sym(truth_sp*mask_erg, mask_erg) # paper's literal transform
x_sp_reg_sym = cs_reconstruct_sym(truth_sp*mask_reg, mask_reg)
# kriging with the power-law covariance is misspecified for a 12-mode signal; shown as the
# symmetric counterpart to CS-on-the-random-field so every process appears on both truths
GP_SP = gp_fit_params(truth_sp*mask_erg, mask_erg)
x_sp_gp, sig_sp_raw = gp_reconstruct(truth_sp*mask_erg, mask_erg, *GP_SP)
x_sp_reg_gp = gp_reconstruct(truth_sp*mask_reg, mask_reg, *GP_SP)[0]
for name, x in [('sparse+spline', x_sp_spl), ('sparse+CS', x_sp_reg),
('sparse+CS-symlet', x_sp_reg_sym), ('sparse+kriging', x_sp_reg_gp),
('ergodic+spline', x_sp_erg_spl), ('ergodic+CS', x_sp_erg),
('ergodic+CS-symlet', x_sp_sym), ('ergodic+kriging', x_sp_gp)]:
print(f"{name:18s} leakage={leakage(x, truth_sp):+.3f} rms={rms(x, truth_sp):.3f}")
# tier 2: potential-field truth: not l1-sparse; report against the oracle floor
print(f"\n--- tier 2: potential-field truth (beta={BETA}, source depth {Z_SRC}) ---")
print(f"oracle {NS}-term Fourier floor: rms={rms(fourier_oracle(truth_pl), truth_pl):.3f}")
BETA_HAT, Z_HAT = gp_fit_params(truth_pl*mask_erg, mask_erg)
print(f"GP ML estimates: beta_hat={BETA_HAT:.1f} z_hat={Z_HAT:.1f} "
f"(generator beta={BETA}, z={Z_SRC})")
CORR_PL = apriori_level_correction(BETA_HAT, Z_HAT, mask_erg)
print(f"a-priori sigma level correction for the normalized truth generator: "
f"x{CORR_PL[~mask_erg].mean():.4f} (spread {CORR_PL[~mask_erg].std():.4f})")
y_erg = truth_pl*mask_erg
x_pl_spl = spline_interp(truth_pl, mask_reg)
x_pl_erg_spl = spline_interp(truth_pl, mask_erg)
x_pl_reg = cs_reconstruct(truth_pl*mask_reg, mask_reg)
x_pl_reg_sym = cs_reconstruct_sym(truth_pl*mask_reg, mask_reg)
x_pl_reg_gp = gp_reconstruct(truth_pl*mask_reg, mask_reg, BETA_HAT, Z_HAT)[0]
x_erg = cs_reconstruct(y_erg, mask_erg) # primary reconstruction (paper's method)
x_erg_sym = cs_reconstruct_sym(y_erg, mask_erg) # symlet-wavelet comparison
x_gp, sig_gp_raw = gp_reconstruct(y_erg, mask_erg, BETA_HAT, Z_HAT)
print("method leakage rms -rms/std(truth) (leakage ~ -rms/std when error is unrecovered signal)")
for name, x in [('sparse+spline', x_pl_spl),
('sparse+CS', x_pl_reg),
('sparse+CS-symlet', x_pl_reg_sym),
('sparse+kriging', x_pl_reg_gp),
('ergodic+spline', x_pl_erg_spl),
('ergodic+CS-Fourier', x_erg),
('ergodic+CS-symlet', x_erg_sym),
('ergodic+kriging', x_gp)]:
print(f"{name:18s} {leakage(x, truth_pl):+.3f} {rms(x, truth_pl):.3f} {-rms(x, truth_pl)/truth_pl.std():+.3f}")
# ---------------- predicted sigma (truth-free: from y + mask only) ----------------
# LOO calibration: reconstruct without one sample, compare prediction at that point to its value
def loo_rms(y_img, mask, rec_fn, m=LOO_M):
on = np.argwhere(mask)
rs = []
for i in rng.choice(len(on), m, replace=False):
m2 = mask.copy(); m2[tuple(on[i])] = False
rs.append(rec_fn(y_img*m2, m2)[tuple(on[i])] - y_img[tuple(on[i])])
return np.sqrt(np.mean(np.square(rs)))
def recalibrate(sig, loo):
return sig*loo/np.sqrt(np.mean(sig**2))
# a single realization |err| = sigma*|z| caps the achievable correlation even for a perfect map:
def pearson_ceiling(s):
return np.sqrt((2/np.pi)*s.var()/((s**2).mean() - (2/np.pi)*s.mean()**2))
def spearman_ceiling(s, ndraw=50):
return np.mean([spearmanr(s.ravel()*np.abs(rng.standard_normal(s.size)), s.ravel())[0]
for _ in range(ndraw)])
OFF = ~mask_erg # evaluate z-scores at unsampled pixels only (err and sigma are both 0 at samples)
def reliability(s, a, nb=10, reduce=np.mean):
q = np.quantile(s, np.linspace(0, 1, nb+1)); q[-1] += 1e-12
m, r = [], []
for lo, hi in zip(q[:-1], q[1:]):
sel = (s >= lo) & (s < hi)
if sel.any():
m.append(s[sel].mean()); r.append(reduce(a[sel]))
return np.array(m), np.array(r)
def rms_reduce(a): return np.sqrt(np.mean(a**2))
# ---------------- plots (no abbreviations in any figure text) ----------------
C_ENS, C_ONE = '#1f77b4', '#ff7f0e' # blue = ensemble comparison, orange = single realization
# --- figures 1a/1b: one figure per ground truth: reconstructions, errors, sampling locations ---
COLS = ['Ground truth',
'Sparse + spline interpolation\n(naive control)',
'Sparse + compressive sensing\n(Fourier)',
'Sparse + compressive sensing\n(symlet wavelet, from paper)',
'Sparse + kriging\n(Gaussian process)',
'Ergodic + spline interpolation',
'Ergodic + compressive sensing\n(Fourier)',
'Ergodic + compressive sensing\n(symlet wavelet, from paper)',
'Ergodic + kriging\n(Gaussian process)']
COL_MASKS = [np.ones_like(mask_reg), mask_reg, mask_reg, mask_reg, mask_reg,
mask_erg, mask_erg, mask_erg, mask_erg]
for fname, gt, imgs, note in [
('plt1a', truth_sp,
[truth_sp, x_sp_spl, x_sp_reg, x_sp_reg_sym, x_sp_reg_gp,
x_sp_erg_spl, x_sp_erg, x_sp_sym, x_sp_gp],
f'Fourier-sparse truth ({N_MODES} modes); compressive sensing with the ergodic pattern '
'recovers the signal exactly'),
('plt1b', truth_pl,
[truth_pl, x_pl_spl, x_pl_reg, x_pl_reg_sym, x_pl_reg_gp,
x_pl_erg_spl, x_erg, x_erg_sym, x_gp],
f'potential-field truth (spectral exponent {BETA}, source depth {Z_SRC} cells); '
'kriging is the best reconstruction for this field class')]:
fig, axes = plt.subplots(3, 9, figsize=(36, 12.4), constrained_layout=True)
v = np.abs(gt).max()
for c, (a, x, t) in enumerate(zip(axes[0], imgs, COLS)):
im = a.imshow(x.T, origin='lower', vmin=-v, vmax=v, cmap='RdBu_r')
sub = '' if c == 0 else (f'\nleakage {leakage(x, gt):+.2f} · '
f'RMS error {rms(x, gt):.3f}')
a.set_title(t + sub, fontsize=9)
axes[0, 0].set_ylabel('Reconstruction\n\ngrid y (cells)')
fig.colorbar(im, ax=axes[0], shrink=0.9,
label='field value\n(units of the truth standard deviation)')
errors = [gt - x for x in imgs[1:]]
ve = max(np.abs(e).max() for e in errors)
axes[1, 0].axis('off')
for a, e in zip(axes[1, 1:], errors):
im2 = a.imshow(e.T, origin='lower', vmin=-ve, vmax=ve, cmap='RdBu_r')
axes[1, 1].set_ylabel('Error: ground truth reconstruction\n\ngrid y (cells)')
fig.colorbar(im2, ax=axes[1, 1:], shrink=0.9,
label='error\n(units of the truth standard deviation)')
for a, mm in zip(axes[2], COL_MASKS):
a.scatter(*np.argwhere(mm).T, s=2, c='k')
a.set_aspect(1); a.set_xlim(-1, NX); a.set_ylim(-1, NY)
a.set_title(f'{int(mm.sum())} of {N} samples', fontsize=10)
a.set_xlabel('grid x (cells)')
axes[2, 0].set_ylabel('Sampling locations\n\ngrid y (cells)')
fig.suptitle(f'Reconstruction from {DELTA:.1%} of the samples: {note}')
plt.savefig(f'./{fname}.png', dpi=120)
# --- figures 2a1/2b1 (kriging) and 2a2/2b2 (symlet compressive sensing): predicted standard
# --- deviation versus actual error, one figure per ground truth and reconstruction method ---
def sigma_study(fname, desc, make_truth, truth, x_rec, sig_raw, rec_fn, rec_name, corr=None,
ens_rec_fn=None):
y_img = truth*mask_erg
loo = loo_rms(y_img, mask_erg, rec_fn)
if corr is None:
sig = recalibrate(sig_raw, loo); level_note = 'sigma level set by leave-one-out'
else:
sig = sig_raw*corr; level_note = 'sigma level corrected a priori'
err = x_rec - truth; aerr = np.abs(err)
print(f"\n--- sigma study ({fname}): {desc} ---")
print(f"LOO error rms (truth-free check): {loo:.3f} "
f"actual rms: {np.sqrt((err**2).mean()):.3f} ({level_note})")
pr = pearsonr(aerr.ravel(), sig.ravel())[0]; sr = spearmanr(aerr.ravel(), sig.ravel())[0]
pc, sc = pearson_ceiling(sig), spearman_ceiling(sig)
print(f"|err| vs sigma_GP: pearson={pr:+.3f} (ceiling {pc:.3f}, ratio {pr/pc:+.2f}) "
f"spearman={sr:+.3f} (ceiling {sc:.3f}, ratio {sr/sc:+.2f}) "
f"z-std={np.std(err[OFF]/sig[OFF]):.2f}")
ens_rec = ens_rec_fn if ens_rec_fn is not None else rec_fn
ens_errs = np.array([ens_rec(t*mask_erg, mask_erg) - t
for t in (make_truth() for _ in range(ENS_K))])
sig_emp = np.sqrt((ens_errs**2).mean(0))
pr2 = pearsonr(sig_emp.ravel(), sig.ravel())[0]; sr2 = spearmanr(sig_emp.ravel(), sig.ravel())[0]
print(f"ensemble (K={ENS_K} truths, same mask; sigma_emp has ~{1/np.sqrt(2*ENS_K):.0%} noise): "
f"pearson={pr2:+.3f} spearman={sr2:+.3f} "
f"pooled z-std={np.std(ens_errs[:, OFF]/sig[OFF]):.2f}")
fig, axes = plt.subplots(3, 3, figsize=(14.5, 13.8), constrained_layout=True)
# fixed axes (3.0 truth standard deviations) so the plt2 figures are directly comparable;
# some values clip by design
v = 3.0
im0 = axes[0, 0].imshow(truth.T, origin='lower', vmin=-v, vmax=v, cmap='RdBu_r')
axes[0, 0].set_title('Ground truth', fontsize=10)
axes[0, 1].imshow(x_rec.T, origin='lower', vmin=-v, vmax=v, cmap='RdBu_r')
axes[0, 1].set_title(f'Reconstruction: ergodic + {rec_name}\n'
f'leakage {leakage(x_rec, truth):+.2f} · RMS error '
f'{rms(x_rec, truth):.3f}', fontsize=10)
axes[0, 2].imshow(err.T, origin='lower', vmin=-v, vmax=v, cmap='RdBu_r')
axes[0, 2].set_title('Signed error:\nreconstruction - truth', fontsize=10)
fig.colorbar(im0, ax=axes[0], shrink=0.9,
label='field value or error\n(units of the truth standard deviation)')
for a, x, t in zip(axes[1],
[sig, sig_emp, aerr],
['Predicted standard deviation:\nGaussian process posterior (truth-free)',
f'Empirical standard deviation\n({ENS_K} independent truths, '
'same sampling pattern)',
'Actual absolute error, this realization:\n|reconstruction - truth|']):
im1 = a.imshow(x.T, origin='lower', vmin=0, vmax=v, cmap='turbo')
a.set_title(t, fontsize=10)
a.set_xlabel('grid x (cells)')
fig.colorbar(im1, ax=axes[1], shrink=0.9,
label='standard deviation or absolute error\n(units of the truth standard deviation)')
for r in (0, 1):
axes[r, 0].set_ylabel('grid y (cells)')
a = axes[2, 0] # the sampling pattern that produced everything above
a.scatter(*np.argwhere(mask_erg).T, s=2, c='k')
a.set_aspect(1); a.set_xlim(-1, NX); a.set_ylim(-1, NY)
a.set_title(f'Ergodic sampling locations\n({NS} of {N} samples)', fontsize=10)
a.set_xlabel('grid x (cells)'); a.set_ylabel('grid y (cells)')
a = axes[2, 1] # scatter with decile-binned means: does the prediction rank the error?
a.scatter(sig[OFF], aerr[OFF], s=2, alpha=.2, color='#9ecae1')
bm, bv = reliability(sig[OFF], aerr[OFF])
a.plot(bm, bv, 'o-', color='k', lw=1.6)
xs = np.linspace(0, sig[OFF].max(), 50)
a.plot(xs, np.sqrt(2/np.pi)*xs, '--', color='k', lw=1.2)
a.set_xlabel('predicted standard deviation'); a.set_ylabel('absolute error, this realization')
a.set_ylim(0, v)
a.ticklabel_format(style='sci', scilimits=(-2, 3))
a.set_title('Prediction versus error at unsampled cells\n(single-realization Pearson '
f'correlation is capped\nat about {pc:.2f} even for a perfect prediction)',
fontsize=10)
a.legend(handles=[plt.Line2D([], [], marker='o', ls='none', ms=4, color='#9ecae1',
label='unsampled cells'),
plt.Line2D([], [], marker='o', ls='-', ms=5, color='k',
label='mean absolute error\nwithin prediction decile'),
plt.Line2D([], [], ls='--', color='k',
label='expected mean absolute error\nfor a perfect prediction')],
loc='upper left', fontsize=8)
a = axes[2, 2] # merged calibration panel: against one realization and against the ensemble
for target, lab, cc in [(aerr, 'against this single realization', C_ONE),
(sig_emp, f'against the {ENS_K}-truth ensemble', C_ENS)]:
mm, rr = reliability(sig, target, reduce=rms_reduce)
a.plot(mm, rr, 'o-', color=cc, label=lab)
a.plot([0, 1], [0, 1], 'k--', lw=.8, label='perfect calibration')
a.set_xlim(0, 1.0); a.set_ylim(0, 1.0); a.set_aspect(1)
a.set_xlabel('mean predicted standard deviation within decile')
a.set_ylabel('RMS error or\nempirical standard deviation within decile')
a.ticklabel_format(style='sci', scilimits=(-2, 3))
a.set_title('Calibration of the Gaussian process prediction\n(cells binned into deciles of '
'the prediction)', fontsize=10)
a.legend(loc='upper left', fontsize=8)
fig.suptitle('Statistical analysis of reconstruction error\n'
f'({desc}; ergodic sampling; {rec_name} reconstruction)\n'
,fontsize=11)
plt.savefig(f'./{fname}.png', dpi=120)
DESC_SP = f'Fourier-sparse truth ({N_MODES} modes)'
DESC_PL = f'potential-field truth (spectral exponent {BETA}, source depth {Z_SRC} cells)'
REC_GP, REC_SYM = 'kriging (Gaussian process)', 'compressive sensing (symlet wavelet, from paper)'
KRIG_FAST_SP = gp_predictor(mask_erg, *GP_SP)
KRIG_FAST_PL = gp_predictor(mask_erg, BETA_HAT, Z_HAT)
sigma_study('plt2a1', DESC_SP, lambda: fourier_sparse_field(N_MODES), truth_sp, x_sp_gp, sig_sp_raw,
lambda y, m: gp_reconstruct(y, m, *GP_SP)[0], REC_GP, ens_rec_fn=KRIG_FAST_SP)
sigma_study('plt2a2', DESC_SP, lambda: fourier_sparse_field(N_MODES), truth_sp, x_sp_sym, sig_sp_raw,
cs_reconstruct_sym, REC_SYM)
sigma_study('plt2b1', DESC_PL, lambda: powerlaw_field(BETA), truth_pl, x_gp, sig_gp_raw,
lambda y, m: gp_reconstruct(y, m, BETA_HAT, Z_HAT)[0], REC_GP, corr=CORR_PL,
ens_rec_fn=KRIG_FAST_PL)
sigma_study('plt2b2', DESC_PL, lambda: powerlaw_field(BETA), truth_pl, x_erg_sym, sig_gp_raw,
cs_reconstruct_sym, REC_SYM)
# ---------------- Monte-Carlo RMS table over truth classes ----------------
N_TRIALS = 25
def white_noise_field():
# incompressible control: no method can beat predicting zero at unsampled cells,
# whose rms is sqrt(1 - DELTA) ~ 0.94
f = rng.standard_normal((NX, NY))
return (f - f.mean())/f.std()
def trial_methods(truth, krig_sparse, krig_erg):
return {
'sparse+spline': spline_interp(truth, mask_reg),
'sparse+CS-Fourier': cs_reconstruct(truth*mask_reg, mask_reg),
'sparse+kriging': krig_sparse(truth*mask_reg),
'ergodic+spline': spline_interp(truth, mask_erg),
'ergodic+CS-Fourier': cs_reconstruct(truth*mask_erg, mask_erg),
'ergodic+CS-symlet': cs_reconstruct_sym(truth*mask_erg, mask_erg),
'ergodic+kriging': krig_erg(truth*mask_erg),
}
CLASSES = [('white noise (control)', white_noise_field),
('Fourier-sparse', lambda: fourier_sparse_field(N_MODES)),
('potential-field', lambda: powerlaw_field(BETA))]
print(f"\n--- Monte-Carlo rms error, mean +/- std over {N_TRIALS} trials per truth class ---")
results = {}
for cname, gen in CLASSES:
t0 = gen()
gp_params = gp_fit_params(t0*mask_erg, mask_erg) # truth-free model selection, once per class
krig_sparse = gp_predictor(mask_reg, *gp_params) # weights precomputed once per class
krig_erg = gp_predictor(mask_erg, *gp_params)
rr = {}
for k in range(N_TRIALS):
truth = t0 if k == 0 else gen()
for m, x in trial_methods(truth, krig_sparse, krig_erg).items():
rr.setdefault(m, []).append(rms(x, truth))
results[cname] = {m: (np.mean(v), np.std(v)) for m, v in rr.items()}
print('method'.ljust(20) + ''.join(c.rjust(24) for c, _ in CLASSES))
for m in results[CLASSES[0][0]]:
print(m.ljust(20) + ''.join(f"{results[c][m][0]:.3f} +/- {results[c][m][1]:.3f}".rjust(24)
for c, _ in CLASSES))
# --- figure 3: the Monte-Carlo table as a chart with error bars ---
DISPLAY = {'sparse+spline': 'Sparse + spline interpolation',
'sparse+CS-Fourier': 'Sparse + compressive sensing (Fourier)',
'sparse+kriging': 'Sparse + kriging (Gaussian process)',
'ergodic+spline': 'Ergodic + spline interpolation',
'ergodic+CS-Fourier': 'Ergodic + compressive sensing (Fourier)',
'ergodic+CS-symlet': 'Ergodic + compressive sensing (symlet wavelet)',
'ergodic+kriging': 'Ergodic + kriging (Gaussian process)'}
methods = list(DISPLAY)
ypos = np.arange(len(methods))[::-1]
xmax = max(results[c][m][0] + results[c][m][1] for c, _ in CLASSES for m in methods)
fig, axes = plt.subplots(1, 3, figsize=(15.5, 5.8), sharey=True, constrained_layout=True)
for a, (cname, _) in zip(axes, CLASSES):
mus = np.array([results[cname][m][0] for m in methods])
sds = np.array([results[cname][m][1] for m in methods])
# bar color encodes the sampling pattern (blue = sparse grid, orange = ergodic); the method
# labels already carry this, so no legend is needed
a.barh(ypos, mus, xerr=sds, height=0.6,
color=['#9ecae1' if m.startswith('sparse') else '#fdbe85' for m in methods],
error_kw=dict(ecolor='k', lw=1, capsize=3))
for yp, mu, sd in zip(ypos, mus, sds):
a.text(mu + sd + 0.02*xmax, yp, f'{mu:.3f} ± {sd:.3f}', va='center', fontsize=8)
t = cname[0].upper() + cname[1:] + ' truth'*(not cname.endswith(')'))
if cname == CLASSES[0][0]:
t += f'\n(dashed line: best possible = predict zero everywhere, {np.sqrt(1 - DELTA):.2f})'
a.set_title(t, fontsize=10)
a.set_xlabel('RMS error\n(units of the truth standard deviation)')
a.set_xlim(0, xmax*1.32)
a.grid(axis='x', lw=0.4, alpha=0.4)
a.set_axisbelow(True)
axes[0].set_yticks(ypos, [DISPLAY[m] for m in methods])
axes[0].axvline(np.sqrt(1 - DELTA), ls='--', lw=1, color='k')
fig.suptitle('Monte-Carlo RMS reconstruction error: '
f'mean ± standard deviation over {N_TRIALS} trials per truth class '
f'({DELTA:.1%} of samples, error bars = ± one standard deviation)')
plt.savefig('./plt3.png', dpi=120)
# ---------------- sanity check: is the Gaussian process sigma statistically valid? ----------------
# With the mask fixed, kriging is a linear map, so a large ensemble is cheap: precompute the
# weight matrix once and run SANITY_K fresh truths through it. Under the hypothesis that the
# posterior standard deviation is correct, the empirical standard deviation over K trials should
# differ from it only by sampling noise (relative size about 1/sqrt(2K) per cell). The unsampled
# cells are spatially correlated within a single realization, so the honest global test averages
# the normalized squared error within each trial and treats the SANITY_K per-trial means as the
# independent observations.
RUN_SANITY_CHECK = False # plt4 served its purpose; flip to True to re-run the check
if RUN_SANITY_CHECK:
from scipy.stats import norm as normal_dist
SANITY_K = 512
Kss_s, Kxs_s = gp_matrices(mask_erg, BETA_HAT, Z_HAT)
W_krig = Kxs_s @ np.linalg.inv(Kss_s)
var_unit = np.maximum(1.0 - np.einsum('ij,ji->i', Kxs_s, np.linalg.solve(Kss_s, Kxs_s.T)), 0)
sig_pred_unit = np.sqrt(var_unit).reshape(NX, NY) # model posterior sigma, unit-variance field
sanity_errs = np.empty((SANITY_K, NX, NY))
for k in range(SANITY_K):
t = powerlaw_field(BETA)
sanity_errs[k] = (W_krig @ t[mask_erg]).reshape(NX, NY) - t
sig_emp_big = np.sqrt((sanity_errs**2).mean(0))
resid = sig_emp_big - sig_pred_unit
# per-trial mean of error^2/sigma^2 over unsampled cells: one independent number per trial
m_k = (sanity_errs[:, OFF]**2/np.maximum(sig_pred_unit[OFF]**2, 1e-12)).mean(1)
z_glob = (m_k.mean() - 1)/(m_k.std(ddof=1)/np.sqrt(SANITY_K))
p_glob = 2*normal_dist.sf(abs(z_glob))
ratio = sig_emp_big[OFF]/sig_pred_unit[OFF]
print(f"\n--- sanity check: Gaussian process sigma vs {SANITY_K}-trial empirical sigma ---")
print(f"mean of per-trial normalized squared error: {m_k.mean():.4f} (1.0 if sigma is correct)")
print(f"global test over {SANITY_K} independent trials: z = {z_glob:+.2f}, "
f"two-sided p = {p_glob:.3f}")
print(f"per-cell ratio sigma_empirical/sigma_predicted at {OFF.sum()} unsampled cells: "
f"mean {ratio.mean():.4f}, spread {ratio.std():.4f} "
f"(pure sampling noise predicts spread ~ {1/np.sqrt(2*SANITY_K):.4f})")
corr_sanity = pearsonr(sig_pred_unit[OFF], sig_emp_big[OFF])[0]
print(f"correlation(sigma_predicted, sigma_empirical) = {corr_sanity:.4f}")
# control: truths drawn exactly from the fitted Gaussian model with a deterministic scale. The
# pipeline generator uses fixed-amplitude random phases and normalizes each realization by its
# sample standard deviation, which perturbs the ensemble covariance away from the model; the
# exact-model truths isolate the question "is the posterior itself constructed correctly?"
SQ_S = np.sqrt(np.where(KRAD > 0, KRAD, 1.0)**(-BETA_HAT)*np.exp(-4*np.pi*KRAD*Z_HAT))
SQ_S[0, 0] = 0.0
C0_S = np.fft.ifft2(SQ_S**2).real[0, 0]
ctrl_errs = np.empty((SANITY_K, NX, NY))
for k in range(SANITY_K):
t = np.fft.ifft2(SQ_S*np.fft.fft2(rng.standard_normal((NX, NY)))).real/np.sqrt(C0_S)
ctrl_errs[k] = (W_krig @ t[mask_erg]).reshape(NX, NY) - t
m2_k = (ctrl_errs[:, OFF]**2/np.maximum(sig_pred_unit[OFF]**2, 1e-12)).mean(1)
z_ctrl = (m2_k.mean() - 1)/(m2_k.std(ddof=1)/np.sqrt(SANITY_K))
p_ctrl = 2*normal_dist.sf(abs(z_ctrl))
ratio_ctrl = np.sqrt((ctrl_errs**2).mean(0))[OFF]/sig_pred_unit[OFF]
print(f"control with exact-model truths (no per-realization normalization): "
f"mean ratio {m2_k.mean():.4f}, z = {z_ctrl:+.2f}, two-sided p = {p_ctrl:.3f}")
# same pipeline ensemble scored against the a-priori corrected sigma
sig_corr_unit = sig_pred_unit*CORR_PL
m3_k = (sanity_errs[:, OFF]**2/np.maximum(sig_corr_unit[OFF]**2, 1e-12)).mean(1)
z_corr = (m3_k.mean() - 1)/(m3_k.std(ddof=1)/np.sqrt(SANITY_K))
p_corr = 2*normal_dist.sf(abs(z_corr))
ratio_corr = np.sqrt((sanity_errs**2).mean(0))[OFF]/sig_corr_unit[OFF]
print(f"pipeline truths with the a-priori level correction: "
f"mean ratio {m3_k.mean():.4f}, z = {z_corr:+.2f}, two-sided p = {p_corr:.3f}")
fig, axes = plt.subplots(1, 3, figsize=(16.5, 4.9), constrained_layout=True)
vr = np.abs(resid).max()
imR = axes[0].imshow(np.where(OFF, resid, np.nan).T, origin='lower', vmin=-vr, vmax=vr,
cmap='RdBu_r')
axes[0].set_title(f'Empirical ({SANITY_K} trials) predicted standard deviation\n'
'(sampled cells blanked; structure here would mean a construction error)',
fontsize=10)
axes[0].set_xlabel('grid x (cells)'); axes[0].set_ylabel('grid y (cells)')
fig.colorbar(imR, ax=axes[0], shrink=0.9,
label='difference\n(units of the truth standard deviation)')
a = axes[1]
a.scatter(sig_pred_unit[OFF], sig_emp_big[OFF], s=2, alpha=.25, color='#9ecae1')
lim = max(sig_pred_unit[OFF].max(), sig_emp_big[OFF].max())*1.05
a.plot([0, lim], [0, lim], 'k--', lw=1, label='equality')
a.set_xlim(0, lim); a.set_ylim(0, lim); a.set_aspect(1)
a.set_xlabel('predicted standard deviation')
a.set_ylabel(f'empirical standard deviation ({SANITY_K} trials)')
a.set_title('Per-cell agreement at unsampled cells\n'
f'(correlation {corr_sanity:.3f})', fontsize=10)
a.legend(loc='upper left', fontsize=8)
a = axes[2]
zn = np.sqrt(2*SANITY_K)*(ratio - 1) # approximately standard normal if sigma is correct
zn_ctrl = np.sqrt(2*SANITY_K)*(ratio_ctrl - 1)
zn_corr = np.sqrt(2*SANITY_K)*(ratio_corr - 1)
a.hist(zn, bins=50, density=True, color='#9ecae1', label='pipeline truths')
a.hist(zn_corr, bins=50, density=True, histtype='step', color='#2ca02c', lw=1.5,
label='pipeline truths,\na-priori corrected sigma')
a.hist(zn_ctrl, bins=50, density=True, histtype='step', color='#ff7f0e', lw=1.5,
label='exact-model truths (control)')
xs = np.linspace(min(zn_corr.min(), -4.5), max(zn.max(), 4.5), 200)
a.plot(xs, normal_dist.pdf(xs), 'k-', lw=1.2, label='expected from sampling\nnoise alone')
a.set_xlabel('normalized per-cell deviation of the ratio\nempirical/predicted standard deviation')
a.set_ylabel('probability density')
a.set_title(f'Deviations vs pure sampling noise\n(pipeline: z = {z_glob:+.2f}; corrected: '
f'z = {z_corr:+.2f}; control: z = {z_ctrl:+.2f})', fontsize=10)
a.legend(loc='upper right', fontsize=8)
fig.suptitle('Validity check of the Gaussian process standard deviation over '
f'{SANITY_K} independent truths: shape agreement {corr_sanity:.3f}.\n'
'The level offset for pipeline truths comes from their per-realization '
'normalization, not from the posterior: it is computable a priori from the '
'spectrum and mask,\nand correcting for it centers the test '
f'(z = {z_corr:+.2f}, p = {p_corr:.2f}); truths drawn exactly from the Gaussian '
f'model also pass (z = {z_ctrl:+.2f}, p = {p_ctrl:.2f}).', fontsize=11)
plt.savefig('./plt4.png', dpi=120)