diff --git a/bouquet/TokaMaker_interface.py b/bouquet/TokaMaker_interface.py index 4d620b3..e9a04d5 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,9 @@ def perturb_kinetic_equilibrium( spike_profile_recon_cached=None, proxy_bias_warmstart=None, pin_jphi=False, + verbose_interval=200, + worker_id=None, + **kwargs ): r"""Perturb kinetic and current-density profiles and iterate to match :math:`I_p` and :math:`l_i` targets. @@ -710,6 +714,15 @@ def perturb_kinetic_equilibrium( pressure matching and equilibrium solving. Returned perturbed profiles are on ``psi_N_kinetic``. If ``None``, ``psi_N`` is used for everything (original behaviour). + max_proxy_draws : int + Maximum number of proxy-space draws attempted per :math:`l_i` + iteration before raising ``RuntimeError`` (default 500). + verbose_interval : int + Print pressure-matching progress every this many iterations + (default 200). + worker_id : int or None + Worker identifier prepended to log messages when running inside + a multiprocessing pool. ``None`` (default) disables the prefix. Returns ------- @@ -743,6 +756,7 @@ def perturb_kinetic_equilibrium( # Kinetic grid: either the user-supplied extended grid or psi_N psi_kin = psi_N_kinetic if psi_N_kinetic is not None else psi_N _dual_grid = psi_N_kinetic is not None + _pfx = f"[Worker {worker_id}] " if worker_id is not None else "" def _kin_to_eq(arr_kin): """Interpolate a profile from kinetic grid onto equilibrium grid.""" @@ -757,14 +771,26 @@ 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. _p_thresh_pct = float(p_thresh) * 100.0 - print("Searching for pressure profile match...") + print(f"{_pfx}Searching for pressure profile match...") while p_err > _p_thresh_pct: p_iter += 1 + if (p_iter % verbose_interval == 0): + print(f"{_pfx} 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 +799,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 +1016,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 +1219,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") @@ -1602,7 +1630,7 @@ def _probe(label): f"widen l_i_tolerance/PRESCREEN_MARGIN or check sigma_jphi") dt_proxy = time.perf_counter() - t_phase - print(f" [li_iter={li_iter}] GPR draw " + print(f"{_pfx} [li_iter={li_iter}] GPR draw " f"({_gpr_try} tries, {_n_skipped} pre-screen-skipped, " f"{dt_proxy:.1f}s)") @@ -1636,8 +1664,8 @@ def _probe(label): _, q_pre, _, _, _, _ = mygs.get_q(npsi=npsi, psi_pad=psi_pad) if q_pre[0] < 1.0: dt_scale = time.perf_counter() - t_scale - print(f" [li_iter={li_iter}] find_optimal_scale: {dt_scale:.1f}s") - print("Skipping this equilibrium, q_0 < 1.0 (pre-check)") + print(f"{_pfx} [li_iter={li_iter}] find_optimal_scale: {dt_scale:.1f}s") + print(f"{_pfx}Skipping this equilibrium, q_0 < 1.0 (pre-check)") l_i = np.inf continue @@ -1645,13 +1673,13 @@ def _probe(label): # solver holds Ip to target natively, so Ip_target is used unscaled # downstream (no Ip-scale secant). dt_scale = time.perf_counter() - t_scale - print(f" [li_iter={li_iter}] find_optimal_scale (j0 only): {dt_scale:.1f}s") + print(f"{_pfx} [li_iter={li_iter}] find_optimal_scale (j0 only): {dt_scale:.1f}s") # ---- 5d. Definitive sawtooth constraint (after Ip scaling) -- if constrain_sawteeth: _, q, _, _, _, _ = mygs.get_q(npsi=npsi, psi_pad=psi_pad) if q[0] < 1.0: - print("Skipping this equilibrium, q_0 < 1.0") + print(f"{_pfx}Skipping this equilibrium, q_0 < 1.0") l_i = np.inf continue @@ -1688,7 +1716,7 @@ def _probe(label): rtol=0.05, verbose=False, ) if _n_corr > 2: - print(f" [jphi correction] {_n_corr} iterations, " + print(f"{_pfx} [jphi correction] {_n_corr} iterations, " f"edge RMS: {_corr_hist[0]/1e6:.4f} → {_corr_hist[-1]/1e6:.4f} MA/m²") if diagnostic_plots: @@ -1741,14 +1769,14 @@ def _probe(label): if l_i > 0 and np.isfinite(l_i): proxy_target = final_li_proxy * (l_i_target / l_i) - print(f" l_i target (equil): {l_i_target:.4f}") - print(f" proxy target: {proxy_target:.4f} (corrected)") - print(f" matched l_i (equil): {l_i:.4f}") - print(f" matched l_i (proxy): {final_li_proxy:.4f}") - print(f" Ip error vs target: {Ip_err:.3f}%") - print(f" proxy vs real l_i: {proxy_vs_real:+.2f}%") + print(f"{_pfx} l_i target (equil): {l_i_target:.4f}") + print(f"{_pfx} proxy target: {proxy_target:.4f} (corrected)") + print(f"{_pfx} matched l_i (equil): {l_i:.4f}") + print(f"{_pfx} matched l_i (proxy): {final_li_proxy:.4f}") + print(f"{_pfx} Ip error vs target: {Ip_err:.3f}%") + print(f"{_pfx} proxy vs real l_i: {proxy_vs_real:+.2f}%") _li_pct_err = 100.0 * abs(l_i - l_i_target) / l_i_target if l_i_target != 0 else float('inf') - print(f" l_i error: {_li_pct_err:.2f}% (tolerance: {_li_tol_pct:.2f}%)") + print(f"{_pfx} l_i error: {_li_pct_err:.2f}% (tolerance: {_li_tol_pct:.2f}%)") iteration_l_is.append(l_i) iteration_Ips.append(Ip) @@ -1854,6 +1882,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 +1900,9 @@ def generate_bouquet( jphi_baseline=True, seed=None, pin_jphi=False, + keep_geqdsk=False, + worker_id=None, + **kwargs ): r"""Generate a batch of perturbed equilibria and archive to HDF5. @@ -2016,6 +2048,29 @@ def generate_bouquet( Soft-reg weight for the ``#VSC`` channel (default 1.0). Kept much lower than ``soft_reg_weight`` so the VSC has freedom to do vertical-mode control work without being heavily penalized. + keep_geqdsk : bool + If ``True``, the temporary per-equilibrium ``.geqdsk`` files written + by ``mygs.save_eqdsk`` are kept on disk after being archived into the + HDF5 database. Useful for manual inspection or debugging. + Default is ``False`` (files are deleted after archiving). + psi_N_kinetic : ndarray or None + Optional extended kinetic-profile grid (starting at 0, ending at + :math:`\hat{\psi} \geq 1`). When provided, ``ne``, ``te``, + ``ni``, ``ti`` and their sigmas must be on this grid; + profiles are interpolated onto ``psi_N`` before the GS solve. + Returned perturbed profiles are on ``psi_N_kinetic``. + ``None`` uses ``psi_N`` for everything. + max_proxy_draws : int + Maximum proxy draws per :math:`l_i` iteration before + ``RuntimeError`` (default 500). Forwarded to + :func:`perturb_kinetic_equilibrium`. + verbose_interval : int + Print pressure-matching progress every this many iterations + (default 200). Forwarded to + :func:`perturb_kinetic_equilibrium`. + worker_id : int or None + Worker identifier prepended to log messages. ``None`` (default) + disables the prefix. Returns ------- @@ -2028,6 +2083,7 @@ def generate_bouquet( np.random.seed(int(seed)) all_diagnostics = [] + _pfx = f"[Worker {worker_id}] " if worker_id is not None else "" # self-consistent pressure for baseline

# When kinetic profiles are on a different grid, interpolate @@ -2756,7 +2812,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() @@ -2837,14 +2893,14 @@ def _k2e(a): remaining = avg_s * (n_equils - count) eta_min = remaining / 60.0 eta_str = f" ETA: {eta_min:.1f} min" - print(f"\n{'='*60}") - print(f" Equilibrium {count+1}/{n_equils} " + print(f"\n{_pfx}{'='*60}") + print(f"{_pfx} Equilibrium {count+1}/{n_equils} " f"(scale_jBS={scale_jBS:.4f}){eta_str}") if l_i_uncertainty > 0.0: _dev_pct = 100.0 * (l_i_target_draw - l_i_target) / l_i_target - print(f" l_i_target sampled: {l_i_target_draw:.4f} " - f"({_dev_pct:+.2f}% vs recon, σ={100*l_i_uncertainty:.1f}%)") - print(f"{'='*60}") + print(f"{_pfx} l_i_target sampled: {l_i_target_draw:.4f} " + f"{_pfx}({_dev_pct:+.2f}% vs recon, σ={100*l_i_uncertainty:.1f}%)") + print(f"{_pfx}{'='*60}") t_start = time.perf_counter() # ---- Warm-start restore ---- @@ -2978,6 +3034,9 @@ def _k2e(a): spike_profile_recon_cached=_diff_spike_recon, proxy_bias_warmstart=_proxy_bias_warmstart, pin_jphi=pin_jphi, + verbose_interval=verbose_interval, + worker_id=worker_id, + **kwargs ) except Exception as e: # Catch ANY exception during a perturbed solve -- ValueError @@ -3435,7 +3494,7 @@ def _build_bounds(drift_F, drift_VSC): elapsed = time.perf_counter() - t_start elapsed_times.append(elapsed) total_elapsed = time.perf_counter() - t_batch_start - print(f" Wall-clock time: {elapsed:.1f}s " + print(f"{_pfx} Wall-clock time: {elapsed:.1f}s " f"(total: {total_elapsed/60:.1f} min, " f"avg: {np.mean(elapsed_times):.1f}s/eq)") @@ -3463,7 +3522,7 @@ def _build_bounds(drift_F, drift_VSC): # ---- save geqdsk to a temporary file, archive, delete ------- eqdsk_filename = f"{header}_count={count}.geqdsk" full_path = os.path.abspath(eqdsk_filename) - print(f" Saving to: {full_path}") + print(f"{_pfx} Saving to: {full_path}") # safe_save_eqdsk: snapshot mygs equilibrium before save, restore # after. Prevents save_eqdsk's q-profile tracer from shifting @@ -3614,7 +3673,7 @@ def _build_bounds(drift_F, drift_VSC): perturbed_pfile_bytes = pf.to_bytes() except Exception as exc: import traceback - print(f" WARNING: could not build perturbed p-file: {exc}") + print(f"{_pfx} WARNING: could not build perturbed p-file: {exc}") traceback.print_exc() store_equilibrium( @@ -3650,11 +3709,14 @@ def _build_bounds(drift_F, drift_VSC): ) # Clean up on-disk eqdsk after archiving - try: - os.remove(full_path) - print(f" Deleted temporary file: {full_path}") - except OSError as exc: - print(f" WARNING: could not delete {full_path}: {exc}") + if keep_geqdsk: + print(f"{_pfx} Keeping temporary file: {full_path}") + else: + try: + os.remove(full_path) + print(f"{_pfx} Deleted temporary file: {full_path}") + except OSError as exc: + print(f"{_pfx} WARNING: could not delete {full_path}: {exc}") all_diagnostics.append(diagnostics) @@ -3681,10 +3743,11 @@ def _build_bounds(drift_F, drift_VSC): # ==================================================================== # Single-equilibrium reconstruction from geqdsk + kinetic profiles # ==================================================================== -def reconstruct_equilibrium(mygs, eqdsk, ne, te, ni, ti, Zeff, +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): + guess_jinductive, n_k, psi_bridge, rescale_j_BS, + shelf_psi_N, initialize_psi=True, + psi_N_kinetic=None, **kwargs): r"""Reconstruct a single Grad-Shafranov equilibrium from a geqdsk reference and kinetic profiles, matching the EFIT :math:`l_i(1)`. @@ -3717,8 +3780,13 @@ def reconstruct_equilibrium(mygs, eqdsk, ne, te, ni, ti, Zeff, Ion density on ``eqdsk.psi_N`` [m\ :sup:`-3`]. ti : ndarray Ion temperature on ``eqdsk.psi_N`` [eV]. - Zeff : ndarray - Effective charge on ``eqdsk.psi_N``. + Zeff : dict or ndarray + Effective ion charge profile. Either: + * a dictionary ``{'x': psi_grid, 'y': values}`` giving the + profile on an arbitrary normalised psi grid, or + * a scalar float / 1-D array on ``eqdsk.psi_N`` (length + ``len(eqdsk.psi_N)``) or on ``psi_N_kinetic`` (length + ``len(psi_N_kinetic)`` when psi_N_kinetic is provided). isoflux_pts : ndarray, shape (N, 2) :math:`(R, Z)` coordinates of isoflux constraint points [m]. Passed to ``mygs.set_isoflux``. @@ -3747,6 +3815,13 @@ def reconstruct_equilibrium(mygs, eqdsk, ne, te, ni, ti, Zeff, If ``True`` (default), call ``mygs.init_psi`` using LCFS geometry estimated from the geqdsk boundary. Set to ``False`` to skip initialisation (e.g. when reusing a prior solution). + psi_N_kinetic : ndarray or None + Optional kinetic-profile grid (starting at 0, ending at + :math:`\hat{\psi} \geq 1`). When provided, ``ne``, ``te``, + ``ni`` and ``ti`` are expected on this + grid and are interpolated onto ``eqdsk.psi_N`` before the GS + solve. Mirrors the same parameter in :func:`generate_bouquet`. + ``guess_jinductive`` is always on ``eqdsk.psi_N``. Returns ------- @@ -3757,6 +3832,76 @@ def reconstruct_equilibrium(mygs, eqdsk, ne, te, ni, ti, Zeff, from OpenFUSIONToolkit.TokaMaker.util import create_power_flux_fun from OpenFUSIONToolkit.TokaMaker.bootstrap import solve_with_bootstrap + # --- Grid sanity checks --- + _psi = eqdsk.psi_N + _dpsi = np.diff(_psi) + assert (np.isclose(_psi[0], 0.0) and np.isclose(_psi[-1], 1.0)), f"eqdsk.psi_N must run from 0 to 1; got [{_psi[0]:.6g}, {_psi[-1]:.6g}]" + assert np.allclose(_dpsi, _dpsi[0]), "eqdsk.psi_N not uniformly sampled" + + _eq_len = len(_psi) + _dual_grid = psi_N_kinetic is not None + _kin_len = len(psi_N_kinetic) if _dual_grid else _eq_len + + # psi_N_kinetic bounds and same-length ambiguity (mirrors generate_bouquet) + if _dual_grid: + if not (np.isclose(psi_N_kinetic[0], 0.0) and psi_N_kinetic[-1] >= 1.0): + raise ValueError( + "psi_N_kinetic must start at 0 and end at psi_N >= 1; " + f"got [{psi_N_kinetic[0]:.6g}, {psi_N_kinetic[-1]:.6g}]" + ) + if _kin_len == _eq_len: + if np.allclose(psi_N_kinetic, _psi): + warn( + "psi_N_kinetic has the same length and endpoints as eqdsk.psi_N; " + "providing a separate kinetic grid of identical length is redundant. " + "This usage is deprecated.", + DeprecationWarning, stacklevel=2, + ) + elif not isinstance(Zeff, dict) and np.ndim(Zeff) > 0: + raise ValueError( + "psi_N_kinetic and eqdsk.psi_N have the same length but differ: " + "it is ambiguous which grid array-valued Zeff belongs to. " + "Use dict-format Zeff to specify the psi grid explicitly." + ) + + if len(guess_jinductive) != _eq_len: + raise ValueError( + f"guess_jinductive has length {len(guess_jinductive)} " + f"but eqdsk.psi_N has length {_eq_len}" + ) + for _name, _arr in {'ne': ne, 'te': te, 'ni': ni, 'ti': ti}.items(): + if len(_arr) != _kin_len: + _grid_name = 'psi_N_kinetic' if _dual_grid else 'eqdsk.psi_N' + raise ValueError( + f"{_name} has length {len(_arr)} but expected " + f"{_kin_len} ({_grid_name})" + ) + if not isinstance(Zeff, dict): + Zeff = np.asarray(Zeff) + if Zeff.ndim > 0 and len(Zeff) not in (_kin_len, _eq_len): + raise ValueError( + f"Zeff has length {len(Zeff)} but expected either " + f"eqdsk.psi_N ({_eq_len})" + + (f" or psi_N_kinetic ({_kin_len})" if _dual_grid else "") + ) + + # Interpolate kinetic profiles from psi_N_kinetic onto the equilibrium + # grid eqdsk.psi_N when a separate kinetic grid is supplied. + # Mirrors the _kin_to_eq logic in generate_bouquet. + if _dual_grid: + from scipy.interpolate import interp1d as _interp1d_kin + def _kin_to_eq(_arr): + return _interp1d_kin( + psi_N_kinetic, _arr, kind='linear', + bounds_error=False, fill_value=(_arr[0], _arr[-1]) + )(_psi) + ne = _kin_to_eq(ne) + te = _kin_to_eq(te) + ni = _kin_to_eq(ni) + ti = _kin_to_eq(ti) + if not isinstance(Zeff, dict) and Zeff.ndim > 0 and len(Zeff) == _kin_len: + Zeff = _kin_to_eq(Zeff) + if initialize_psi: # Estimate shape parameters from geqdsk LCFS geometry geo = eqdsk.geometry @@ -3765,6 +3910,9 @@ def reconstruct_equilibrium(mygs, eqdsk, ne, te, ni, ti, Zeff, a = geo['a'][-1] kappa = geo['kappa'][-1] delta = geo['delta'][-1] + ffp_prof = create_power_flux_fun(40,1.5,2.0) + pp_prof = create_power_flux_fun(40,4.0,1.0) + mygs.set_profiles(ffp_prof=ffp_prof,pp_prof=pp_prof,foffset=kwargs.get('F0', None)) # Need to reset flux profiles to prevent old jphi-linterp or jphi-split-bootstrap ffp_profs throwing errors mygs.init_psi(R0, Z0, a, kappa, delta) eqdsk_jtor = abs(eqdsk.j_tor_averaged_direct) @@ -3776,6 +3924,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'] @@ -3876,6 +4025,7 @@ def reconstruct_equilibrium(mygs, eqdsk, ne, te, ni, ti, Zeff, "x": eqdsk.psi_N, } + mygs.set_targets(Ip=abs(eqdsk.Ip), pax=pres_tmp[0]) mygs.set_profiles(ffp_prof=ffp_prof, pp_prof=pp_prof) mygs.solve() diff --git a/bouquet/io/ida.py b/bouquet/io/ida.py new file mode 100644 index 0000000..cfdab8c --- /dev/null +++ b/bouquet/io/ida.py @@ -0,0 +1,396 @@ +""" +IDA-lite NetCDF profile reader +============================== + +Reads kinetic profiles from IDA-lite ``.cdf`` files produced by the +Integrated Data Analysis (IDA) code at IPP Garching. + +Two file layouts are supported, detected automatically from the shape of the +``n_e`` array: + +- **Workflow 1** — shape ``(n_times, n_samples, n_radial)``: a single time + point with a full posterior distribution of samples. The central profile + is the median (``uncertainty_method='percentile'``) or mean + (``uncertainty_method='std'``) over the sample axis. +- **Workflow 2** — shape ``(n_times, n_radial)``: multiple time slices with + pre-computed fitted profiles and explicit uncertainty columns + (``n_e_err``, ``T_e_err``, ``T_12C6_err``). + +Requires the ``h5py`` package. +""" + +import numpy as np +from scipy.interpolate import interp1d + +# --------------------------------------------------------------------------- +# Module-level helpers (also used by parallel.IDALiteUncertaintyGenerator) +# --------------------------------------------------------------------------- + +def _interp_to_grid(psin_src, arr, psin_tgt): + """Linear interpolation onto *psin_tgt* with boundary fill values.""" + return interp1d( + psin_src, arr, + kind='linear', bounds_error=False, + fill_value=(arr[0], arr[-1]), + )(psin_tgt) + + +def _summarise_samples(samples, method='percentile'): + """Reduce ``(n_samples, n_radial)`` posterior draws to centre and 1σ. + + Parameters + ---------- + samples : ndarray, shape (n_samples, n_radial) + method : ``'percentile'`` (default) or ``'std'`` + + Returns + ------- + centre : ndarray (n_radial,) + Posterior median (percentile) or mean (std). + sigma : ndarray (n_radial,) + Half-width of the 16th–84th percentile band, or 1-sigma std. + """ + if method == 'std': + centre = samples.mean(axis=0) + sigma = samples.std(axis=0) + else: # percentile + centre = np.median(samples, axis=0) + lo = np.percentile(samples, 16, axis=0) + hi = np.percentile(samples, 84, axis=0) + sigma = (hi - lo) / 2.0 + return centre, sigma + + +def _detect_workflow(pf_ne, profile_file): + """Return ``1`` or ``2`` based on the dimensionality of *pf_ne*.""" + if pf_ne.ndim == 3: + return 1 # (n_times, n_samples, n_radial) — posterior samples + if pf_ne.ndim == 2: + return 2 # (n_times, n_radial) — fitted profiles + error columns + raise ValueError( + f"Unexpected n_e shape {pf_ne.shape} in {profile_file!r}. " + "Expected 2D (n_times, n_radial) or 3D (n_times, n_samples, n_radial)." + ) + + +def _select_time_index(time_arr_ms, workflow_num, time_idx, sim_time_ms): + """Return the integer time index to use for a given CDF file.""" + if workflow_num == 1 or sim_time_ms is None: + return time_idx + # Workflow 2 with explicit target time + pf_time = np.asarray(time_arr_ms) / 1e3 # ms → s + return int(np.argmin(np.abs(pf_time - sim_time_ms / 1e3))) + + +# --------------------------------------------------------------------------- +# Public reader class +# --------------------------------------------------------------------------- + +class IDALiteProfileReader: + r"""Read kinetic profiles from an IDA-lite NetCDF (.cdf) file. + + Two file layouts are detected automatically from the shape of ``n_e``: + + - **Workflow 1** — ``(n_times, n_samples, n_radial)``: Bayesian posterior + samples at a single time point. The central profile is computed as + the median (``uncertainty_method='percentile'``) or mean + (``uncertainty_method='std'``) across the sample axis. + - **Workflow 2** — ``(n_times, n_radial)``: fitted profiles for multiple + time slices with explicit uncertainty columns. + + ``Zeff`` is read directly from the CDF variable ``Zeff``. + + By default, the main-ion (deuterium) density is computed using the + measured carbon density via quasi-neutrality: + + .. math:: + + n_i = \max\bigl(n_e - Z_C\,n_C,\; 0\bigr) + + where :math:`n_C` is read from CDF variable ``n_12C6``. Setting + *carbon_quasi_neutrality* to ``False`` falls back to the simpler + approximation :math:`n_i \approx n_e`. + + Parameters + ---------- + time_idx : int + Time index to use for workflow 1 (typically 0 for a single-time file) + or as a fallback for workflow 2 when *sim_time_ms* is not provided. + sim_time_ms : float or None + Target time in milliseconds. When provided and the file is workflow 2, + the nearest time slice in the CDF ``time`` variable is selected. + Ignored for workflow 1. + Z_C : int + Charge number of the carbon impurity (default 6). Only used when + *carbon_quasi_neutrality* is ``True``. + uncertainty_method : ``'percentile'`` | ``'std'`` + How to summarise workflow-1 posterior samples. ``'percentile'`` uses + the median and 16th/84th percentile band; ``'std'`` uses mean ± σ. + carbon_quasi_neutrality : bool + If ``True`` (default), compute :math:`n_i = \max(n_e - Z_C n_C, 0)` using the + ``n_12C6`` CDF variable. If ``False``, use :math:`n_i \approx n_e`. + """ + + def __init__(self, time_idx=0, sim_time_ms=None, Z_C=6, + uncertainty_method='percentile', + carbon_quasi_neutrality=True): + self.time_idx = time_idx + self.sim_time_ms = sim_time_ms + self.Z_C = Z_C + self.uncertainty_method = uncertainty_method + self.carbon_quasi_neutrality = carbon_quasi_neutrality + + def __call__(self, geqdsk_file, profile_file, psi_N): + """Read profiles from an IDA-lite CDF and interpolate onto *psi_N*. + + Parameters + ---------- + geqdsk_file : str + Path to the geqdsk (unused; present for interface consistency). + profile_file : str + Path to the IDA-lite ``.cdf`` file. + psi_N : array-like + Normalised flux grid onto which profiles are interpolated. + + Returns + ------- + ne_SI, te_SI, ni_SI, ti_SI : ndarray + Profiles in SI units (m^-3, eV). + Zeff_eq : ndarray + Effective charge (clipped to ≥ 1). + profile_bytes : bytes + Raw bytes of the profile file for HDF5 archival. + """ + import h5py + + with h5py.File(profile_file, 'r') as f: + # Read raw data, converting to working units (10^20 m^-3, keV) + pf_ne = f['n_e'][:] / 1e20 + pf_te = f['T_e'][:] / 1e3 + pf_ti = f['T_12C6'][:] / 1e3 + pf_zeff = f['Zeff'][:] + time_arr = f['time'][:] + pf_psin = f['psi_n'][:] + if self.carbon_quasi_neutrality: + pf_nc = f['n_12C6'][:] / 1e20 + + workflow_num = _detect_workflow(pf_ne, profile_file) + t_idx = _select_time_index(time_arr, workflow_num, self.time_idx, + self.sim_time_ms) + + if workflow_num == 1: + # Radial grid is the same for all samples + psin_ida = pf_psin[t_idx, 0, :] + ne_centre, _ = _summarise_samples(pf_ne[t_idx], self.uncertainty_method) + te_centre, _ = _summarise_samples(pf_te[t_idx], self.uncertainty_method) + ti_centre, _ = _summarise_samples(pf_ti[t_idx], self.uncertainty_method) + zeff_centre = np.clip(pf_zeff[t_idx].mean(axis=0), 1.0, None) + if self.carbon_quasi_neutrality: + nc_centre, _ = _summarise_samples(pf_nc[t_idx], self.uncertainty_method) + else: # workflow 2 + psin_ida = pf_psin + ne_centre = pf_ne[t_idx].copy() + te_centre = pf_te[t_idx].copy() + ti_centre = pf_ti[t_idx].copy() + zeff_centre = np.clip(pf_zeff[t_idx], 1.0, None) + if self.carbon_quasi_neutrality: + nc_centre = pf_nc[t_idx].copy() + + # Derive ni: either ne - Z_C*n_C (carbon QN) or ni ≈ ne (simple) + if self.carbon_quasi_neutrality: + ni_centre = np.maximum(ne_centre - self.Z_C * nc_centre, 0.0) + else: + ni_centre = ne_centre + + # Interpolate onto target grid + ne_SI = _interp_to_grid(psin_ida, ne_centre, psi_N) * 1e20 + te_SI = _interp_to_grid(psin_ida, te_centre, psi_N) * 1e3 + ni_SI = _interp_to_grid(psin_ida, ni_centre, psi_N) * 1e20 + ti_SI = _interp_to_grid(psin_ida, ti_centre, psi_N) * 1e3 + Zeff_eq = np.clip(_interp_to_grid(psin_ida, zeff_centre, psi_N), 1.0, None) + + with open(profile_file, 'rb') as fh: + profile_bytes = fh.read() + + return ne_SI, te_SI, ni_SI, ti_SI, Zeff_eq, profile_bytes + +# --------------------------------------------------------------------------- +# Uncertainty generator +# --------------------------------------------------------------------------- + +class IDALiteUncertaintyGenerator: + r"""Compute 1σ kinetic-profile uncertainty envelopes from an IDA-lite CDF. + + Two file layouts are handled automatically (same detection as + :class:`IDALiteProfileReader`): + + - **Workflow 1** — ``(n_times, n_samples, n_radial)``: Bayesian posterior + samples at a single time point. The 1σ half-width is derived from the + posterior distribution: + + - ``uncertainty_method='percentile'`` (default) → half-width of the + 16th–84th percentile band. + - ``uncertainty_method='std'`` → standard deviation across samples. + + - **Workflow 2** — ``(n_times, n_radial)``: fitted profiles with explicit + per-variable uncertainty columns in the CDF. The sigma arrays are read + directly from the configurable variable names (*ne_sigma_var*, etc.). + + For ``sigma_jphi``, pass ``None`` — IDA-lite files do not carry current + density information. The caller is responsible for computing it from the + reconstructed :math:`j_\phi` (see ``jphi_uncertainty_gen`` in + :class:`~bouquet.parallel.load_IDA_file_obj`). + + Parameters + ---------- + time_idx : int + Time index for workflow 1 (typically 0) or fallback for workflow 2 + when *sim_time_ms* is not provided. + sim_time_ms : float or None + Target time in milliseconds for workflow 2 time-slice selection. + The nearest slice in the CDF ``time`` variable is chosen. + Ignored for workflow 1. + Z_C : int + Charge of the carbon impurity (default 6). Used for error propagation + when *carbon_quasi_neutrality* is ``True``. + uncertainty_method : ``'percentile'`` | ``'std'`` + How to summarise workflow 1 posterior samples. + carbon_quasi_neutrality : bool + If ``True``, propagate the carbon-density uncertainty into + :math:`\sigma_{n_i}` via + + .. math:: + + \sigma_{n_i} = \sqrt{\sigma_{n_e}^2 + (Z_C\,\sigma_{n_C})^2} + + - **Workflow 1**: :math:`\sigma_{n_C}` is derived from the ``n_12C6`` + posterior samples using *uncertainty_method*. + - **Workflow 2**: :math:`\sigma_{n_C}` is read from *nc_sigma_var*. + + When ``False`` the simpler :math:`\sigma_{n_i} \approx + \sigma_{n_e}` approximation is used (or *ni_sigma_var* when set). + ne_sigma_var : str + CDF variable name for the electron-density 1σ [m\ :sup:`-3`] + (workflow 2 only). + te_sigma_var : str + CDF variable name for the electron-temperature 1σ [eV] + (workflow 2 only). + ti_sigma_var : str + CDF variable name for the ion-temperature 1σ [eV] + (workflow 2 only). + ni_sigma_var : str or None + CDF variable name for an explicit main-ion density 1σ + [m\ :sup:`-3`] (workflow 2 only). When set this takes priority over + both *carbon_quasi_neutrality* propagation and the *ne_sigma_var* + fallback. + nc_sigma_var : str + CDF variable name for the carbon-density 1σ [m\ :sup:`-3`] + (workflow 2 only; used only when *carbon_quasi_neutrality* is + ``True`` and *ni_sigma_var* is ``None``). + """ + + def __init__( + self, + time_idx=0, + sim_time_ms=None, + Z_C=6, + uncertainty_method='percentile', + carbon_quasi_neutrality=True, + ne_sigma_var='n_e_err', + te_sigma_var='T_e_err', + ti_sigma_var='T_12C6_err', + ni_sigma_var=None, + nc_sigma_var='n_12C6_err', + ): + self.time_idx = time_idx + self.sim_time_ms = sim_time_ms + self.Z_C = Z_C + self.uncertainty_method = uncertainty_method + self.carbon_quasi_neutrality = carbon_quasi_neutrality + self.ne_sigma_var = ne_sigma_var + self.te_sigma_var = te_sigma_var + self.ti_sigma_var = ti_sigma_var + self.ni_sigma_var = ni_sigma_var + self.nc_sigma_var = nc_sigma_var + + def __call__(self, profile_file, profile_reader, psi_N, j_phi_fit=None): + """Compute 1σ arrays from the CDF and interpolate onto *psi_N*. + + Parameters + ---------- + profile_file : str + Path to the IDA-lite ``.cdf`` file. + profile_reader : callable + Unused; present for interface consistency with the bouquet + uncertainty-generator protocol. + psi_N : ndarray + Normalised flux grid onto which sigmas are interpolated. + j_phi_fit : ndarray or None + Unused; retained for interface consistency. + + Returns + ------- + sigma_ne, sigma_te, sigma_ni, sigma_ti : ndarray + Absolute 1σ uncertainties in SI units (m\ :sup:`-3`, eV). + sigma_jphi : None + Always ``None``; current-density uncertainty is not available + from IDA-lite files and must be supplied by the caller. + """ + import h5py + + with h5py.File(profile_file, 'r') as f: + pf_ne = f['n_e'][:] / 1e20 # → 10^20 m^-3 + time_arr = f['time'][:] + + workflow_num = _detect_workflow(pf_ne, profile_file) + t_idx = _select_time_index(time_arr, workflow_num, self.time_idx, + self.sim_time_ms) + + if workflow_num == 1: + # ---- Posterior samples: derive sigma from the distribution ---- + # Radial grid is the same for all samples (shape: n_radial) + psin_ida = f['psi_n'][t_idx, 0, :] + + _, sigma_ne_sc = _summarise_samples( + pf_ne[t_idx], self.uncertainty_method) # 10^20 m^-3 + _, sigma_te_sc = _summarise_samples( + f['T_e'][t_idx] / 1e3, self.uncertainty_method) # keV + _, sigma_ti_sc = _summarise_samples( + f['T_12C6'][t_idx] / 1e3, self.uncertainty_method) # keV + + if self.carbon_quasi_neutrality: + # ni = ne - Z_C * n_C → σ(ni) = sqrt(σ(ne)^2 + (Z_C σ(n_C))^2) + _, sigma_nc_sc = _summarise_samples( + f['n_12C6'][t_idx] / 1e20, self.uncertainty_method) + sigma_ni_sc = np.sqrt(sigma_ne_sc**2 + (self.Z_C * sigma_nc_sc)**2) + else: + sigma_ni_sc = sigma_ne_sc.copy() # σ(ni) ≈ σ(ne) + + # Convert scaled units → SI + sigma_ne_si = sigma_ne_sc * 1e20 # m^-3 + sigma_te_si = sigma_te_sc * 1e3 # eV + sigma_ni_si = sigma_ni_sc * 1e20 # m^-3 + sigma_ti_si = sigma_ti_sc * 1e3 # eV + + else: + # ---- Workflow 2: read explicit error columns (already SI) ---- + psin_ida = f['psi_n'][:] + sigma_ne_si = f[self.ne_sigma_var][t_idx].copy() + sigma_te_si = f[self.te_sigma_var][t_idx].copy() + sigma_ti_si = f[self.ti_sigma_var][t_idx].copy() + if self.ni_sigma_var is not None: + # Explicit override: use the named column directly + sigma_ni_si = f[self.ni_sigma_var][t_idx].copy() + elif self.carbon_quasi_neutrality: + # ni = ne - Z_C * n_C → σ(ni) = sqrt(σ(ne)^2 + (Z_C σ(n_C))^2) + sigma_nc_si = f[self.nc_sigma_var][t_idx].copy() + sigma_ni_si = np.sqrt(sigma_ne_si**2 + (self.Z_C * sigma_nc_si)**2) + else: + sigma_ni_si = sigma_ne_si.copy() # σ(ni) ≈ σ(ne) fallback + + sigma_ne = _interp_to_grid(psin_ida, sigma_ne_si, psi_N) + sigma_te = _interp_to_grid(psin_ida, sigma_te_si, psi_N) + sigma_ni = _interp_to_grid(psin_ida, sigma_ni_si, psi_N) + sigma_ti = _interp_to_grid(psin_ida, sigma_ti_si, psi_N) + + return sigma_ne, sigma_te, sigma_ni, sigma_ti, None diff --git a/bouquet/parallel.py b/bouquet/parallel.py new file mode 100644 index 0000000..0225c41 --- /dev/null +++ b/bouquet/parallel.py @@ -0,0 +1,1019 @@ +"""Method-agnostic parallel bouquet runner +======================== + +Distributes ``(input_files, load_files_obj, bouquet_method)`` +across available CPU cores. Each case (ie. each timeslice, kinetic equilibrium, shot) +has an associated tuple of input file names, which are read into python using the specified +load method, and then passed to the bouquet method. parallel_runner distributes +these cases across available CPU cores and runs them in parallel. + +'Non atomic' input files with multiple timeslices/kinetic equilibria (eg. IDA files) +are supported by optional atomic_input_recast and atomic_load_files methods inside load_files_obj. + +Basic pfile example: + input_files = (eqdsk, pfile) + load_files_obj.load_files = load_eqdsk_pfile + bouquet_method = re_generate_bouquet + +@authors Stuart Benjamin +@date June 2026 +""" + +########################################################################################################### +# General parallel functions +########################################################################################################### + +import os +import sys +import queue +import shutil +import socket +import traceback +import pickle as pkl +import multiprocessing +import numpy as np +from threadpoolctl import threadpool_limits + +# Module-level state populated by _init_worker in each spawned worker process. +_worker_state: dict = {} + +class _IndexMap: + """Picklable map_object: ``map_object(idx)`` returns ``flat_list[idx]``. + + ``map_object`` that can be pickled & saved to disk by ``parallel_runner``. + """ + def __init__(self, flat_list): + self.flat_list = flat_list + + def __call__(self, idx): + return self.flat_list[idx] + + def __len__(self): + return len(self.flat_list) + + def __iter__(self): + return iter(self.flat_list) + +class FractionalUncertainty: + """Picklable callable that returns ``frac * |x|``. + + Use as ``config['jphi_uncertainty_gen']`` when the j_phi uncertainty + is a fixed fraction of the fitted current profile:: + + config["jphi_uncertainty_gen"] = FractionalUncertainty(0.10) # 10 % + + """ + def __init__(self, frac): + self.frac = frac + + def __call__(self, x): + return self.frac * np.abs(x) + +def parallel_runner(all_input_files, load_files_obj, bouquet_method, master_working_dir, + chunksize='automatic', use_logical_cpus=True, n_cpus_override=None, + verbose=False, keep_output=False): + """Run a bouquet method in parallel across available CPU cores (single node). + all_input_files must be a list of tuples, where each tuple contains the input files for a single 'case', + matching the expected input of load_files_obj.load_files. + + Parameters + ---------- + verbose : bool + ``False`` (default): each worker's output is redirected to a per-worker + log file (``/worker_N.log``); the terminal only + shows brief per-worker status lines from the parent process. + ``True``: no redirection — all worker output streams directly to the + terminal (asynchronously, used for debugging). + """ + + #=================================================================================== + # Chunking logic + #=================================================================================== + + if n_cpus_override is not None: + n_cpus, nthreads = n_cpus_override, 1 + else: + n_cpus, nthreads = _get_num_cpus(use_logical=use_logical_cpus) + n_runs, map_object = load_files_obj.total_runs(all_input_files) + if n_runs == 0: + print("[bouquet_parallel] No runs to execute.") + return {}, {} + n_workers = min(n_cpus, n_runs) + print( + f"[bouquet_parallel] Distributing {n_runs} runs across " + f"{n_workers} workers ({n_cpus} CPUs available, {nthreads} thread(s)/worker)." + ) + + if not load_files_obj.is_atomic: + _all_input_files = load_files_obj.atomic_input_recast(all_input_files) + load_files = load_files_obj.atomic_load_files + else: + _all_input_files = all_input_files + load_files = load_files_obj.load_files + assert len(_all_input_files) == n_runs, ( + f"Expected {n_runs} runs from load_files_obj.total_runs, but got " + f"{len(_all_input_files)} from load_files_obj.atomic_input_recast" + ) + + if chunksize == 'automatic': + # Heuristic: 10x more tasks than workers, but no more than 1000 tasks per chunk + chunksize = max(1, min(1000, n_runs // (10 * n_workers))) + print(f"[bouquet_parallel] Using chunksize={chunksize} for dynamic scheduling.") + else: + print(f"[bouquet_parallel] Using user-specified chunksize={chunksize} for dynamic scheduling.") + + # Save map_object so users can look up input files by idx after the run. + map_object_path = os.path.join(master_working_dir, "map_object.pkl") + with open(map_object_path, "wb") as f: + pkl.dump(map_object, f) + print(f"[bouquet_parallel] Saved input file map to {map_object_path}") + + #=================================================================================== + # Pool setup + #=================================================================================== + + os.makedirs(master_working_dir, exist_ok=True) + + # 'spawn' avoids fork-safety issues with Fortran shared libraries in OFT + ctx = multiprocessing.get_context("spawn") + errors = {} + outputs = {} + + # Each worker reports (worker_id, None) on success or + # (worker_id, traceback_str) on failure via this queue. We wait for + # all n_workers to report before dispatching any tasks so that a + # broken initializer causes an immediate, clean failure instead of a + # silent hang in imap_unordered. + + init_status_queue = ctx.Queue() + + # Hand each spawned worker a unique ID via a pre-loaded queue. + worker_id_queue = ctx.Queue() + for w in range(n_workers): + worker_id_queue.put(w) + + # Inject nthreads into a config copy so _init_OFT sets thread counts correctly. + _config = dict(load_files_obj.config) + _config["_nthreads"] = nthreads + _config["_verbose"] = verbose + + _pool = ctx.Pool( + processes=n_workers, + initializer=load_files_obj.init_worker, + initargs=(worker_id_queue, master_working_dir, _config, init_status_queue), + ) + + # Barrier: wait for every worker to finish initialising, terminate if there's a failure + init_failures = [] + for _ in range(n_workers): + try: + wid, tb = init_status_queue.get(timeout=120) # 2 min per worker + except queue.Empty: + init_failures.append((-1, "Worker initialisation timed out (> 120 s)")) + else: + if tb is not None: + init_failures.append((wid, tb)) + else: + if verbose: + print(f"[bouquet_parallel] Worker {wid} ready.", flush=True) + else: + log = os.path.join(master_working_dir, f"worker_{wid}.log") + print(f"[bouquet_parallel] Worker {wid} ready (log: {log})", flush=True) + + if init_failures: + _pool.terminate() + _pool.join() + msgs = "\n".join( + f" Worker {wid}:\n{tb}" for wid, tb in init_failures + ) + raise RuntimeError( + f"[bouquet_parallel] FATAL: {len(init_failures)} worker(s) failed " + f"to initialise:\n{msgs}" + ) + + + #=================================================================================== + # Task dispatch + #=================================================================================== + + per_run_args = [(i, _all_input_files[i], load_files, bouquet_method) for i in range(n_runs)] + + try: + with _pool: + for idx, success, err_msg, output in _pool.imap_unordered(_run_one, per_run_args, + chunksize=chunksize): + if not success: + errors[idx] = err_msg + print( + f"[bouquet_parallel] WARNING: run {idx} " + f"({_all_input_files[idx]}) failed:\n{err_msg}" + ) + else: + if not keep_output: + output = None + outputs[idx] = output + except KeyboardInterrupt: + _pool.join() + raise + except Exception as _exc: + _pool.join() + raise RuntimeError( + f"[bouquet_parallel] FATAL error during task dispatch:\\n" + f"{traceback.format_exc()}" + ) from _exc + + n_success = n_runs - len(errors) + print(f"[bouquet_parallel] Completed: {n_success}/{n_runs} runs succeeded.") + + if errors: + error_path = os.path.join(master_working_dir, "errors.pkl") + with open(error_path, "wb") as f: + pkl.dump(errors, f) + print(f"[bouquet_parallel] Error details saved to {error_path}") + + return errors, outputs + +def _get_num_cpus(use_logical=True): + """Return ``(n_workers, nthreads_per_worker)`` for spawning OFT workers. + + Puportedly works on Linux HPC cluster (SLURM, PBS, LSF, SGE) and degrades + gracefully on non-Linux systems (macOS, Windows). + + Parameters + ---------- + use_logical : bool + ``True`` (default): one worker per logical CPU (hyperthread), + ``nthreads=1``. + + ``False``: one worker per physical core, ``nthreads = logical/physical``. + Uses OFT's OpenMP intra-core parallelism. + + Returns + ------- + n_workers : int + nthreads_per_worker : int + """ + # --- Logical CPU count from OS affinity (Linux) or cpu_count (other) --- + try: + affinity = os.sched_getaffinity(0) # Linux: respects cgroup/taskset + n_logical = len(affinity) + except AttributeError: + affinity = None + n_logical = os.cpu_count() or 1 # macOS / Windows fallback + + # --- Physical core count via Linux sysfs --- + n_physical = None + if affinity is not None: + core_ids = set() + for cpu in affinity: + try: + with open(f"/sys/devices/system/cpu/cpu{cpu}/topology/physical_package_id") as _f: + pkg = _f.read().strip() + with open(f"/sys/devices/system/cpu/cpu{cpu}/topology/core_id") as _f: + core = _f.read().strip() + core_ids.add((pkg, core)) + except OSError: + pass + if core_ids: + n_physical = len(core_ids) + if n_physical is None: + n_physical = n_logical # sysfs unavailable: assume no SMT + nthreads_per_core = max(1, n_logical // n_physical) + + if use_logical: + # Scheduler-specific CPU count env vars (used as a cap to avoid + # over-subscription when the affinity set is wider than the job's + # CPU reservation — observed on some SLURM configurations). + _SCHEDULER_CPU_VARS = ( + "SLURM_CPUS_PER_TASK", # SLURM + "PBS_NUM_PPN", # PBS (CPUs per node) + "LSB_DJOB_NUMPROC", # IBM LSF + "NSLOTS", # SGE / Grid Engine + ) + for var in _SCHEDULER_CPU_VARS: + val = os.environ.get(var) + if val is not None: + n_logical = min(n_logical, int(val)) + break + return n_logical, 1 + else: + return n_physical, nthreads_per_core + +def _run_one(run_args): + """Worker function: run one case of the bouquet method.""" + idx, input_files, load_files, bouquet_method = run_args + + nthreads = _worker_state.get("config", {}).get("_nthreads", 1) + with threadpool_limits(limits=nthreads): + try: + data = load_files(input_files, idx) + output = bouquet_method(data) + return idx, True, None, output + except Exception as exc: + tb_str = traceback.format_exc() + return idx, False, tb_str, None + +########################################################################################################### +# bouquet_method 're_generate_bouquet' +########################################################################################################### + +_RE_GENERATE_BOUQUET_REQUIRED_KEYS = ( + "idx", "eqdsk", "profile_bytes", + "psi_N", # equilibrium (GS) grid from eqdsk + "psi_N_kinetic", # kinetic profile grid (None → psi_N assumed) + # profiles defined on psi_N_kinetic (or psi_N should psi_N_kinetic be None) + "ne_SI", "te_SI", "ni_SI", "ti_SI", + "sigma_ne", "sigma_te", "sigma_ni", "sigma_ti", + "Zeff", # either scalar, or dictionary of psi_normalised 'x' values and Zeff 'y' values + "w_ExB", # STUB (currently unused) on psi_N equilibrium grid + "sigma_jphi", # on psi_N (equilibrium grid) since j_phi is an equilibrium quantity +) + +def _check_data_keys(data, required_keys, tag=""): + """Raise ValueError listing all missing keys if *data* is incomplete.""" + missing = [k for k in required_keys if k not in data] + if missing: + raise ValueError( + f"{tag} data dict is missing required keys: {missing}" + ) + +def re_generate_bouquet(data): + # Unpack shared worker state functions. We do this because we want load errors caught by _init_worker. + reconstruct_equilibrium = _worker_state["reconstruct_equilibrium"] + generate_bouquet = _worker_state["generate_bouquet"] + store_equilibrium = _worker_state["store_equilibrium"] + create_power_flux_fun = _worker_state["create_power_flux_fun"] + initialize_equilibrium_database = _worker_state["initialize_equilibrium_database"] + # Unpack shared worker state objects + config = _worker_state["config"] + worker_id = _worker_state["worker_id"] + mygs = _worker_state["mygs"] + + # Reset mygs for new equilibrium, keeping mesh and coils + mygs.set_targets() + + # --- Sanity check input data dict --- + idx = data['idx'] + _tag = f"[Worker {worker_id} | run {idx}]" + _check_data_keys(data, _RE_GENERATE_BOUQUET_REQUIRED_KEYS, _tag) + + # Per-run header: one HDF5 database per equilibrium, named by idx. + header = f"{config['header']}_idx{idx}" + initialize_equilibrium_database(header) + + # --- Reconstruct equilibrium --- + print(f"{_tag} Reconstructing equilibrium...", flush=True) + isoflux_pts = np.column_stack([data['eqdsk'].boundary_R, data['eqdsk'].boundary_Z]) + isoflux_weights = np.ones(len(data['eqdsk'].boundary_R)) * config["isoflux_weight"] + mygs.set_isoflux(isoflux_pts, weights=isoflux_weights) + guess_jinductive = create_power_flux_fun(len(data['psi_N']), 1.5, 1.5)["y"] + result = reconstruct_equilibrium( + mygs, + data['eqdsk'], + data['ne_SI'], + data['te_SI'], + data['ni_SI'], + data['ti_SI'], + data['Zeff'], + isoflux_pts, + isoflux_weights, + config["psi_pad"], + guess_jinductive=guess_jinductive, + rescale_j_BS=False, + shelf_psi_N=0.0, + initialize_psi=True, + psi_N_kinetic=data['psi_N_kinetic'], + F0=abs(data['eqdsk'].R_center * data['eqdsk'].B_center), + **config.get("reconstruct_equilibrium_kwargs", {}), + ) + + # --- Save the reconstructed geqdsk, read raw bytes, store profiles in HDF5, then clean up --- + eqdsk_out = f"{header}.geqdsk" + eqdsk_out_abs = os.path.abspath(eqdsk_out) + mygs.save_eqdsk(eqdsk_out, nr=257, nz=257, truncate_eq=False, lcfs_pad=config["psi_pad"]) + with open(eqdsk_out_abs, 'rb') as _fh: + baseline_eqdsk_raw = _fh.read() + li1 = mygs.get_stats(lcfs_pad=config["psi_pad"], li_normalization="std")["l_i"] + li3 = mygs.get_stats(lcfs_pad=config["psi_pad"], li_normalization="iter")["l_i"] + store_equilibrium( + header, + 0, + eqdsk_out_abs, + data['psi_N'], + result["j_phi_fit"], + result["j_BS_used"], + result["j_inductive_fit"], + data['ne_SI'], + data['te_SI'], + data['ni_SI'], + data['ti_SI'], + data['w_ExB'], + li1, + li3, + Zeff=data["Zeff"], + psi_N_kinetic=data['psi_N_kinetic'], + ) + if config.get("keep_geqdsk", False): + print(f"{_tag} Keeping reconstruction geqdsk: {eqdsk_out_abs}", flush=True) + else: + os.remove(eqdsk_out) + print(f"{_tag} Reconstruction done — li_final={result['li_final']:.4f}, li1={li1:.4f}", flush=True) + + # --- j_phi uncertainty --- + if data['sigma_jphi'] is None: + data['sigma_jphi'] = config['jphi_uncertainty_gen'](result["j_phi_fit"]) + + # --- Generate perturbed equilibrium family --- + print(f"{_tag} Generating {config['n_equils']} perturbed equilibria...", flush=True) + mygs.set_isoflux(result["isoflux_pts"], weights=result["weights"]) + diagnostics = generate_bouquet( + mygs, + data['psi_N'], + config["n_equils"], + header, + result["j_phi_fit"], + data['ne_SI'], + data['te_SI'], + data['ni_SI'], + data['ti_SI'], + data['sigma_ne'], + data['sigma_te'], + data['sigma_ni'], + data['sigma_ti'], + data['sigma_jphi'], + config["n_ls"], + config["t_ls"], + config["j_ls"], + abs(data['eqdsk'].Ip), + result["li_final"], + data["Zeff"], + input_jinductive=result["j_inductive_fit"], + psi_N_kinetic=data['psi_N_kinetic'], + pfile_bytes=data['profile_bytes'], + baseline_eqdsk_bytes=baseline_eqdsk_raw, + baseline_pfile_bytes=data['profile_bytes'], + diagnostic_plots=False, + **config.get("generate_bouquet_kwargs", {}), + ) + + # --- Final reporting --- + print(f"{_tag} Done — {len(diagnostics)} equilibria archived.", flush=True) + return { + "li_final": result["li_final"], + "li1": li1, + "li3": li3, + "n_equils_generated": len(diagnostics), + } + +########################################################################################################### +# Generic load files object (load_files_obj) for 're_generate_bouquet' +########################################################################################################### + +class load_files_obj: + """Base interface for load_files objects used by parallel_runner. + + Subclasses must set ``is_atomic`` and implement ``load_files`` (or + ``atomic_load_files`` + ``atomic_input_recast`` if not atomic), + ``total_runs``, ``init_worker``, and ``config``. + + Atomic = one equilibrium per input file tuple, so no recasting needed. + Non-atomic = multiple equilibria per input file tuple, so recasting needed. + """ + is_atomic: bool + config: dict + + def total_runs(self, all_input_files): + """Return ``(n_runs, map_object)`` where ``map_object(idx)`` gives the + atomic input-files tuple for run *idx*.""" + raise NotImplementedError + + def load_files(self, input_files, idx): + """Load one case and return a data dict for bouquet_method. Used when is_atomic=True.""" + raise NotImplementedError + + def atomic_input_recast(self, all_input_files) -> list: + """Expand non-atomic inputs into a flat list of input tuples for use by atomic_load_files.""" + raise NotImplementedError + + @property + def atomic_load_files(self): + """Load one case using inputs from atomic_input_recast. Used when is_atomic=False.""" + raise NotImplementedError + +#################################################################### +# Atomic load_profile_obj +#################################################################### + +class load_profile_obj(load_files_obj): + """Load generic (geqdsk, kinetic_profile_file) pairs for re_generate_bouquet. + + Each entry in all_input_files is a tuple ``(geqdsk_path, profile_path)``, + one per run. Input is already atomic so no recasting is needed. + + The profile_reader and uncertainty_generator in config must match the + specific type of kinetic profile file used (e.g. p-file). + + Parameters + ---------- + config : dict + ... + """ + is_atomic = True + + def __init__(self, config): + self.config = config + + def total_runs(self, all_input_files): + # Assume all_input_files is a vector of (geqdsk, pfile) pairs + return len(all_input_files), _IndexMap(all_input_files) + + def load_files(self, input_files, idx): + # Take one (geqdsk, kinetic_profile_file) pair, returns data dict for bouquet method + geqdsk_file, profile_file = input_files + worker_id = _worker_state["worker_id"] + read_geqdsk = _worker_state["read_geqdsk"] + profile_reader = self.config["profile_reader"] + uncertainty_gen = self.config["uncertainty_generator"] + + _tag = f"[Worker {worker_id} | run {idx} | {os.path.basename(geqdsk_file)}]" + print(f"{_tag} Starting — host={socket.gethostname()}, PID={os.getpid()}, cwd={os.getcwd()}", flush=True) + + # Copy input files into the worker's private working directory so that + # every file read or write by TokaMaker stays within a single directory. + # The idx prefix prevents collisions when a worker processes multiple + # equilibria that share the same base filename. + _local_geqdsk = os.path.join(os.getcwd(), f"idx{idx}_{os.path.basename(geqdsk_file)}") + _local_profile = os.path.join(os.getcwd(), f"idx{idx}_{os.path.basename(profile_file)}") + shutil.copy2(geqdsk_file, _local_geqdsk) + shutil.copy2(profile_file, _local_profile) + geqdsk_file = _local_geqdsk + profile_file = _local_profile + print(f"{_tag} Copied input files to worker directory.", flush=True) + + # --- Load equilibrium --- + eqdsk = read_geqdsk(geqdsk_file) + psi_N = eqdsk.psi_N + + # --- Read kinetic profiles via the pluggable reader --- + # Returns profiles, Zeff, raw bytes for HDF5 archival + ne_SI, te_SI, ni_SI, ti_SI, Zeff, psi_N_kinetic, profile_bytes = profile_reader( + profile_file, self.config["profile_reader_kwargs"] + ) + + # --- Generate profile uncertainties via the pluggable generator --- + # Returns sigma_ne, sigma_te, sigma_ni, sigma_ti on the kinetic profile grid (psi_N_kinetic), + # and optionally sigma_jphi on the equilibrium grid (psi_N). + sigma_ne, sigma_te, sigma_ni, sigma_ti, psi_N_kinetic, sigma_jphi = uncertainty_gen( + profile_file, profile_reader, psi_N, self.config["profile_reader_kwargs"], self.config["uncertainty_generator_kwargs"] + ) + + return { + "idx": idx, + "eqdsk": eqdsk, + "psi_N": psi_N, + "psi_N_kinetic": psi_N_kinetic, + "ne_SI": ne_SI, + "te_SI": te_SI, + "ni_SI": ni_SI, + "ti_SI": ti_SI, + "w_ExB": np.zeros_like(psi_N), + "Zeff": Zeff, + "profile_bytes": profile_bytes, + "sigma_ne": sigma_ne, + "sigma_te": sigma_te, + "sigma_ni": sigma_ni, + "sigma_ti": sigma_ti, + "sigma_jphi": sigma_jphi, + } + +################################## +# pfile specific reader and uncertainty generator for atomic load_profile_obj +################################## +def pfile_reader(pfile_file, reader_kwargs): + """Read an Osborne p-file and return SI kinetic profiles plus raw bytes. + + Matches the ``config['profile_reader']`` contract expected by + :class:`~bouquet.parallel.load_profile_obj`. + + All profiles are remapped onto ``ne``'s ``psinorm`` grid before use + (p-files allow each profile to carry its own independent grid). + Interpolation onto the equilibrium grid is handled downstream via the + ``psi_N_kinetic`` argument to + :func:`~bouquet.TokaMaker_interface.generate_bouquet`. + + Parameters + ---------- + pfile_file : str + Path to the p-file. + reader_kwargs : dict + Must contain ``ion_N``, ``ion_Z``, ``ion_A`` — the number of ions + per formula unit, charge state, and mass number of the main ion + species (e.g. ``{"ion_N": 1, "ion_Z": 1, "ion_A": 2}`` for + deuterium). + + Returns + ------- + ne_SI, te_SI, ni_SI, ti_SI : ndarray + Kinetic profiles in SI units (m\ :sup:`-3` and eV) on *psi_N_kinetic*. + Zeff_eq : ndarray + Effective ion charge on *psi_N_kinetic*, clipped to >= 1. + psi_N_kinetic : ndarray + Normalised poloidal flux grid of the p-file (``psinorm``). + profile_bytes : bytes + Raw p-file content for HDF5 archival. + """ + ion_N = reader_kwargs['ion_N'] + ion_Z = reader_kwargs['ion_Z'] + ion_A = reader_kwargs['ion_A'] + + from bouquet.io.pfile import read_pfile + pf = read_pfile(pfile_file) + pf = pf.remap(key='ne') + + if pf.ion_species is None: + pf.set_ion_species(N=ion_N, Z=ion_Z, A=ion_A) + pf.compute_quasineutrality() + psi_N_kinetic, Zeff = pf.compute_zeff() + + ne_SI = pf.ne * 1e20 # 10^20 m^-3 -> m^-3 + te_SI = pf.te * 1e3 # keV -> eV + ni_SI = pf.ni * 1e20 + ti_SI = pf.ti * 1e3 + Zeff_eq = np.clip(Zeff, 1.0, None) + + with open(pfile_file, 'rb') as fh: + profile_bytes = fh.read() + + return ne_SI, te_SI, ni_SI, ti_SI, Zeff_eq, psi_N_kinetic, profile_bytes + +def pfile_uncertainty_gen(profile_file, profile_reader_fn, psi_N, reader_kwargs, uncertainty_kwargs): + """Build radially-varying 1-sigma uncertainties from the p-file profiles. + + Matches the ``config['uncertainty_generator']`` contract expected by + :class:`~bouquet.parallel.load_profile_obj`. + + Parameters + ---------- + profile_file : str + Path to the p-file. + profile_reader_fn : callable + A ``pfile_reader``-style callable used to load baseline profiles so + that fractional uncertainties can be converted to absolute ones. + Called as ``profile_reader_fn(profile_file, reader_kwargs)``. + psi_N : ndarray + Equilibrium normalised-flux grid. Not used for kinetic sigmas + (those live on *psi_N_kinetic*); available for ``sigma_jphi`` if + needed. + reader_kwargs : dict + Forwarded to *profile_reader_fn* unchanged. + uncertainty_kwargs : dict + Must contain: ``frac_ne``, ``frac_te``, ``frac_ni``, ``frac_ti`` + (fractional 1-sigma levels); ``falloff_ne``, ``falloff_te``, + ``falloff_ni``, ``falloff_ti`` (radial falloff exponents); ``shelf`` + (minimum fractional floor). + + Returns + ------- + sigma_ne, sigma_te, sigma_ni, sigma_ti : ndarray + Absolute 1-sigma arrays in SI units on *psi_N_kinetic*. + psi_N_kinetic : ndarray + Normalised flux grid for the kinetic sigma arrays (from the p-file). + sigma_jphi : None + Deferred — computed inside :func:`~bouquet.parallel.re_generate_bouquet` + from ``j_phi_fit`` via ``config['jphi_uncertainty_gen']``. + """ + frac_ne = uncertainty_kwargs['frac_ne'] + frac_te = uncertainty_kwargs['frac_te'] + frac_ni = uncertainty_kwargs['frac_ni'] + frac_ti = uncertainty_kwargs['frac_ti'] + falloff_ne = uncertainty_kwargs['falloff_ne'] + falloff_te = uncertainty_kwargs['falloff_te'] + falloff_ni = uncertainty_kwargs['falloff_ni'] + falloff_ti = uncertainty_kwargs['falloff_ti'] + shelf = uncertainty_kwargs['shelf'] + + from bouquet.uncertainties import new_uncertainty_profiles + ne_SI, te_SI, ni_SI, ti_SI, _, psi_N_kinetic, _ = profile_reader_fn(profile_file, reader_kwargs) + sigma_ne = new_uncertainty_profiles(psi_N_kinetic, frac_ne, falloff_exp=falloff_ne, shelf=shelf) * ne_SI + sigma_te = new_uncertainty_profiles(psi_N_kinetic, frac_te, falloff_exp=falloff_te, shelf=shelf) * te_SI + sigma_ni = new_uncertainty_profiles(psi_N_kinetic, frac_ni, falloff_exp=falloff_ni, shelf=shelf) * ni_SI + sigma_ti = new_uncertainty_profiles(psi_N_kinetic, frac_ti, falloff_exp=falloff_ti, shelf=shelf) * ti_SI + return sigma_ne, sigma_te, sigma_ni, sigma_ti, psi_N_kinetic, None + +#################################################################### +# Non-atomic load_profile_obj types +#################################################################### + +class load_IDA_file_obj(load_files_obj): + """Load timeslices from IDA-lite CDF files for re_generate_bouquet. + + Each entry in all_input_files is a tuple ``(cdf_path, geqdsk_paths)``. + A single CDF may contain multiple timeslices, so ``is_atomic=False``: + ``atomic_input_recast`` expands each CDF into ``(cdf_path, geqdsk_paths, time_idx)`` + triples and ``atomic_load_files`` loads one such triple. + + ``geqdsk_paths`` must be a list/tuple of paths with one geqdsk per timeslice + in the CDF, in time-index order. + + jphi uncertainty is not calculated, so the config dict must contain: + + - ``jphi_uncertainty_gen`` : callable, e.g. ``FractionalUncertainty(0.15)`` + + Optional config keys: + + - ``profile_reader_kwargs`` : dict of keyword arguments forwarded to + :class:`~bouquet.io.ida.IDALiteProfileReader` (excluding ``time_idx``). + - ``uncertainty_generator_kwargs`` : dict of keyword arguments forwarded to + :class:`~bouquet.io.ida.IDALiteUncertaintyGenerator` (excluding ``time_idx``). + + Parameters + ---------- + config : dict + ... + """ + is_atomic = False + + def __init__(self, config): + self.config = config + + def total_runs(self, all_input_files): + """Sum timeslice counts across all CDF files and build a map from flat + idx to ``(cdf_path, geqdsk_path, time_idx)`` atomic input triple.""" + import h5py + atomic = [] + for cdf_path, geqdsk_paths in all_input_files: + with h5py.File(cdf_path, 'r') as f: + n_times = f['n_e'].shape[0] + if len(geqdsk_paths) != n_times: + raise ValueError( + f"geqdsk_paths has {len(geqdsk_paths)} entries but " + f"'{cdf_path}' contains {n_times} timeslice(s). " + "Provide exactly one geqdsk path per timeslice." + ) + for t in range(n_times): + atomic.append((cdf_path, geqdsk_paths[t], t)) + return len(atomic), _IndexMap(atomic) + + def atomic_input_recast(self, all_input_files) -> list: + """Expand each (cdf_path, geqdsk_paths) into flat (cdf_path, geqdsk, time_idx) triples.""" + n_runs, map_object = self.total_runs(all_input_files) + return [map_object(i) for i in range(n_runs)] + + def load_files_atomic(self, input_files, idx): + """Load one (cdf_path, geqdsk_file, time_idx) triple and return a data dict.""" + from bouquet.io.ida import IDALiteProfileReader, IDALiteUncertaintyGenerator + cdf_path, geqdsk_file, time_idx = input_files + worker_id = _worker_state["worker_id"] + read_geqdsk = _worker_state["read_geqdsk"] + profile_reader = IDALiteProfileReader( + **self.config.get("profile_reader_kwargs", {}), time_idx=time_idx + ) + uncertainty_gen = IDALiteUncertaintyGenerator( + **self.config.get("uncertainty_generator_kwargs", {}), time_idx=time_idx + ) + + _tag = ( + f"[Worker {worker_id} | run {idx} | " + f"{os.path.basename(cdf_path)} t={time_idx}]" + ) + print( + f"{_tag} Starting — host={socket.gethostname()}, " + f"PID={os.getpid()}, cwd={os.getcwd()}", + flush=True, + ) + + # Copy input files into the worker's private working directory. + _local_geqdsk = os.path.join(os.getcwd(), f"idx{idx}_{os.path.basename(geqdsk_file)}") + _local_cdf = os.path.join(os.getcwd(), os.path.basename(cdf_path)) + shutil.copy2(geqdsk_file, _local_geqdsk) + if not os.path.exists(_local_cdf): + shutil.copy2(cdf_path, _local_cdf) + geqdsk_file = _local_geqdsk + cdf_path = _local_cdf + print(f"{_tag} Copied input files to worker directory.", flush=True) + + # Load equilibrium + eqdsk = read_geqdsk(geqdsk_file) + psi_N = eqdsk.psi_N + + # IDALiteProfileReader interpolates profiles directly onto psi_N, + # so psi_N_kinetic=None. + ne_SI, te_SI, ni_SI, ti_SI, Zeff_eq, _ = profile_reader( + geqdsk_file, cdf_path, psi_N + ) + profile_bytes = None + + # Compute kinetic-profile sigmas. sigma_jphi is deferred: + # re_generate_bouquet will call config['jphi_uncertainty_gen']. + sigma_ne, sigma_te, sigma_ni, sigma_ti, _ = uncertainty_gen( + cdf_path, None, psi_N, np.ones_like(psi_N) + ) + + return { + "idx": idx, + "eqdsk": eqdsk, + "psi_N": psi_N, + "psi_N_kinetic": None, + "ne_SI": ne_SI, + "te_SI": te_SI, + "ni_SI": ni_SI, + "ti_SI": ti_SI, + "w_ExB": np.zeros_like(psi_N), + "Zeff": Zeff_eq, + "profile_bytes": profile_bytes, + "sigma_ne": sigma_ne, + "sigma_te": sigma_te, + "sigma_ni": sigma_ni, + "sigma_ti": sigma_ti, + "sigma_jphi": None, + } + + @property + def atomic_load_files(self): + return self.load_files_atomic + +########################################################################################################### +# utils +########################################################################################################### + +def _mesh_config_simp(mygs, config, local_mesh_file): + """Load mesh and configure TokaMaker with optional VSC and coil regularisation. + + A simple mesh configuration function suitable for passing as + ``config['mesh_config_function']``. Loads the worker-local mesh copy, + sets up the FE mesh and conductor/coil regions, applies solver settings + from *config*, and optionally configures a vertical stability coil and + coil current regularisation targets. + + Parameters + ---------- + mygs : TokaMaker + TokaMaker instance to configure (already constructed, not yet set up). + config : dict + Shared configuration dict. Expected keys: ``oft_order``, + ``oft_maxits``. Optional keys: ``vsc_coil_def``, ``target_currents``. + local_mesh_file : str + Absolute path to the worker-local copy of the mesh HDF5 file. + """ + from OpenFUSIONToolkit.TokaMaker.meshing import load_gs_mesh + mesh_pts, mesh_lc, mesh_reg, coil_dict, cond_dict = load_gs_mesh(local_mesh_file) + mygs.setup_mesh(mesh_pts, mesh_lc, mesh_reg) + mygs.setup_regions(cond_dict=cond_dict, coil_dict=coil_dict) + + mygs.setup(order=config["oft_order"]) + mygs.settings.maxits = config["oft_maxits"] + mygs.settings.pm = config.get("oft_pm", False) + mygs.update_settings() + + vsc_coil_def = config.get("vsc_coil_def") + if vsc_coil_def is not None: + mygs.set_coil_vsc(vsc_coil_def) + + target_currents = config.get("target_currents") + if target_currents is not None: + reg_terms = [] + for coil_name, val_ma in target_currents.items(): + reg_terms.append( + mygs.coil_reg_term({coil_name: 1.0}, target=val_ma * 1e6, weight=1.0) + ) + reg_terms.append( + mygs.coil_reg_term({"#VSC": 1.0}, target=0.0, weight=1e-2) + ) + mygs.set_coil_reg(reg_terms=reg_terms) + +def _init_OFT(worker_id_queue, master_working_dir, config, init_status_queue): + """Pool initialiser: set up OFT/TokaMaker once per spawned worker process. + + Called automatically by ``Pool`` (via ``load_files_obj.init_worker``) before + any tasks are dispatched. Each worker claims a unique ID from + *worker_id_queue*, creates a private working directory, copies the mesh + file locally, initialises OFT and TokaMaker, and stores all shared state + in the module-level ``_worker_state`` dict for use by ``_run_one``. + + On success posts ``(worker_id, None)`` to *init_status_queue*. + On failure posts ``(worker_id, traceback_str)``. + + Parameters + ---------- + worker_id_queue : multiprocessing.Queue + Pre-loaded with integers 0…n_workers−1. Each worker pops one value + to claim its unique ID. + master_working_dir : str + Root directory under which per-worker subdirectories are created. + config : dict + Shared configuration dict (general options, not worker-specific state). + Must contain ``mesh_file``, ``header``, ``mesh_config_function``, and + any keys required by that function (e.g. ``oft_order``). + init_status_queue : multiprocessing.Queue + Used to signal initialisation success or failure back to the main + process barrier. + """ + global _worker_state + worker_id = -1 # fallback if queue.get() itself fails + try: + nthreads = config.get("_nthreads", 1) + os.environ["OMP_NUM_THREADS"] = str(nthreads) + os.environ["MKL_NUM_THREADS"] = str(nthreads) + os.environ["OPENBLAS_NUM_THREADS"] = str(nthreads) + os.environ["NUMEXPR_NUM_THREADS"] = str(nthreads) + print( + f"[Worker {worker_id}] thread env: " + f"OMP_NUM_THREADS={os.environ.get('OMP_NUM_THREADS')} " + f"MKL_NUM_THREADS={os.environ.get('MKL_NUM_THREADS')} " + f"OPENBLAS_NUM_THREADS={os.environ.get('OPENBLAS_NUM_THREADS')} " + f"NUMEXPR_NUM_THREADS={os.environ.get('NUMEXPR_NUM_THREADS')}", + flush=True, + ) + + # Use a timeout so replacement workers (spawned after a crash) fail fast + # rather than blocking forever and deadlocking the parent imap_unordered. + try: + worker_id = worker_id_queue.get(timeout=60) + except Exception: + raise RuntimeError( + "[bouquet_parallel] Worker ID queue empty — this is a pool replacement " + "for a dead worker. Cannot initialise." + ) + working_dir = os.path.abspath(os.path.join(master_working_dir, f"worker_{worker_id}")) + os.makedirs(working_dir, exist_ok=True) + master_working_dir = os.path.abspath(master_working_dir) # ← anchor before chdir + os.chdir(working_dir) + + # Redirect this worker's stdout/stderr to a per-worker log file. + # os.dup2 at the file-descriptor level also captures + # output written directly to fd 1/2 by Fortran/C extensions (e.g. OFT). + # Skipped when config["_verbose"] is True so output goes to the terminal. + log_path = os.path.join(master_working_dir, f"worker_{worker_id}.log") + if not config.get("_verbose", False): + _log_fh = open(log_path, "w", buffering=1) # line-buffered + os.dup2(_log_fh.fileno(), 1) + os.dup2(_log_fh.fileno(), 2) + sys.stdout = _log_fh + sys.stderr = _log_fh + + # Add the OFT python directory to sys.path if supplied. + # This is required when using spawned processes because sys.path + # modifications in the parent process are not inherited by children. + oft_python_path = config.get("oft_python_path") + if oft_python_path and oft_python_path not in sys.path: + sys.path.insert(0, oft_python_path) + + # Copy the mesh HDF5 into this worker's private directory so that + # concurrent HDF5 opens by multiple workers do not trigger file-locking + # conflicts in a serial HDF5 build. working_dir is already absolute so + # local_mesh_file is absolute regardless of what os.getcwd() is now. + local_mesh_file = os.path.join(working_dir, os.path.basename(config["mesh_file"])) + shutil.copy2(config["mesh_file"], local_mesh_file) + + from OpenFUSIONToolkit import OFT_env + from OpenFUSIONToolkit.TokaMaker import TokaMaker + from OpenFUSIONToolkit.TokaMaker.util import create_power_flux_fun + + from bouquet import ( + read_geqdsk, + reconstruct_equilibrium, + generate_bouquet, + initialize_equilibrium_database, + store_equilibrium, + ) + + myOFT = OFT_env(nthreads=nthreads) + mygs = TokaMaker(myOFT) + + config['mesh_config_function'](mygs, config, local_mesh_file) + + print( + f"[Worker {worker_id}] OFT initialised — " + f"host={socket.gethostname()}, PID={os.getpid()}, " + f"cwd={working_dir}", + flush=True, + ) + + _worker_state.update({ + "worker_id": worker_id, + "working_dir": working_dir, + "log_path": log_path, + "config": config, + "mygs": mygs, + "read_geqdsk": read_geqdsk, + "reconstruct_equilibrium": reconstruct_equilibrium, + "generate_bouquet": generate_bouquet, + "store_equilibrium": store_equilibrium, + "initialize_equilibrium_database": initialize_equilibrium_database, + "create_power_flux_fun": create_power_flux_fun, + }) + + init_status_queue.put((worker_id, None)) # signal success to main process + + except Exception: + tb = traceback.format_exc() + print(f"[Worker {worker_id}] INIT FAILED:\n{tb}", flush=True) + try: + init_status_queue.put((worker_id, tb)) + except Exception: + pass + raise # kill this worker process + + +# Assign init_worker after _init_OFT is defined to avoid forward-reference error. +load_files_obj.init_worker = staticmethod(_init_OFT) \ No newline at end of file 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/examples/D3D-like/parallel_IDA_example.py b/examples/D3D-like/parallel_IDA_example.py new file mode 100644 index 0000000..9cc19dd --- /dev/null +++ b/examples/D3D-like/parallel_IDA_example.py @@ -0,0 +1,294 @@ +""" +Parallel bouquet runner — DIII-D-like IDA example +============================================== + +Runs `re_generate_bouquet` on a D3D IDA file and eqdsk list. +""" + +import os +import sys +import numpy as np +import matplotlib +matplotlib.use('Agg') # headless — remove for interactive use +import matplotlib.pyplot as plt + +# file options +PLOT_ONLY = False +remake_dir = True # If true, deletes pre-existing working directory on re-runs +use_python_solve = False # Use python bootstrap solve +verbose=True # If false, worker outputs are printed to individual log files +use_logical_cpus=True # Multi-thread based on hardware (use with caution if you're not on linux) + +# --------------------------------------------------------------------------- +# OFT / TokaMaker path — adjust to your installation +# --------------------------------------------------------------------------- +OFT_PATH = '/home/stubenj9/src/OpenFUSIONToolkit/builds/install_release/python' +if OFT_PATH: + sys.path.append(OFT_PATH) + +# Add bouquet root so the package is importable when run directly +sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..')) + +# ============================================================================ +# 1. Input files +# ============================================================================ +# All paths are relative to this script's directory so the example works +# regardless of where you run it from. + +HERE = os.path.dirname(os.path.realpath(__file__)) + +geqdsks = ["","","",...,""] +IDA_filename = "" + +# DIII-D mesh from OFT examples +MESH_FILE = os.path.join( + HERE, '../../../OpenFUSIONToolkit/src/examples/TokaMaker/DIIID', + 'DIIID_mesh.h5', +) + +# Output directory — each worker gets its own subdirectory +OUTPUT_DIR = os.path.join(HERE, 'output_parallel_IDA') + +# HDF5 database base name (one per worker:

_worker.h5) +HEADER = 'TkMkr_D3Dlike_Hmode_parallel_IDA' + +# ============================================================================ +# 2. Load bouquet_method (general per-run worker function) +# ============================================================================ + +from bouquet.parallel import ( + re_generate_bouquet, load_IDA_file_obj, _mesh_config_simp, + FractionalUncertainty, +) + +# bouquet_method is the per-run worker function called by parallel_runner +bouquet_method = re_generate_bouquet + +# ============================================================================ +# 3. Load and configure load_files_obj (most method-specific settings go here) +# ============================================================================ + +# j_phi fractional uncertainty (applied to j_phi_fit in re_generate_bouquet) +frac_jphi = 0.10 # 10 % on j_phi + +# IDALiteProfileReader keyword overrides (all optional — defaults are reasonable) +ida_reader_kwargs = { + 'carbon_quasi_neutrality': True, # use n_12C6 CDF variable for ni +} + +# IDALiteUncertaintyGenerator keyword overrides (all optional) +ida_uncertainty_kwargs = {} + +target_currents = { + 'ECOILA': -0.977888676757812 / 61.0, + 'ECOILB': -0.962711173828125 / 61.0, + 'F1A': 0.115971984375, 'F1B': 0.128368578125, + 'F2A': 0.05980789453125, 'F2B': 0.0763701328125, + 'F3A': -0.03076001171875, 'F3B': -0.02234583203125, + 'F4A': -0.0401077421875, 'F4B': -0.13314096875, + 'F5A': 0.000723009033203125, 'F5B': 0.000399709045410156, + 'F6A': -0.1178582578125, 'F6B': -0.15356990625, + 'F7A': 0.0341264296875, 'F7B': 0.0415082109375, + 'F8A': -0.05660116015625, 'F8B': -0.05138975390625, + 'F9A': 0.236625375, 'F9B': 0.252380265625, +} + +n_equils = 5 # perturbed equilibria per baseline +n_ls = 0.5 # GPR correlation length — density (psi_N units) +t_ls = 0.4 # GPR correlation length — temperature +j_ls = 0.25 # GPR correlation length — current density +jBS_scale_range = [0.9, 1.1] # uniform random scale on j_BS per sample +pad_psi = 1e-4 # LCFS psi padding for TokaMaker queries + +# reconstruct_equilibrium settings +n_k = 5 +psi_bridge = 0.99 +l_i_tolerance = 5.0 # percent +constrain_sawteeth = True +recalculate_j_BS = True +jphi_baseline = True + +# ---- coil-drift / homotopy / in-spec (DIII-D +/-2% measurement spec) ---- +# Everything is a FRACTION (decimal), never a percentage. +coil_drift = 0.01 # +/-1% hard coil-current bound +homotopy_passes = [(0.1, 0.10), (0.02, 0.05), (0.015, 0.03)] # (F, VSC) limits +inspec_F_max = 0.02 # in-spec non-VSC F-coil drift +inspec_VSC_max = 0.02 # in-spec VSC (F9A/F9B) drift +vsc_soft_reg_weight = 1.0 +p_thresh = 0.05 # pressure-match tolerance + +isoflux_weight = 500.0 # uniform weight on all isoflux boundary points + +# Set True to keep per-equilibrium .geqdsk files after archiving to HDF5. +# Useful for manual inspection or debugging. +KEEP_GEQDSK = True + +config = { + # --- TokaMaker / mesh --- + "mesh_file": MESH_FILE, + "header": HEADER, + "mesh_config_function": _mesh_config_simp, + "oft_order": 3, + "oft_maxits": 50, + "oft_python_path": OFT_PATH, + # --- Profile I/O (IDA) --- + "profile_reader_kwargs": ida_reader_kwargs, + "uncertainty_generator_kwargs": ida_uncertainty_kwargs, + # --- Bouquet sampling --- + "n_equils": n_equils, + "n_ls": n_ls, + "t_ls": t_ls, + "j_ls": j_ls, + "psi_pad": pad_psi, + "isoflux_weight": isoflux_weight, + "jphi_uncertainty_gen": FractionalUncertainty(frac_jphi), + "keep_geqdsk": KEEP_GEQDSK, + # --- Optional coil regularisation --- + "target_currents": target_currents, + # --- reconstruct_equilibrium keyword overrides --- + "reconstruct_equilibrium_kwargs": { + "n_k": n_k, + "psi_bridge": psi_bridge, + #"taper_edge_jBS": False, + "use_python_solve": use_python_solve, + }, + "generate_bouquet_kwargs": { + "l_i_tolerance": l_i_tolerance, + "psi_pad": pad_psi, + "constrain_sawteeth": constrain_sawteeth, + "jBS_scale_range": jBS_scale_range, + "recalculate_j_BS": recalculate_j_BS, + #"taper_edge_jBS": False, + "use_python_solve": use_python_solve, + "jphi_baseline": jphi_baseline, + "coil_drift": coil_drift, + "homotopy_passes": homotopy_passes, + "inspec_F_max": inspec_F_max, + "inspec_VSC_max": inspec_VSC_max, + "vsc_soft_reg_weight": vsc_soft_reg_weight, + "p_thresh": p_thresh, + } +} + +load_files_obj = load_IDA_file_obj(config) + +# ============================================================================ +# 4. Pre-flight checks +# ============================================================================ + +def _compute(): + from bouquet.parallel import _get_num_cpus as _parallel_get_num_cpus + n_cpus, _ = _parallel_get_num_cpus() + if n_cpus < 2: + print(f'Warning: only {n_cpus} CPU core(s) available. ' + 'Parallel run requires at least 2 cores.') + sys.exit(0) + print(f'Running parallel run with {n_cpus} CPU core(s) available.') + + # Clean output directory for a fresh run + if os.path.exists(OUTPUT_DIR) and remake_dir: + import shutil as _shutil + _shutil.rmtree(OUTPUT_DIR) + os.makedirs(OUTPUT_DIR, exist_ok=True) + + print('\nChecking required input files...') + missing = [f for f in geqdsks + [IDA_filename, MESH_FILE] if not os.path.exists(f)] + if missing: + print('ERROR: the following files were not found:') + for f in missing: + print(f' {f}') + print('\nAdjust OFT_PATH / MESH_FILE and retry.') + sys.exit(1) + print(f' {len(geqdsks)} geqdsk file(s) and IDA CDF found.') + print(f' Mesh: {MESH_FILE}') + +# ============================================================================ +# 5. Compute in parallel +# ============================================================================ + print(f'\nLaunching parallel run into: {OUTPUT_DIR}') + print(f' {len(geqdsks)} equilibria, {n_equils} perturbed samples each') + print() + + from bouquet.parallel import parallel_runner + errors, outputs = parallel_runner( + [(IDA_filename, geqdsks)], + load_files_obj, + bouquet_method, + OUTPUT_DIR, + use_logical_cpus=True, + verbose=True, + ) + +# ============================================================================ +# 6. Error report +# ============================================================================ + if errors: + print(f'\nWARNING: {len(errors)} run(s) failed:') + for idx, tb in errors.items(): + print(f' [run {idx}] {tb.splitlines()[-1]}') + else: + print(f'\nAll runs completed successfully.') + return errors + +if __name__ == '__main__': + + if PLOT_ONLY: + errors = {} + print(f'\nPLOT_ONLY mode: skipping compute, loading results from:\n {OUTPUT_DIR}') + else: + errors = _compute() + + # ============================================================================ + # 7. Visualize results + # ============================================================================ + # parallel_runner saves one HDF5 database per run, named: + # {HEADER}_idx{idx}.h5 + # located in worker subdirectories under OUTPUT_DIR. + + import glob + import pickle as pkl + from collections import defaultdict + from bouquet import read_geqdsk, plot_bouquet, plot_geqdsk_bouquet + + # Collect all HDF5 databases + h5_files = sorted(glob.glob(os.path.join(OUTPUT_DIR, '**', f'{HEADER}_idx*.h5'), recursive=True)) + + if not h5_files: + print('\nNo HDF5 result files found; skipping plots.') + sys.exit(0 if not errors else 1) + + print(f'\n{"="*60}') + print(f'Visualizing results ({len(h5_files)} equilibrium/database file(s))') + print('=' * 60) + + # Baseline g-file shapes + print('\nPlotting baseline g-files...') + plot_geqdsk_bouquet(geqdsks, x_coord='rho') + out = os.path.join(OUTPUT_DIR, f'{HEADER}_baseline_geqdsk.png') + plt.savefig(out, dpi=150, bbox_inches='tight') + plt.close() + print(f' Saved: {out}') + + # Per-database plots (one per run/equilibrium) + for i, h5_path in enumerate(h5_files): + tag = f'idx{i}' + + try: + print(f'\nPlotting bouquet — {tag}...') + plot_bouquet(h5_path, scan_value=None, mode='all') + out = os.path.join(OUTPUT_DIR, f'{HEADER}_bouquet_{tag}.png') + plt.savefig(out, dpi=150, bbox_inches='tight') + plt.close() + print(f' Saved: {out}') + except KeyError as e: + print(f' Skipped (empty file): {e}') + + try: + print(f'Plotting perturbed g-files — {tag}...') + plot_geqdsk_bouquet(h5path=h5_path, x_coord='rho') + out = os.path.join(OUTPUT_DIR, f'{HEADER}_perturbed_geqdsk_{tag}.png') + plt.savefig(out, dpi=150, bbox_inches='tight') + plt.close() + print(f' Saved: {out}') + except (KeyError, ValueError) as e: + print(f' Skipped (empty file): {e}') diff --git a/examples/D3D-like/parallel_pfile_example.py b/examples/D3D-like/parallel_pfile_example.py new file mode 100644 index 0000000..9212bd0 --- /dev/null +++ b/examples/D3D-like/parallel_pfile_example.py @@ -0,0 +1,336 @@ +""" +Parallel bouquet runner — DIII-D-like example +============================================== + +Runs :func:`re_generate_bouquet` on the D3D-like +H-mode equilibrium/profile pairs bundled in this directory. +""" + +import os +import sys +import numpy as np +import matplotlib +matplotlib.use('Agg') # headless — remove for interactive use +import matplotlib.pyplot as plt + +# file options +PLOT_ONLY = False +remake_dir = True # If true, deletes pre-existing working directory on re-runs +use_python_solve = False # Use python bootstrap solve +verbose=True # If false, worker outputs are printed to individual log files +use_logical_cpus=True # Multi-thread based on hardware (use with caution if you're not on linux) + +# --------------------------------------------------------------------------- +# OFT / TokaMaker path — adjust to your installation +# --------------------------------------------------------------------------- +OFT_PATH = '/home/stubenj9/src/OpenFUSIONToolkit/builds/install_release/python' +if OFT_PATH: + sys.path.append(OFT_PATH) + +# Add bouquet root so the package is importable when run directly +sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..')) + +# ============================================================================ +# 1. Input files +# ============================================================================ +# All paths are relative to this script's directory so the example works +# regardless of where you run it from. + +HERE = os.path.dirname(os.path.abspath(__file__)) + +geqdsks = [ + os.path.join(HERE, 'D3Dlike_Hmode_baseline.geqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.geqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.geqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.geqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.geqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.geqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.geqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.geqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.geqdsk'), +] +pfiles = [ + os.path.join(HERE, 'D3Dlike_Hmode_baseline.peqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.peqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.peqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.peqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.peqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.peqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.peqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.peqdsk'), + os.path.join(HERE, 'D3Dlike_Hmode_baseline.peqdsk'), +] + +# DIII-D mesh from OFT examples +MESH_FILE = os.path.join( + HERE, '../../../OpenFUSIONToolkit/src/examples/TokaMaker/DIIID', + 'DIIID_mesh.h5', +) + +# Output directory — each worker gets its own subdirectory +OUTPUT_DIR = os.path.join(HERE, 'output_pfile_parallel') + +# HDF5 database base name (one per worker:
_worker.h5) +HEADER = 'TkMkr_D3Dlike_Hmode_parallel' + +# ============================================================================ +# 2. Load bouquet_method (general per-run worker function) +# ============================================================================ + +from bouquet.parallel import ( + re_generate_bouquet, load_profile_obj, _mesh_config_simp, + pfile_reader, pfile_uncertainty_gen, FractionalUncertainty, +) + +# bouquet_method is the per-run worker function called by parallel_runner +bouquet_method = re_generate_bouquet + +# ============================================================================ +# 3. Load and configure load_files_obj (most method-specific settings go here) +# ============================================================================ + +# Ion species: deuterium main ion (adjust Z/A for other machines or impurity models) +reader_kwargs = { + 'ion_N': 1, # ions per formula unit + 'ion_Z': 1, # charge state + 'ion_A': 2, # mass number (deuterium) +} + +# Fractional 1-sigma uncertainties and radial falloff shapes for kinetic profiles +frac_jphi = 0.10 # 10 % on j_phi (applied to j_phi_fit in re_generate_bouquet) +uncertainty_kwargs = { + 'frac_ne': 0.05, # 5% on electron density + 'frac_te': 0.05, # 5% on electron temperature + 'frac_ni': 0.05, # 5% on ion density + 'frac_ti': 0.05, # 5% on ion temperature + 'falloff_ne': 2.0, # power-law radial falloff exponent + 'falloff_te': 2.0, + 'falloff_ni': 2.0, + 'falloff_ti': 2.0, + 'shelf': 0.05, # minimum fractional uncertainty floor (prevents zero at edge) +} + +target_currents = { + 'ECOILA': -0.977888676757812 / 61.0, + 'ECOILB': -0.962711173828125 / 61.0, + 'F1A': 0.115971984375, 'F1B': 0.128368578125, + 'F2A': 0.05980789453125, 'F2B': 0.0763701328125, + 'F3A': -0.03076001171875, 'F3B': -0.02234583203125, + 'F4A': -0.0401077421875, 'F4B': -0.13314096875, + 'F5A': 0.000723009033203125, 'F5B': 0.000399709045410156, + 'F6A': -0.1178582578125, 'F6B': -0.15356990625, + 'F7A': 0.0341264296875, 'F7B': 0.0415082109375, + 'F8A': -0.05660116015625, 'F8B': -0.05138975390625, + 'F9A': 0.236625375, 'F9B': 0.252380265625, +} + +n_equils = 5 # perturbed equilibria per baseline +n_ls = 0.5 # GPR correlation length — density (psi_N units) +t_ls = 0.4 # GPR correlation length — temperature +j_ls = 0.25 # GPR correlation length — current density +jBS_scale_range = [0.9, 1.1] # uniform random scale on j_BS per sample +pad_psi = 1e-4 # LCFS psi padding for TokaMaker queries + +# settings +n_k = 5 +psi_bridge = 0.99 +l_i_tolerance = 0.05 +constrain_sawteeth = True +recalculate_j_BS = True +jphi_baseline = True + +# ---- coil-drift / homotopy / in-spec (DIII-D +/-2% measurement spec) ---- +# Everything is a FRACTION (decimal), never a percentage. +coil_drift = 0.01 # +/-1% hard coil-current bound +homotopy_passes = [(0.1, 0.10), (0.02, 0.05), (0.015, 0.03)] # (F, VSC) limits +inspec_F_max = 0.02 # in-spec non-VSC F-coil drift +inspec_VSC_max = 0.02 # in-spec VSC (F9A/F9B) drift +vsc_soft_reg_weight = 1.0 +p_thresh = 0.05 # pressure-match tolerance + +isoflux_weight = 500.0 # uniform weight on all isoflux boundary points + +# Set True to keep per-equilibrium .geqdsk files after archiving to HDF5. +# Useful for manual inspection or debugging. +KEEP_GEQDSK = True + +config = { + # --- TokaMaker / mesh --- + "mesh_file": MESH_FILE, + "header": HEADER, + "mesh_config_function": _mesh_config_simp, + "oft_order": 3, + "oft_maxits": 50, + "oft_python_path": OFT_PATH, + # --- Profile I/O --- + "profile_reader": pfile_reader, + "profile_reader_kwargs": reader_kwargs, + "uncertainty_generator": pfile_uncertainty_gen, + "uncertainty_generator_kwargs": uncertainty_kwargs, + # --- Bouquet sampling --- + "n_equils": n_equils, + "n_ls": n_ls, + "t_ls": t_ls, + "j_ls": j_ls, + "psi_pad": pad_psi, + "isoflux_weight": isoflux_weight, + "jphi_uncertainty_gen": FractionalUncertainty(frac_jphi), + "keep_geqdsk": KEEP_GEQDSK, + # --- Optional coil regularisation --- + "target_currents": target_currents, + # --- reconstruct_equilibrium keyword overrides --- + "reconstruct_equilibrium_kwargs": { + "n_k": n_k, + "psi_bridge": psi_bridge, + "use_python_solve": use_python_solve, + }, + "generate_bouquet_kwargs": { + "l_i_tolerance": l_i_tolerance, + "psi_pad": pad_psi, + "constrain_sawteeth": constrain_sawteeth, + "jBS_scale_range": jBS_scale_range, + "recalculate_j_BS": recalculate_j_BS, + "use_python_solve": use_python_solve, + "jphi_baseline": jphi_baseline, + "coil_drift": coil_drift, + "homotopy_passes": homotopy_passes, + "inspec_F_max": inspec_F_max, + "inspec_VSC_max": inspec_VSC_max, + "vsc_soft_reg_weight": vsc_soft_reg_weight, + "p_thresh": p_thresh, + } +} + +load_files_obj = load_profile_obj(config) + +# ============================================================================ +# 4. Pre-flight checks +# ============================================================================ + +def _compute(): + from bouquet.parallel import _get_num_cpus as _parallel_get_num_cpus + n_cpus, _ = _parallel_get_num_cpus() + if n_cpus < 2: + print(f'Warning: only {n_cpus} CPU core(s) available. ' + 'Parallel run requires at least 2 cores.') + sys.exit(0) + print(f'Running parallel run with {n_cpus} CPU core(s) available.') + + # Clean output directory for a fresh run + if os.path.exists(OUTPUT_DIR) and remake_dir: + import shutil as _shutil + _shutil.rmtree(OUTPUT_DIR) + os.makedirs(OUTPUT_DIR, exist_ok=True) + + print('\nChecking required input files...') + missing = [f for f in geqdsks + pfiles + [MESH_FILE] if not os.path.exists(f)] + if missing: + print('ERROR: the following files were not found:') + for f in missing: + print(f' {f}') + print('\nAdjust OFT_PATH / MESH_FILE and retry.') + sys.exit(1) + print(f' All {len(geqdsks)} geqdsk + {len(pfiles)} p-file pairs found.') + print(f' Mesh: {MESH_FILE}') + +# ============================================================================ +# 5. Compute in parallel +# ============================================================================ + print(f'\nLaunching parallel run into: {OUTPUT_DIR}') + print(f' {len(geqdsks)} equilibria, {n_equils} perturbed samples each') + print() + + from bouquet.parallel import parallel_runner + errors, outputs = parallel_runner( + list(zip(geqdsks, pfiles)), + load_files_obj, + bouquet_method, + OUTPUT_DIR, + use_logical_cpus=use_logical_cpus, + verbose=verbose, + ) + +# ============================================================================ +# 6. Error report +# ============================================================================ + if errors: + print(f'\nWARNING: {len(errors)} run(s) failed:') + for idx, tb in errors.items(): + print(f' [{idx}] {geqdsks[idx]}') + print(f' {tb.splitlines()[-1]}') + else: + print(f'\nAll {len(geqdsks)} runs completed successfully.') + return errors + +if __name__ == '__main__': + + if PLOT_ONLY: + errors = {} + print(f'\nPLOT_ONLY mode: skipping compute, loading results from:\n {OUTPUT_DIR}') + else: + errors = _compute() + + # ============================================================================ + # 7. Visualize results + # ============================================================================ + # parallel_runner saves one HDF5 database per run, named: + # {HEADER}_idx{idx}.h5 + # located in worker subdirectories under OUTPUT_DIR. + + import glob + import pickle as pkl + from collections import defaultdict + from bouquet import read_geqdsk, plot_bouquet, plot_geqdsk_bouquet, plot_pfile_bouquet + + # Collect all HDF5 databases + h5_files = sorted(glob.glob(os.path.join(OUTPUT_DIR, '**', f'{HEADER}_idx*.h5'), recursive=True)) + + if not h5_files: + print('\nNo HDF5 result files found; skipping plots.') + sys.exit(0 if not errors else 1) + + print(f'\n{"="*60}') + print(f'Visualizing results ({len(h5_files)} equilibrium/database file(s))') + print('=' * 60) + + # Baseline g-file shapes + print('\nPlotting baseline g-files...') + plot_geqdsk_bouquet(geqdsks, x_coord='rho') + out = os.path.join(OUTPUT_DIR, f'{HEADER}_baseline_geqdsk.png') + plt.savefig(out, dpi=150, bbox_inches='tight') + plt.close() + print(f' Saved: {out}') + + # Baseline p-files + print('\nPlotting baseline p-files...') + eq_ref = read_geqdsk(geqdsks[0]) + plot_pfile_bouquet(pfiles[0], x_coord='rho', eq=eq_ref) + out = os.path.join(OUTPUT_DIR, f'{HEADER}_baseline_pfile.png') + plt.savefig(out, dpi=150, bbox_inches='tight') + plt.close() + print(f' Saved: {out}') + + # Per-database plots (one per run/equilibrium) + for i, h5_path in enumerate(h5_files): + tag = f'idx{i}' + + print(f'\nPlotting bouquet — {tag}...') + plot_bouquet(h5_path, mode='all') + out = os.path.join(OUTPUT_DIR, f'{HEADER}_bouquet_{tag}.png') + plt.savefig(out, dpi=150, bbox_inches='tight') + plt.close() + print(f' Saved: {out}') + + print(f'Plotting perturbed g-files — {tag}...') + plot_geqdsk_bouquet(h5path=h5_path, x_coord='rho') + out = os.path.join(OUTPUT_DIR, f'{HEADER}_perturbed_geqdsk_{tag}.png') + plt.savefig(out, dpi=150, bbox_inches='tight') + plt.close() + print(f' Saved: {out}') + + print(f'Plotting perturbed p-files — {tag}...') + plot_pfile_bouquet(h5path=h5_path, x_coord='psi_N') + out = os.path.join(OUTPUT_DIR, f'{HEADER}_perturbed_pfile_{tag}.png') + plt.savefig(out, dpi=150, bbox_inches='tight') + plt.close() + print(f' Saved: {out}') 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