diff --git a/bouquet/TokaMaker_interface.py b/bouquet/TokaMaker_interface.py index 4d620b3..7bf7daa 100644 --- a/bouquet/TokaMaker_interface.py +++ b/bouquet/TokaMaker_interface.py @@ -25,6 +25,7 @@ import matplotlib.pyplot as plt from .sampling import ( + GPRProfilePerturber, generate_perturbed_GPR, calc_cylindrical_li_proxy, get_li_proxy_geometry, @@ -626,6 +627,8 @@ def perturb_kinetic_equilibrium( spike_profile_recon_cached=None, proxy_bias_warmstart=None, pin_jphi=False, + verbose_interval=200, + **kwargs ): r"""Perturb kinetic and current-density profiles and iterate to match :math:`I_p` and :math:`l_i` targets. @@ -757,6 +760,16 @@ def _kin_to_eq(arr_kin): # ---------------------------------------------------------------- inp_avg = mygs.flux_integral(psi_N, pressure) + # Pre-compute GPR eigenfactors for the four kinetic profiles. + _ne_gpr = GPRProfilePerturber(kernel_func="rbf", length_scale=n_ls) + _ne_gpr.precompute_factor(psi_kin, sigma_ne / ne[0]) + _te_gpr = GPRProfilePerturber(kernel_func="rbf", length_scale=t_ls) + _te_gpr.precompute_factor(psi_kin, sigma_te / te[0]) + _ni_gpr = GPRProfilePerturber(kernel_func="rbf", length_scale=n_ls) + _ni_gpr.precompute_factor(psi_kin, sigma_ni / ni[0]) + _ti_gpr = GPRProfilePerturber(kernel_func="rbf", length_scale=t_ls) + _ti_gpr.precompute_factor(psi_kin, sigma_ti / ti[0]) + p_err = np.inf p_iter = 0 # p_thresh is a FRACTION (e.g. 0.05 == 5%); p_err is computed in percent. @@ -765,6 +778,8 @@ def _kin_to_eq(arr_kin): while p_err > _p_thresh_pct: p_iter += 1 + if (p_iter % verbose_interval == 0): + print(f" pressure match: iter={p_iter}, err={p_err:.3f}% (threshold {p_thresh}%)") if p_iter > max_pressure_iter: raise RuntimeError( f"Pressure match not found within {max_pressure_iter} iterations " @@ -773,19 +788,19 @@ def _kin_to_eq(arr_kin): # GPR sampling on psi_kin (kinetic grid, may include SOL) ne_perturb = _draw_monotonic_perturbation( - psi_kin, ne / ne[0], sigma_ne / ne[0], n_ls + psi_kin, ne / ne[0], sigma_ne / ne[0], n_ls, perturber=_ne_gpr ) * ne[0] te_perturb = _draw_monotonic_perturbation( - psi_kin, te / te[0], sigma_te / te[0], t_ls + psi_kin, te / te[0], sigma_te / te[0], t_ls, perturber=_te_gpr ) * te[0] ni_perturb = _draw_monotonic_perturbation( - psi_kin, ni / ni[0], sigma_ni / ni[0], n_ls + psi_kin, ni / ni[0], sigma_ni / ni[0], n_ls, perturber=_ni_gpr ) * ni[0] ti_perturb = _draw_monotonic_perturbation( - psi_kin, ti / ti[0], sigma_ti / ti[0], t_ls + psi_kin, ti / ti[0], sigma_ti / ti[0], t_ls, perturber=_ti_gpr ) * ti[0] # Pressure matching on equilibrium grid (psi_N, confined only) @@ -990,6 +1005,7 @@ def _probe(label): isolate_edge_jBS=isolate_edge_jBS, diagnostic_plots=False, verbose=False, + **kwargs ) finally: if _stashed_bounds is not None: @@ -1192,6 +1208,7 @@ def _probe(label): # SWB H-mode self-consistency iterations (default 3). Env # SWB_ITERS lets us trim for speed (2 is usually enough). iterations=int(os.environ.get('SWB_ITERS', '3')), + **kwargs ) if os.environ.get('PROFILE', '0') == '1': print(f" [profile] SWB call: {time.perf_counter()-_t_swb0:.1f}s") @@ -1854,6 +1871,7 @@ def generate_bouquet( baseline_pfile_bytes=None, psi_N_kinetic=None, max_proxy_draws=500, + verbose_interval=200, coil_drift=0.01, coil_drift_floor_A=50.0, vsc_coils=('F9A', 'F9B'), @@ -1871,6 +1889,7 @@ def generate_bouquet( jphi_baseline=True, seed=None, pin_jphi=False, + **kwargs ): r"""Generate a batch of perturbed equilibria and archive to HDF5. @@ -2756,7 +2775,7 @@ def _k2e(a): initial_Ip_target, _swb_seed_cache, scale_jBS=1.0, isolate_edge_jBS=isolate_edge_jBS, - diagnostic_plots=False, verbose=False, + diagnostic_plots=False, verbose=False,**kwargs ) _diff_spike_recon = np.asarray( _cache_results["isolated_j_BS"]).copy() @@ -2978,6 +2997,8 @@ def _k2e(a): spike_profile_recon_cached=_diff_spike_recon, proxy_bias_warmstart=_proxy_bias_warmstart, pin_jphi=pin_jphi, + verbose_interval=verbose_interval, + **kwargs ) except Exception as e: # Catch ANY exception during a perturbed solve -- ValueError @@ -3684,7 +3705,7 @@ def _build_bounds(drift_F, drift_VSC): def reconstruct_equilibrium(mygs, eqdsk, ne, te, ni, ti, Zeff, isoflux_pts, weights, psi_pad, guess_jinductive,n_k,psi_bridge,rescale_j_BS, - shelf_psi_N,initialize_psi=True): + shelf_psi_N,initialize_psi=True,**kwargs): r"""Reconstruct a single Grad-Shafranov equilibrium from a geqdsk reference and kinetic profiles, matching the EFIT :math:`l_i(1)`. @@ -3776,6 +3797,7 @@ def reconstruct_equilibrium(mygs, eqdsk, ne, te, ni, ti, Zeff, scale_jBS=1.0, isolate_edge_jBS=True, diagnostic_plots=False, + **kwargs ) j_BS_isolated = results['isolated_j_BS'] diff --git a/bouquet/sampling.py b/bouquet/sampling.py index f861013..ba834a8 100644 --- a/bouquet/sampling.py +++ b/bouquet/sampling.py @@ -151,6 +151,72 @@ def _matern52_kernel(self, X1: np.ndarray, X2: np.ndarray) -> np.ndarray: return prefactor * (1.0 + s + s**2 / 3.0) * np.exp(-s) # ---- core sampling method ---------------------------------------- + def precompute_factor( + self, + psi_N: np.ndarray, + sigma_profile: np.ndarray, + ) -> None: + r"""Pre-compute and cache GPR eigen-factor. + + After calling this, use :meth:`draw_from_factor` to generate + samples without repeating the O(n³) eigendecomposition -> np.linalg.eigh + + Parameters + ---------- + psi_N : ndarray + 1-D normalised flux grid. + sigma_profile : ndarray + 1-D experimental uncertainty **in profile units** -- this + becomes the GP's marginal standard deviation at every + grid point. + """ + # 1. Unit-variance base kernel + K = self._kernel(psi_N, psi_N) # K(x,x) = 1 + + # 2. Scale by σ(x): C_ij = σ_i · σ_j · K_ij + # ⟹ C(x,x) = σ(x)² ⟹ marginal std = σ(x) ✓ + S = np.outer(sigma_profile, sigma_profile) + K_scaled = K * S + + # 3. Eigen-decomposition (symmetric → eigh) + vals, vecs = np.linalg.eigh(K_scaled) + vals = np.maximum(vals, 0.0) + self._cached_vecs = vecs + self._cached_sqrt_vals = np.sqrt(vals) + + def draw_from_factor( + self, + input_profile: np.ndarray, + n_samples: int, + rng: np.random.Generator, + ) -> np.ndarray: + r"""Draw perturbed profiles whose pointwise :math:`1\sigma` + matches the uncertainty supplied to :meth:`precompute_factor` exactly. + + Call :meth:`precompute_factor` first. + + Parameters + ---------- + input_profile : ndarray + 1-D baseline profile (GP mean). + n_samples : int + Number of independent profile draws. + rng : numpy.random.Generator + Random generator. + + Returns + ------- + ndarray, shape ``(n_samples, len(input_profile))`` + """ + n = len(input_profile) + + # 4. Sample: δf = V diag(√λ) z, z ~ N(0, I) + z = rng.standard_normal((n, n_samples)) + perturbations = self._cached_vecs @ (self._cached_sqrt_vals[:, None] * z) + + # 5. Perturbed profiles → (n_samples, n_points) + return input_profile[None, :] + perturbations.T + def generate_profiles( self, psi_N: np.ndarray, @@ -372,6 +438,33 @@ def generate_perturbed_GPR( # ==================================================================== # Statistics verification # ==================================================================== +def _gpr_stats_from_samples(samples, profile, uncertainty_prof, confidence_band): + """Compute empirical statistics from a ``(n_samples, n_points)`` array. + + Helper shared by the two sampling paths in :func:`verify_gpr_statistics`. + + Returns + ------- + dict with keys ``empirical_mean``, ``empirical_std``, + ``exceedance_per_point``, ``avg_exceedance``. + """ + sigma_theory = uncertainty_prof + empirical_mean = np.mean(samples, axis=0) + empirical_std = np.std(samples, axis=0) + + residuals = samples - profile[None, :] + outside = np.abs(residuals) > confidence_band * sigma_theory[None, :] + exceedance_per_point = np.mean(outside, axis=0) + avg_exceedance = np.mean(exceedance_per_point) + + return { + "empirical_mean": empirical_mean, + "empirical_std": empirical_std, + "exceedance_per_point": exceedance_per_point, + "avg_exceedance": avg_exceedance, + } + + def verify_gpr_statistics( psi_N, profile, @@ -388,6 +481,8 @@ def verify_gpr_statistics( 2. Pointwise std :math:`\approx` ``uncertainty_prof`` (variance check) 3. Fraction of samples outside :math:`\pm k\sigma` :math:`\approx` theoretical (tail check) + Calls both the draw_from_factor and generate_profiles methods for a peace-of-mind comparison. + Parameters ---------- psi_N : ndarray @@ -414,6 +509,12 @@ def verify_gpr_statistics( length_scale=length_scale, ) + # ---- Path A: re-draw (precompute_factor + draw_from_factor) ----- + perturber.precompute_factor(psi_N, uncertainty_prof) + rng_a = np.random.default_rng(42) + samples_a = perturber.draw_from_factor(profile, n_verification, rng_a) + + # ---- Path B: generate_profiles ------------- rng = np.random.default_rng(42) samples = perturber.generate_profiles( psi_N, profile, uncertainty_prof, @@ -424,17 +525,25 @@ def verify_gpr_statistics( # Marginal std equals sigma_profile exactly by construction sigma_theory = uncertainty_prof - # ---- empirical statistics --------------------------------------- - empirical_mean = np.mean(samples, axis=0) - empirical_std = np.std(samples, axis=0) - - # ---- pointwise exceedance rate ---------------------------------- - residuals = samples - profile[None, :] # (n_verification, n_points) - outside = np.abs(residuals) > confidence_band * sigma_theory[None, :] - # Fraction of samples outside the band at each point - exceedance_per_point = np.mean(outside, axis=0) - # Overall average exceedance rate - avg_exceedance = np.mean(exceedance_per_point) + stats_a = _gpr_stats_from_samples(samples_a, profile, sigma_theory, confidence_band) + samples = samples_a + empirical_mean = stats_a["empirical_mean"] + empirical_std = stats_a["empirical_std"] + exceedance_per_point = stats_a["exceedance_per_point"] + avg_exceedance = stats_a['avg_exceedance'] + + stats_b = _gpr_stats_from_samples(samples_b, profile, sigma_theory, confidence_band) + + # ---- cross-check: both paths must agree on empirical std -------- + denom = np.maximum(sigma_theory, 1e-30) + std_rel = np.abs(stats_a["empirical_std"] - stats_b["empirical_std"]) / denom + max_rel = np.max(std_rel) + if max_rel > 0.02: + raise RuntimeError( + f"Re-draw and generate_profiles empirical stds differ by " + f"{max_rel:.3f} (> 2 %). The two paths are not statistically " + f"equivalent — check GPRProfilePerturber implementation." + ) # ---- theoretical exceedance for a Gaussian ---------------------- theoretical_exceedance = 2.0 * norm.sf(confidence_band) # two-tailed @@ -456,6 +565,8 @@ def verify_gpr_statistics( # (a) Empirical vs theoretical std axes[0, 0].plot(psi_N, sigma_theory, "k-", lw=2, label="Theory") axes[0, 0].plot(psi_N, empirical_std, "r--", lw=1.5, label="Empirical") + axes[0, 0].plot(psi_N, stats_b["empirical_std"], "b:", lw=1.5, + label="Empirical b") axes[0, 0].set_ylabel(r"$\sigma(x)$") axes[0, 0].set_title("Pointwise std: theory vs empirical") axes[0, 0].legend() @@ -466,6 +577,8 @@ def verify_gpr_statistics( label=f"Theory ({theoretical_exceedance:.3f})") axes[0, 1].plot(psi_N, exceedance_per_point, "r-", alpha=0.7, label="Empirical") + axes[0, 1].plot(psi_N, stats_b["exceedance_per_point"], "b:", alpha=0.7, + label="Empirical b") axes[0, 1].set_ylabel(f"Fraction outside ±{confidence_band:.0f}σ") axes[0, 1].set_title("Exceedance rate per grid point") axes[0, 1].legend() @@ -512,11 +625,13 @@ def verify_gpr_statistics( plt.show() return { - "empirical_mean": empirical_mean, - "empirical_std": empirical_std, + "stats_a": stats_a, + "stats_b": stats_b, + "empirical_mean": stats_a['empirical_mean'], + "empirical_std": stats_a['empirical_std'], "sigma_theory": sigma_theory, - "exceedance_per_point": exceedance_per_point, - "avg_exceedance": avg_exceedance, + "exceedance_per_point": stats_a['exceedance_per_point'], + "avg_exceedance": stats_a['avg_exceedance'], "theoretical_exceedance": theoretical_exceedance, } @@ -720,6 +835,8 @@ def _draw_monotonic_perturbation( sigma_profile, length_scale, max_draws=_MAX_MONOTONIC_DRAWS, + perturber=None, + rng=None, ): r"""Repeatedly sample a GPR perturbation until the draw is monotonically decreasing. @@ -737,6 +854,12 @@ def _draw_monotonic_perturbation( GPR correlation length (scalar or spatially-varying). max_draws : int Safety cap on the number of attempts. + perturber : GPRProfilePerturber or None + Optional pre-initialised perturber with + :meth:`~GPRProfilePerturber.precompute_factor` already called. + rng : numpy.random.Generator or None + Optional shared random generator. A fresh generator is created + if ``None``. Returns ------- @@ -748,15 +871,14 @@ def _draw_monotonic_perturbation( RuntimeError If no monotonic draw is found within *max_draws* attempts. """ + if perturber is None: + perturber = GPRProfilePerturber(kernel_func="rbf", length_scale=length_scale) + perturber.precompute_factor(psi_N, sigma_profile) + if rng is None: + rng = np.random.default_rng() + for _ in range(max_draws): - sample = generate_perturbed_GPR( - psi_N, - normalised_profile, - sigma_profile=sigma_profile, - length_scale=length_scale, - n_samples=1, - diag_plot=False, - ) + sample = perturber.draw_from_factor(normalised_profile, 1, rng)[0] if np.all(np.diff(sample) <= 0.0): return sample diff --git a/tests/test_core.py b/tests/test_core.py index 1f16a16..ceb7a59 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -16,6 +16,7 @@ from bouquet.sampling import ( GPRProfilePerturber, generate_perturbed_GPR, + _draw_monotonic_perturbation, ) from bouquet.utils import ( initialize_equilibrium_database, @@ -157,6 +158,153 @@ def test_matern52_runs(self): result = p.generate_profiles(psi_N, profile, sigma, n_samples=3) assert result.shape == (3, len(psi_N)) + def test_output_shape_single_redraw(self): + psi_N = np.linspace(0, 1, 51) + profile = 1.0 - psi_N + sigma = 0.05 * np.ones_like(psi_N) + p = GPRProfilePerturber(kernel_func="rbf", length_scale=0.2) + p.precompute_factor(psi_N, sigma) + rng = np.random.default_rng(0) + result = p.draw_from_factor(profile, 1, rng) + assert result.shape == (1, len(psi_N)) + + def test_output_shape_multi_redraw(self): + psi_N = np.linspace(0, 1, 51) + profile = 1.0 - psi_N + sigma = 0.05 * np.ones_like(psi_N) + p = GPRProfilePerturber(kernel_func="rbf", length_scale=0.2) + p.precompute_factor(psi_N, sigma) + rng = np.random.default_rng(0) + result = p.draw_from_factor(profile, 10, rng) + assert result.shape == (10, len(psi_N)) + + def test_zero_sigma_returns_mean_redraw(self): + psi_N = np.linspace(0, 1, 51) + profile = 1.0 - psi_N + sigma = np.zeros_like(psi_N) + p = GPRProfilePerturber(kernel_func="rbf", length_scale=0.2) + p.precompute_factor(psi_N, sigma) + rng = np.random.default_rng(0) + result = p.draw_from_factor(profile, 5, rng) + # Every row should match the input profile + for i in range(result.shape[0]): + np.testing.assert_allclose(result[i], profile, atol=1e-10) + + def test_marginal_std_matches_sigma_redraw(self): + """Empirical pointwise std should match the input sigma.""" + rng = np.random.default_rng(42) + psi_N = np.linspace(0, 1, 41) + profile = np.ones_like(psi_N) + sigma = 0.1 * np.ones_like(psi_N) + p = GPRProfilePerturber(kernel_func="rbf", length_scale=0.2) + p.precompute_factor(psi_N, sigma) + samples = p.draw_from_factor(profile, 5000, rng) + empirical_std = np.std(samples, axis=0) + np.testing.assert_allclose(empirical_std, sigma, rtol=0.1) + + def test_redraw_matches_generate_profiles_redraw(self): + """Re-draw and generate_profiles must produce the same marginal std.""" + psi_N = np.linspace(0, 1, 41) + profile = np.ones_like(psi_N) + sigma = 0.1 * np.ones_like(psi_N) + p = GPRProfilePerturber(kernel_func="rbf", length_scale=0.2) + p.precompute_factor(psi_N, sigma) + + rng_a = np.random.default_rng(7) + samples_a = p.draw_from_factor(profile, 3000, rng_a) + + rng_b = np.random.default_rng(7) + samples_b = p.generate_profiles(psi_N, profile, sigma, n_samples=3000, rng=rng_b) + + std_a = np.std(samples_a, axis=0) + std_b = np.std(samples_b, axis=0) + # Empirical stds agree to within 1 % (Monte-Carlo noise) + np.testing.assert_allclose(std_a, std_b, rtol=0.01) + + def test_matern52_runs_redraw(self): + psi_N = np.linspace(0, 1, 51) + profile = 1.0 - psi_N + sigma = 0.05 * np.ones_like(psi_N) + p = GPRProfilePerturber(kernel_func="matern52", length_scale=0.2) + p.precompute_factor(psi_N, sigma) + rng = np.random.default_rng(0) + result = p.draw_from_factor(profile, 3, rng) + assert result.shape == (3, len(psi_N)) + + def test_precompute_reuse_consistent(self): + """Multiple draw_from_factor calls with the same factor must be i.i.d.""" + psi_N = np.linspace(0, 1, 31) + profile = np.ones_like(psi_N) + sigma = 0.1 * np.ones_like(psi_N) + p = GPRProfilePerturber(kernel_func="rbf", length_scale=0.2) + p.precompute_factor(psi_N, sigma) + rng = np.random.default_rng(99) + # Draw in two separate calls and stack + batch1 = p.draw_from_factor(profile, 3000, rng) + batch2 = p.draw_from_factor(profile, 3000, rng) + # Verify batch1 and batch2 are actually different (independent) + assert not np.allclose(batch1, batch2), "batch1 and batch2 should be different" + combined = np.vstack([batch1, batch2]) + empirical_std = np.std(combined, axis=0) + # Check thath their standard deviation is similar + np.testing.assert_allclose(empirical_std, sigma, rtol=0.01) + + +class TestDrawMonotonicPerturbation: + """Quick tests for _draw_monotonic_perturbation.""" + + # A strictly decreasing parabola — nearly every GPR draw is monotone + PSI = np.linspace(0, 1, 32) + PROFILE = 1.0 - PSI ** 2 + SIGMA = 0.02 * np.ones(32) + + def test_returns_monotone_array(self): + result = _draw_monotonic_perturbation( + self.PSI, self.PROFILE, self.SIGMA, length_scale=0.3, + ) + assert result.shape == (len(self.PSI),) + assert np.all(np.diff(result) <= 0.0) + + def test_pre_built_perturber_accepted(self): + """Passing a pre-built perturber must still return a valid monotone draw.""" + p = GPRProfilePerturber(kernel_func="rbf", length_scale=0.3) + p.precompute_factor(self.PSI, self.SIGMA) + rng = np.random.default_rng(7) + result = _draw_monotonic_perturbation( + self.PSI, self.PROFILE, self.SIGMA, length_scale=0.3, + perturber=p, rng=rng, + ) + assert np.all(np.diff(result) <= 0.0) + + def test_shared_rng_advances_state(self): + """Two calls with the same rng object produce different draws.""" + p = GPRProfilePerturber(kernel_func="rbf", length_scale=0.3) + p.precompute_factor(self.PSI, self.SIGMA) + rng = np.random.default_rng(42) + r1 = _draw_monotonic_perturbation( + self.PSI, self.PROFILE, self.SIGMA, length_scale=0.3, + perturber=p, rng=rng, + ) + r2 = _draw_monotonic_perturbation( + self.PSI, self.PROFILE, self.SIGMA, length_scale=0.3, + perturber=p, rng=rng, + ) + assert not np.allclose(r1, r2) + + def test_exhaustion_raises_runtime_error(self): + """RuntimeError must be raised when no monotone draw is found in max_draws.""" + # Flat profile + huge sigma: draws will not be monotone in 1 attempt. + p = GPRProfilePerturber(kernel_func="rbf", length_scale=0.3) + psi = np.linspace(0, 1, 8) + sigma = 10.0 * np.ones(8) + p.precompute_factor(psi, sigma) + rng = np.random.default_rng(0) + with pytest.raises(RuntimeError, match="monotonically"): + _draw_monotonic_perturbation( + psi, np.ones(8), sigma, length_scale=0.3, + max_draws=1, batch_size=1, perturber=p, rng=rng, + ) + class TestGeneratePerturbedGPR: """Tests for the convenience wrapper.""" diff --git a/tests/test_gpr_timing.py b/tests/test_gpr_timing.py new file mode 100644 index 0000000..17154ee --- /dev/null +++ b/tests/test_gpr_timing.py @@ -0,0 +1,218 @@ +""" +Timing benchmark: re-draw path vs repeated generate_perturbed_GPR calls. + +Run with:: + + pytest tests/test_gpr_timing.py -v -s + +or directly:: + + python -m pytest tests/test_gpr_timing.py -v -s --tb=short + +The benchmark is not a correctness test, it always passes. Its +purpose is to show the wall-clock speedup from amortising the O(n³) +eigendecomposition with ``precompute_factor`` + ``draw_from_factor`` +versus calling ``generate_perturbed_GPR`` (which does a fresh linalg.eigh on +every call). +""" + +import time +from unittest.mock import patch + +import matplotlib +matplotlib.use("Agg") # non-interactive backend — must be set before any other plt import +import numpy as np +import pytest + +from bouquet.sampling import GPRProfilePerturber, generate_perturbed_GPR, verify_gpr_statistics + +# ==================================================================== +# Shared fixtures +# ==================================================================== + +N_GRID = 257 # profile grid points +N_DRAWS = 200 # number of samples to generate in each benchmark +LENGTH = 0.25 # GPR correlation length +SIGMA = 0.05 # flat fractional uncertainty + +N_VERIFY = 2000 # samples for verify_gpr_statistics (reduced from default 5000 + # to keep the test fast; increase for publication-quality checks) + + +@pytest.fixture(scope="module") +def grid(): + psi_N = np.linspace(0, 1, N_GRID) + profile = 1.0 - psi_N + sigma = SIGMA * np.ones_like(psi_N) + return psi_N, profile, sigma + + +# ==================================================================== +# Benchmarks +# ==================================================================== + +class TestGPRTiming: + """Wall-clock comparison between the two sampling strategies.""" + + def test_timing_comparison(self, grid, capsys): + """Print a timing table; always passes.""" + psi_N, profile, sigma = grid + + # ---- Method A: re-draw (precompute_factor + draw_from_factor) ---- + perturber = GPRProfilePerturber(kernel_func="rbf", length_scale=LENGTH) + + rng_a = np.random.default_rng(0) + t0 = time.perf_counter() + perturber.precompute_factor(psi_N, sigma) # O(n³) — included in total + for _ in range(N_DRAWS): + perturber.draw_from_factor(profile, 1, rng_a) + t_redraw = time.perf_counter() - t0 + + # ---- Method B: repeated generate_perturbed_GPR (fresh eigh each call) ---- + rng_b = np.random.default_rng(0) + t0 = time.perf_counter() + for _ in range(N_DRAWS): + generate_perturbed_GPR( + psi_N, profile, + sigma_profile=sigma, + length_scale=LENGTH, + n_samples=1, + rng=rng_b, + ) + t_repeated = time.perf_counter() - t0 + + speedup = t_repeated / t_redraw if t_redraw > 0 else float("inf") + + with capsys.disabled(): + print(f"\n{'='*55}") + print(f" GPR timing ({N_DRAWS} draws, {N_GRID}-point grid)") + print(f"{'='*55}") + print(f" Re-draw (precompute_factor + {N_DRAWS}× draw_from_factor):") + print(f" total : {t_redraw*1e3:.1f} ms") + print(f" per : {t_redraw/N_DRAWS*1e3:.3f} ms/draw (amortised)") + print(f" Repeated generate_perturbed_GPR (fresh eigh each):") + print(f" total : {t_repeated*1e3:.1f} ms") + print(f" per : {t_repeated/N_DRAWS*1e3:.3f} ms/draw") + print(f" Speedup : {speedup:.1f}×") + print(f"{'='*55}") + + # Sanity: re-draw should be at least 2× faster for N_DRAWS >= 10 + if N_DRAWS >= 10: + assert speedup >= 2.0, ( + f"Expected re-draw to be ≥2× faster; got {speedup:.2f}×. " + "This may indicate the eigendecomposition is not being cached." + ) + + def test_timing_batch_vs_loop(self, grid, capsys): + """Batch draw (n_samples > 1) vs loop of single draws for pre_computed factor (minor speedup)""" + psi_N, profile, sigma = grid + + perturber = GPRProfilePerturber(kernel_func="rbf", length_scale=LENGTH) + + rng_a = np.random.default_rng(1) + t0 = time.perf_counter() + perturber.precompute_factor(psi_N, sigma) # O(n³) — included in total + _ = perturber.draw_from_factor(profile, N_DRAWS, rng_a) + t_batch = time.perf_counter() - t0 + + perturber2 = GPRProfilePerturber(kernel_func="rbf", length_scale=LENGTH) + rng_b = np.random.default_rng(1) + t0 = time.perf_counter() + perturber2.precompute_factor(psi_N, sigma) # O(n³) — included in total + for _ in range(N_DRAWS): + perturber2.draw_from_factor(profile, 1, rng_b) + t_loop = time.perf_counter() - t0 + + with capsys.disabled(): + print(f"\n{'='*55}") + print(f" Batch vs loop ({N_DRAWS} draws, {N_GRID}-point grid)") + print(f"{'='*55}") + print(f" Batch draw_from_factor(n={N_DRAWS}) + precompute: {t_batch*1e3:.2f} ms") + print(f" Loop draw_from_factor(n=1)×{N_DRAWS} + precompute: {t_loop*1e3:.2f} ms") + print(f" Ratio batch/loop: {t_batch/t_loop:.2f}") + print(f"{'='*55}") + + def test_timing_grid_scaling(self, capsys): + """Show how timing scales with grid size for precomputed factor vs generate_perturbed_GPR.""" + grid_sizes = [32, 64, 128, 256] + n_draws = 1000 + + rows = [] + for n in grid_sizes: + psi_N = np.linspace(0, 1, n) + profile = 1.0 - psi_N + sigma = SIGMA * np.ones(n) + + # Re-draw + p = GPRProfilePerturber(kernel_func="rbf", length_scale=LENGTH) + rng = np.random.default_rng(0) + t0 = time.perf_counter() + p.precompute_factor(psi_N, sigma) # O(n³) — included in total + for _ in range(n_draws): + p.draw_from_factor(profile, 1, rng) + t_rd = (time.perf_counter() - t0) / n_draws * 1e3 # ms/draw + + # Repeated + rng2 = np.random.default_rng(0) + t0 = time.perf_counter() + for _ in range(n_draws): + generate_perturbed_GPR(psi_N, profile, + sigma_profile=sigma, + length_scale=LENGTH, + n_samples=1, rng=rng2) + t_rep = (time.perf_counter() - t0) / n_draws * 1e3 + + rows.append((n, t_rd, t_rep, t_rep / t_rd if t_rd > 0 else float("inf"))) + + with capsys.disabled(): + print(f"\n{'='*55}") + print(f" Grid-size scaling ({n_draws} draws each)") + print(f" {'n':>5} {'redraw ms':>10} {'repeated ms':>12} {'speedup':>8}") + print(f" {'-'*5} {'-'*10} {'-'*12} {'-'*8}") + for n, rd, rep, sp in rows: + print(f" {n:>5} {rd:>10.3f} {rep:>12.3f} {sp:>8.1f}×") + print(f"{'='*55}") + + +# ==================================================================== +# Statistical verification (both sampling paths) +# ==================================================================== + +class TestVerifyGPRStatistics: + """Run verify_gpr_statistics to cross-check re-draw vs generate_profiles. + + ``plt.show`` is patched to a no-op so this runs non-interactively. + The test passes only if the two paths agree to within 2 %. + """ + + def test_verify_prints_and_agrees(self, capsys): + psi_N = np.linspace(0, 1, N_GRID) + profile = 1.0 - psi_N + sigma = SIGMA * np.ones_like(psi_N) + + with patch("matplotlib.pyplot.show"): + result = verify_gpr_statistics( + psi_N, profile, sigma, + length_scale=LENGTH, + n_verification=N_VERIFY, + confidence_band=2.0, + ) + + # Check return structure + assert "stats_a" in result + assert "stats_b" in result + assert "sigma_theory" in result + assert "theoretical_exceedance" in result + + # Both paths must have avg_exceedance close to the Gaussian prediction + from scipy.stats import norm + theory = 2.0 * norm.sf(2.0) + for key in ("stats_a", "stats_b"): + exc = result[key]["avg_exceedance"] + assert abs(exc - theory)/theory < 0.02, ( + f"{key} avg_exceedance {exc:.4f} deviates from theory " + f"{theory:.4f} by more than 0.02" + ) + + with capsys.disabled(): + print() # verify_gpr_statistics already printed its own report