Batched gpr sampling - updated

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']

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