diff --git a/test/test_calculate_q_profile_from_psi.py b/test/test_calculate_q_profile_from_psi.py new file mode 100644 index 00000000..86d579e3 --- /dev/null +++ b/test/test_calculate_q_profile_from_psi.py @@ -0,0 +1,507 @@ +"""Tests for calculating the safety factor profile q(psi) from psi map and F(psi). + +Verifies Issue #647: +- Benchmark against analytic Solov'ev on-axis limit q_0 = |F_0| / (R_0 * sqrt(psi_RR * psi_ZZ)) +- Benchmark against high-precision 1D adaptive contour quadrature across multiple flux surfaces +- Convergence with grid resolution +- COCOS invariance and orientation sign compliance +- Input flexibility: callable, scalar, array, tuple, and return_details +- Verification that flux_surface_quantities returns 'q' +""" + +from __future__ import annotations + +import numpy as np +import pytest +from scipy.integrate import quad +from scipy.optimize import root_scalar + +from vaft.data.equilibrium import SolovevConstraint +from vaft.formula.equilibrium import q_from_flux_surface_averages +from vaft.process.equilibrium import ( + calculate_q_profile_from_psi, + evaluate_solovev, + flux_surface_quantities, + solve_solovev_constraints, +) + +R0, MINOR, ELONGATION = 1.0, 0.30, 1.6 + + +@pytest.fixture(scope="module") +def solovev_model(): + theta = np.linspace(0, 2 * np.pi, 9, endpoint=False) + solved = solve_solovev_constraints( + [ + SolovevConstraint( + R0 + MINOR * np.cos(t), ELONGATION * MINOR * np.sin(t), "psi", 0.0 + ) + for t in theta + ], + pprime=-1.0e4, + ffprime=0.05, + rref=R0, + psi_boundary=0.0, + ) + assert solved.rank == 5 + return solved + + +def _evaluate_solovev_grid(model, size): + r = np.linspace(R0 - 1.25 * MINOR, R0 + 1.25 * MINOR, size) + z = np.linspace(-1.25 * ELONGATION * MINOR, 1.25 * ELONGATION * MINOR, size) + grid_r, grid_z = np.meshgrid(r, z, indexing="ij") + ev = evaluate_solovev(model, grid_r, grid_z) + psi = np.asarray(ev["psi"], float) + return r, z, psi + + +def _solovev_analytic_q0(model, f0: float) -> float: + """Exact analytic on-axis safety factor from local Hessian.""" + eps = 1e-4 + ev_axis = evaluate_solovev(model, R0, 0.0) + ev_r_p = evaluate_solovev(model, R0 + eps, 0.0) + ev_r_m = evaluate_solovev(model, R0 - eps, 0.0) + ev_z_p = evaluate_solovev(model, R0, eps) + ev_z_m = evaluate_solovev(model, R0, -eps) + + psi_rr = (ev_r_p["psi"] + ev_r_m["psi"] - 2.0 * ev_axis["psi"]) / (eps**2) + psi_zz = (ev_z_p["psi"] + ev_z_m["psi"] - 2.0 * ev_axis["psi"]) / (eps**2) + + return abs(float(f0)) / (R0 * np.sqrt(abs(float(psi_rr) * float(psi_zz)))) + + +def _solovev_quadrature_q(model, psi_target: float, f0: float) -> float: + """High-precision reference q via 1D adaptive contour quadrature.""" + + def get_contour_point(t): + def obj(rad): + return ( + float( + evaluate_solovev( + model, + R0 + rad * np.cos(t), + ELONGATION * rad * np.sin(t), + )["psi"] + ) + - psi_target + ) + + sol = root_scalar(obj, bracket=[0.0, MINOR * 1.5]) + rad_val = sol.root + return R0 + rad_val * np.cos(t), ELONGATION * rad_val * np.sin(t) + + h = 1e-5 + + def integrand(t): + r_t, z_t = get_contour_point(t) + r_p, z_p = get_contour_point(t + h) + r_m, z_m = get_contour_point(t - h) + dr_dt = (r_p - r_m) / (2.0 * h) + dz_dt = (z_p - z_m) / (2.0 * h) + dl_dt = np.hypot(dr_dt, dz_dt) + ev = evaluate_solovev(model, r_t, z_t) + grad_psi = np.hypot(float(ev["dpsi_dr"]), float(ev["dpsi_dz"])) + return dl_dt / (r_t * grad_psi) + + integral = quad(integrand, 0.0, 2.0 * np.pi, epsabs=1e-8, epsrel=1e-8)[0] + return abs(float(f0)) / (2.0 * np.pi) * integral + + +def test_q_from_flux_surface_averages_pure_formula(): + """Verify formula layer units and COCOS factor.""" + # COCOS 1 (Wb/rad): factor is 1 / (2*pi)^2 + q_rad = q_from_flux_surface_averages(1.0, 1.0, 1.0, cocos=1) + assert np.isclose(q_rad, 1.0 / (4.0 * np.pi**2)) + + # COCOS 11 (full Wb): factor is 1 / (2*pi) + q_wb = q_from_flux_surface_averages(1.0, 1.0, 1.0, cocos=11) + assert np.isclose(q_wb, 1.0 / (2.0 * np.pi)) + + # Sign handling + q_neg = q_from_flux_surface_averages(1.0, 1.0, 1.0, cocos=1, sigma_ip=-1, sigma_b0=1) + assert q_neg < 0 + + +def test_analytic_solovev_axis_limit(solovev_model): + """Calculated q_0 matches the analytic Hessian formula on Solov'ev.""" + f0 = 1.2 + r, z, psi = _evaluate_solovev_grid(solovev_model, 129) + psi_axis = float(psi.max()) + psi_edge = 0.0 + + q_calc = calculate_q_profile_from_psi( + psi, + r, + z, + f0, + psi_axis=psi_axis, + psi_boundary=psi_edge, + cocos=1, + axis_rz=(R0, 0.0), + ) + + q0_expected = _solovev_analytic_q0(solovev_model, f0) + rel_error = abs(q_calc[0] - q0_expected) / q0_expected + + # On a 129x129 grid, extrapolated q0 matches within 0.2% + assert rel_error < 2.0e-3 + + +def test_high_precision_quadrature_benchmark(solovev_model): + """Calculated q(psi) matches adaptive 1D contour quadrature across surfaces.""" + f0 = 1.0 + r, z, psi = _evaluate_solovev_grid(solovev_model, 129) + psi_axis = float(psi.max()) + psi_edge = 0.0 + + levels = [0.2, 0.4, 0.6, 0.8] + q_calc = calculate_q_profile_from_psi( + psi, + r, + z, + f0, + psi_axis=psi_axis, + psi_boundary=psi_edge, + levels_norm=levels, + cocos=1, + axis_rz=(R0, 0.0), + ) + + for idx, lvl in enumerate(levels): + psi_target = psi_axis + lvl * (psi_edge - psi_axis) + q_ref = _solovev_quadrature_q(solovev_model, psi_target, f0) + rel_err = abs(q_calc[idx] - q_ref) / q_ref + assert rel_err < 5.0e-4, f"Level {lvl} rel_err={rel_err} exceeds 5e-4" + + +def test_mesh_convergence(solovev_model): + """Refining grid resolution from 65 to 129 to 257 improves error monotonically.""" + f0 = 1.0 + psi_axis = float(evaluate_solovev(solovev_model, R0, 0.0)["psi"]) + psi_edge = 0.0 + lvl = 0.5 + psi_target = psi_axis + lvl * (psi_edge - psi_axis) + q_ref = _solovev_quadrature_q(solovev_model, psi_target, f0) + + errors = [] + for size in (65, 129, 257): + r, z, psi = _evaluate_solovev_grid(solovev_model, size) + q_val = calculate_q_profile_from_psi( + psi, + r, + z, + f0, + psi_axis=psi_axis, + psi_boundary=psi_edge, + levels_norm=[lvl], + cocos=1, + axis_rz=(R0, 0.0), + )[0] + errors.append(abs(q_val - q_ref) / q_ref) + + assert errors[0] > errors[1] > errors[2] + assert errors[2] < 5.0e-5 + + +def test_cocos_invariance(solovev_model): + """The physical safety factor profile is invariant whether input is COCOS 1 or COCOS 11.""" + f0 = 1.0 + r, z, psi_rad = _evaluate_solovev_grid(solovev_model, 65) + psi_axis_rad = float(psi_rad.max()) + psi_edge_rad = 0.0 + + levels = np.linspace(0.0, 1.0, 21) + + # COCOS 1: Wb/rad + q_cocos1 = calculate_q_profile_from_psi( + psi_rad, + r, + z, + f0, + psi_axis=psi_axis_rad, + psi_boundary=psi_edge_rad, + levels_norm=levels, + cocos=1, + axis_rz=(R0, 0.0), + ) + + # COCOS 11: full Weber (multiplied by 2*pi) + psi_wb = psi_rad * (2.0 * np.pi) + psi_axis_wb = psi_axis_rad * (2.0 * np.pi) + psi_edge_wb = psi_edge_rad * (2.0 * np.pi) + + q_cocos11 = calculate_q_profile_from_psi( + psi_wb, + r, + z, + f0, + psi_axis=psi_axis_wb, + psi_boundary=psi_edge_wb, + levels_norm=levels, + cocos=11, + axis_rz=(R0, 0.0), + ) + + np.testing.assert_allclose(q_cocos1, q_cocos11, rtol=1e-12) + + +def test_cocos_signs(solovev_model): + """Safety factor sign follows Sauter Eq. 23 under COCOS conventions.""" + r, z, psi = _evaluate_solovev_grid(solovev_model, 65) + psi_axis = float(psi.max()) + psi_edge = 0.0 + + # Normal orientations: Ip > 0, B0 > 0 -> q > 0 + q_pos = calculate_q_profile_from_psi( + psi, + r, + z, + 1.0, + psi_axis=psi_axis, + psi_boundary=psi_edge, + levels_norm=[0.5], + cocos=11, + sigma_ip=1, + sigma_b0=1, + ) + assert q_pos[0] > 0 + + # Reversed current: Ip < 0, B0 > 0 -> q < 0 + q_neg_ip = calculate_q_profile_from_psi( + psi, + r, + z, + 1.0, + psi_axis=psi_axis, + psi_boundary=psi_edge, + levels_norm=[0.5], + cocos=11, + sigma_ip=-1, + sigma_b0=1, + ) + assert q_neg_ip[0] < 0 + + # Both reversed: Ip < 0, B0 < 0 -> q > 0 + q_both_neg = calculate_q_profile_from_psi( + psi, + r, + z, + -1.0, + psi_axis=psi_axis, + psi_boundary=psi_edge, + levels_norm=[0.5], + cocos=11, + sigma_ip=-1, + sigma_b0=-1, + ) + assert q_both_neg[0] > 0 + + +def test_input_flexibility_and_details(solovev_model): + """Tests callable f, tuple (psi, f), and return_details=True.""" + r, z, psi = _evaluate_solovev_grid(solovev_model, 65) + psi_axis = float(psi.max()) + psi_edge = 0.0 + levels = np.linspace(0.0, 1.0, 21) + + # 1. Scalar F + q_scalar = calculate_q_profile_from_psi( + psi, + r, + z, + 2.0, + psi_axis=psi_axis, + psi_boundary=psi_edge, + levels_norm=levels, + cocos=1, + ) + + # 2. Callable F(psi) + q_callable = calculate_q_profile_from_psi( + psi, + r, + z, + lambda p: 2.0, + psi_axis=psi_axis, + psi_boundary=psi_edge, + levels_norm=levels, + cocos=1, + ) + np.testing.assert_allclose(q_scalar, q_callable, rtol=1e-12) + + # 3. Tuple (psi_f, f_values) + psi_samples = np.linspace(min(psi_axis, psi_edge), max(psi_axis, psi_edge), 10) + q_tuple = calculate_q_profile_from_psi( + psi, + r, + z, + (psi_samples, np.full_like(psi_samples, 2.0)), + psi_axis=psi_axis, + psi_boundary=psi_edge, + levels_norm=levels, + cocos=1, + ) + np.testing.assert_allclose(q_scalar, q_tuple, rtol=1e-12) + + # 4. return_details=True + details = calculate_q_profile_from_psi( + psi, + r, + z, + 2.0, + psi_axis=psi_axis, + psi_boundary=psi_edge, + levels_norm=levels, + cocos=1, + return_details=True, + ) + assert isinstance(details, dict) + assert "q" in details + assert "levels_norm" in details + assert "psi_levels" in details + assert "q_axis" in details + assert "q_95" in details + assert "surfaces" in details + assert np.isclose(details["q_axis"], details["q"][0]) + + +def test_flux_surface_quantities_populates_q(solovev_model): + """flux_surface_quantities populates 'q' when f_profile is given, and NaN when None.""" + r, z, psi = _evaluate_solovev_grid(solovev_model, 65) + psi_axis = float(psi.max()) + psi_edge = 0.0 + levels = np.linspace(0.0, 1.0, 11) + + # Without f_profile -> q is all NaN + surfaces_no_f = flux_surface_quantities( + psi, + r, + z, + psi_axis, + psi_edge, + levels, + f_profile=None, + ) + assert "q" in surfaces_no_f + assert np.isnan(surfaces_no_f["q"]).all() + + # With f_profile -> q is finite + f_vals = np.ones_like(levels) * 1.5 + surfaces_with_f = flux_surface_quantities( + psi, + r, + z, + psi_axis, + psi_edge, + levels, + f_profile=f_vals, + ) + assert "q" in surfaces_with_f + assert np.isfinite(surfaces_with_f["q"]).all() + assert (surfaces_with_f["q"] > 0).all() + + +def test_f_profile_none_raises(solovev_model): + """f_profile=None raises ValueError in calculate_q_profile_from_psi.""" + r, z, psi = _evaluate_solovev_grid(solovev_model, 65) + psi_axis = float(psi.max()) + psi_edge = 0.0 + with pytest.raises(ValueError, match="f_profile must be provided"): + calculate_q_profile_from_psi( + psi, + r, + z, + None, + psi_axis=psi_axis, + psi_boundary=psi_edge, + levels_norm=[0.5], + ) + + +def test_negative_f_profile_infers_sigma_b0(solovev_model): + """Negative f_profile with default sigma_b0=1 infers negative toroidal field.""" + r, z, psi = _evaluate_solovev_grid(solovev_model, 65) + psi_axis = float(psi.max()) + psi_edge = 0.0 + + q_pos = calculate_q_profile_from_psi( + psi, + r, + z, + 1.5, + psi_axis=psi_axis, + psi_boundary=psi_edge, + levels_norm=[0.5], + sigma_ip=1, + sigma_b0=1, + cocos=1, + ) + assert q_pos[0] > 0 + + # Negative F should produce negative q when sigma_ip=1 in COCOS 1 + q_neg = calculate_q_profile_from_psi( + psi, + r, + z, + -1.5, + psi_axis=psi_axis, + psi_boundary=psi_edge, + levels_norm=[0.5], + sigma_ip=1, + sigma_b0=1, + cocos=1, + ) + assert q_neg[0] < 0 + assert np.isclose(q_neg[0], -q_pos[0]) + + +def test_unsorted_and_descending_levels_return_details(solovev_model): + """return_details correctly handles unsorted and descending level inputs.""" + r, z, psi = _evaluate_solovev_grid(solovev_model, 65) + psi_axis = float(psi.max()) + psi_edge = 0.0 + + # Levels descending and omitting 0.0 + desc_levels = [0.95, 0.7, 0.5, 0.2] + res = calculate_q_profile_from_psi( + psi, + r, + z, + 1.0, + psi_axis=psi_axis, + psi_boundary=psi_edge, + levels_norm=desc_levels, + return_details=True, + ) + assert np.isfinite(res["q_axis"]) + assert np.isfinite(res["q_95"]) + assert np.isclose(res["q_95"], res["q"][0]) + # Axis extrapolation should reasonably match q on smallest level + assert res["q_axis"] < res["q"][-1] + + +def test_q0_stability_with_dense_levels(solovev_model): + """Polynomial fit preserves q0 accuracy across varying level densities.""" + r, z, psi = _evaluate_solovev_grid(solovev_model, 129) + psi_axis = float(psi.max()) + psi_edge = 0.0 + q0_ref = _solovev_analytic_q0(solovev_model, 1.0) + + # Test with standard N=33 levels and dense N=129 levels + for n_levels in (33, 129): + levels = np.linspace(0.0, 0.9, n_levels) + q = calculate_q_profile_from_psi( + psi, + r, + z, + 1.0, + psi_axis=psi_axis, + psi_boundary=psi_edge, + levels_norm=levels, + cocos=1, + axis_rz=(R0, 0.0), + ) + rel_err = abs(q[0] - q0_ref) / q0_ref + assert rel_err < 0.015, f"n_levels={n_levels} q0 rel_err={rel_err} exceeds 1.5%" + diff --git a/test/test_formula_catalog.py b/test/test_formula_catalog.py index 4825c1ad..b389744b 100644 --- a/test/test_formula_catalog.py +++ b/test/test_formula_catalog.py @@ -57,7 +57,7 @@ def test_the_catalog_counts_the_known_public_surface(): assert counts == { "constants": 0, "utils": 10, - "equilibrium": 106, + "equilibrium": 107, "stability": 19, "green": 16, "atomic": 3, diff --git a/test/test_process_docstrings.py b/test/test_process_docstrings.py index ed22932a..8c21c9f7 100644 --- a/test/test_process_docstrings.py +++ b/test/test_process_docstrings.py @@ -106,6 +106,7 @@ "sxr_band_signals", "sxr_electron_temperature", # equilibrium (#419) + "calculate_q_profile_from_psi", "calculate_reconstructed_diamagnetic_flux", "convert_cocos", "derive_global_descriptors", @@ -262,6 +263,7 @@ "as_equilibrium", "calculate_average_boundary_poloidal_field", "calculate_diamagnetism", + "calculate_q_profile_from_psi", "calculate_reconstructed_diamagnetic_flux", "check_equilibrium_requirements", "computed_diamagnetism_from_phi", diff --git a/vaft/formula/equilibrium.py b/vaft/formula/equilibrium.py index 4bfd3c42..c6a96b5b 100644 --- a/vaft/formula/equilibrium.py +++ b/vaft/formula/equilibrium.py @@ -332,6 +332,86 @@ def rho_tor_from_phi(phi: Union[np.ndarray, float], # Safety Factor Calculations # ------------------------------------------------------------------ +def q_from_flux_surface_averages( + gm1: Union[np.ndarray, float], + dvolume_dpsi: Union[np.ndarray, float], + f: Union[np.ndarray, float], + *, + cocos: int | None = None, + psi_per_radian: bool | None = None, + sigma_ip: int = 1, + sigma_b0: int = 1, +) -> Union[np.ndarray, float]: + r"""Safety factor $q$ from flux-surface geometric averages and poloidal current. + + $$q(\psi) = \sigma \cdot \frac{F(\psi)}{2\pi} \oint \frac{dl_p}{R^2 B_p} + = \sigma \cdot (2\pi)^{e_{B_p} - 2} \cdot F(\psi) \cdot \left\langle \frac{1}{R^2} \right\rangle \cdot \frac{dV}{d|\psi|}$$ + + Parameters + ---------- + gm1 : float or np.ndarray + Geometric flux-surface average $\langle 1/R^2 \rangle$ [m^-2]. + dvolume_dpsi : float or np.ndarray + Differential volume element $dV/d|\psi|$ [m^3/(Wb/rad) if $e_{B_p}=0$ or m^3/Wb if $e_{B_p}=1$]. + f : float or np.ndarray + Poloidal current function $F = R B_\varphi$ on the flux surface [T m]. + cocos : int or None, optional + COCOS coordinate convention index (1-8, 11-18) [-]. + psi_per_radian : bool or None, optional + Storage family of the flux when ``cocos`` is None [bool]. + ``False`` assumes full-weber flux ($e_{B_p} = 1$); ``True`` and ``None`` + keep the per-radian assumption ($e_{B_p} = 0$). + sigma_ip : int, optional + Sign of plasma current (+1 or -1) in the equilibrium coordinate system [-]. + sigma_b0 : int, optional + Sign of toroidal field (+1 or -1) in the equilibrium coordinate system [-]. + + Returns + ------- + float or np.ndarray + Safety factor profile $q$ on the corresponding flux surfaces [-]. + + Convention + ---------- + In Sauter and Medvedev (2013), $B_p = |\nabla\psi| / (R (2\pi)^{e_{B_p}})$. + The safety factor contour integral is: + $$q = \frac{F}{2\pi} \oint \frac{dl_p}{R^2 B_p} = (2\pi)^{e_{B_p}-1} \frac{F}{2\pi} \oint \frac{dl_p}{R |\nabla\psi|}$$ + Since $dV/d\psi = 2\pi \oint \frac{R dl_p}{|\nabla\psi|}$ and + $\langle 1/R^2 \rangle = \frac{\oint dl_p / (R |\nabla\psi|)}{\oint R dl_p / |\nabla\psi|}$, + this gives: + $$q = \sigma \cdot (2\pi)^{e_{B_p}-2} \cdot |F| \cdot \langle 1/R^2 \rangle \cdot \frac{dV}{d|\psi|}$$ + For $e_{B_p} = 0$ (COCOS 1-8, Wb/rad), $(2\pi)^{0-2} = 1/(4\pi^2)$. + For $e_{B_p} = 1$ (COCOS 11-18, full Wb), $(2\pi)^{1-2} = 1/(2\pi)$. + The sign $\sigma$ is determined by Sauter Eq. 23: $\sigma_q = \sigma_{Ip}\sigma_{B0}\sigma_{\rho\theta\varphi}$. + + References + ---------- + .. [1] O. Sauter and S. Yu. Medvedev, Comput. Phys. Commun. 184 (2013) 293. + """ + gm1_arr = np.asarray(gm1, dtype=float) + dv_arr = np.asarray(dvolume_dpsi, dtype=float) + f_arr = np.asarray(f, dtype=float) + + if cocos is not None: + from vaft.data.cocos import cocos_spec + + spec = cocos_spec(cocos) + exp_bp = spec.exp_bp + target_sign = spec.expected_sign("q", sigma_ip=sigma_ip, sigma_b0=sigma_b0) + else: + exp_bp = 0 if (psi_per_radian is None or psi_per_radian) else 1 + target_sign = 1 if (sigma_ip * sigma_b0) >= 0 else -1 + + factor = (2.0 * np.pi) ** (exp_bp - 2) + q_mag = factor * np.abs(f_arr) * np.abs(gm1_arr) * np.abs(dv_arr) + q = target_sign * q_mag + if np.ndim(gm1) == 0 and np.ndim(dvolume_dpsi) == 0 and np.ndim(f) == 0: + return float(q) + return q + + + + def q_from_phi(psi: np.ndarray, phi: np.ndarray) -> np.ndarray: r"""Safety factor $q$ as the flux derivative $d\Phi/d\psi$. diff --git a/vaft/process/equilibrium.py b/vaft/process/equilibrium.py index 5ed57e80..085d5365 100644 --- a/vaft/process/equilibrium.py +++ b/vaft/process/equilibrium.py @@ -92,6 +92,7 @@ "MIN_FLUX_SURFACE_POINTS", "calculate_average_boundary_poloidal_field", "calculate_diamagnetism", + "calculate_q_profile_from_psi", "calculate_reconstructed_diamagnetic_flux", "computed_diamagnetism_from_phi", "contour_shape_parameters", @@ -2430,6 +2431,7 @@ def _enclosing_segment( "b_field_min", "bp_dl", "length_pol", + "q", ) @@ -2465,7 +2467,7 @@ def flux_surface_quantities( [-]. f_profile : array_like, optional Poloidal current function ``F = R B_phi`` on the same levels. Required for - the mean square field [T m]. + the mean square field and safety factor [T m]. axis_rz : tuple of float, optional Magnetic axis position, used to pick the confined segment when a level traces several [m]. @@ -2482,8 +2484,8 @@ def flux_surface_quantities( *levels_norm*. Volume in cubic metres, areas in square metres, radii and ``length_pol`` in metres, ``bp_dl`` in tesla-metre, ``gm1`` in inverse square metres, ``gm5`` in tesla squared, ``gm8`` in metres, ``gm9`` in - inverse metres, the shape parameters dimensionless, and the two - derivatives per radian [-]. + inverse metres, the shape parameters and safety factor ``q`` dimensionless, + and the two derivatives per radian [-]. Convention ---------- @@ -2583,6 +2585,8 @@ def flux_surface_quantities( else {} ) + from vaft.formula.equilibrium import q_from_flux_surface_averages + for index, level in enumerate(levels): if level == 0.0: if axis_rz is not None: @@ -2650,28 +2654,294 @@ def flux_surface_quantities( out["gm5"][index] = float(np.sum(weight * b_mod**2) / total) out["b_field_max"][index] = float(np.max(b_mod)) out["b_field_min"][index] = float(np.min(b_mod)) + out["q"][index] = float( + q_from_flux_surface_averages( + out["gm1"][index], out["dvolume_dpsi"][index], f_values[index] + ) + ) + + if f_values is not None: + # q is quadratic in radius r ~ sqrt(psi_N), making it smooth and linear in + # psi_N near the magnetic axis. Extrapolate q to axis level (psi_N = 0) in psi_N. + # Use a low-degree polynomial fit over innermost resolved surfaces (psi_N <= 0.25) + # to filter out discrete polygon vertex jitter on tiny near-axis contours. + axis_mask = levels == 0.0 + if axis_mask.any() and not np.isfinite(out["q"][axis_mask]).all(): + finite_idx = np.where(np.isfinite(out["q"]))[0] + if len(finite_idx) >= 2: + sorted_finite = finite_idx[np.argsort(levels[finite_idx])] + fit_idx = [i for i in sorted_finite if levels[i] <= 0.25] + if len(fit_idx) < 3: + fit_idx = sorted_finite[: min(len(sorted_finite), 5)] + p_fit = levels[fit_idx] + q_fit = out["q"][fit_idx] + poly = np.polyfit(p_fit, q_fit, deg=1) + out["q"][axis_mask] = float(np.polyval(poly, 0.0)) + elif len(finite_idx) == 1: + out["q"][axis_mask] = out["q"][finite_idx[0]] # Gaps are filled against sqrt(psi_N), not psi_N: near the axis a flux # surface's linear size goes as sqrt(psi_N), so every quantity that vanishes # there is linear in sqrt and badly curved in psi_N. Interpolating a dropped # innermost level in psi_N underestimates `surface` by a third. - # - # `np.interp` requires an increasing `xp` and returns nonsense rather than - # raising when it does not get one, so the levels are sorted here instead of - # assumed: a psi profile stored boundary-first is a real input, and it - # corrupted only the quantities that happened to need a gap filled. + # q is instead smooth and linear in psi_N, so its gaps are filled against levels. coordinate = np.sqrt(np.clip(levels, 0.0, None)) order = np.argsort(coordinate, kind="stable") + order_linear = np.argsort(levels, kind="stable") for name, values in out.items(): missing = ~np.isfinite(values) if missing.any() and not missing.all(): - good_sorted = order[np.isfinite(values[order])] - values[missing] = np.interp( - coordinate[missing], coordinate[good_sorted], values[good_sorted] - ) + if name == "q": + good_sorted = order_linear[np.isfinite(values[order_linear])] + values[missing] = np.interp( + levels[missing], levels[good_sorted], values[good_sorted] + ) + else: + good_sorted = order[np.isfinite(values[order])] + values[missing] = np.interp( + coordinate[missing], coordinate[good_sorted], values[good_sorted] + ) return out +def calculate_q_profile_from_psi( + psi_grid: np.ndarray, + R: np.ndarray, + Z: np.ndarray, + f_profile: Any, + psi_axis: float | None = None, + psi_boundary: float | None = None, + levels_norm: Any = None, + *, + axis_rz: tuple[float, float] | None = None, + boundary: tuple[np.ndarray, np.ndarray] | None = None, + cocos: int = 11, + sigma_ip: int = 1, + sigma_b0: int = 1, + min_points: int = MIN_FLUX_SURFACE_POINTS, + return_details: bool = False, +) -> np.ndarray | dict[str, Any]: + """Calculate the safety factor profile q(psi) from a 2D poloidal flux map and F(psi). + + Parameters + ---------- + psi_grid : array_like + Poloidal flux on the grid, shaped ``(len(R), len(Z))``, in the convention + ``cocos`` specifies [Wb or Wb/rad]. + R : array_like + Major-radius grid axis [m]. + Z : array_like + Height grid axis [m]. + f_profile : float, array_like, callable, or tuple of array_like + Poloidal current function ``F = R B_phi``, as a scalar float, 1D array, + callable ``f(psi)``, or tuple ``(psi_f, f_values)`` [T m]. + psi_axis : float, optional + Poloidal flux at the magnetic axis, in the convention ``cocos`` specifies [Wb or Wb/rad]. + psi_boundary : float, optional + Poloidal flux at the plasma boundary/LCFS, in the convention ``cocos`` specifies [Wb or Wb/rad]. + levels_norm : sequence of float, optional + Normalized flux levels to evaluate on, 0 on axis and 1 at the boundary [-]. + axis_rz : tuple of float, optional + Magnetic axis coordinates ``(R_axis, Z_axis)`` [m]. + boundary : tuple of np.ndarray, optional + Boundary outline ``(R_bdry, Z_bdry)`` [m]. + cocos : int, optional + COCOS coordinate convention index (1-8, 11-18) describing the input + conventions and target sign for ``q`` [-]. + sigma_ip : int, optional + Sign of plasma current (+1 or -1) in the equilibrium coordinate system [-]. + sigma_b0 : int, optional + Sign of toroidal field (+1 or -1) in the equilibrium coordinate system [-]. + min_points : int, optional + Fewest contour vertices a level needs before it is treated as resolved [-]. + return_details : bool, optional + Whether to return diagnostic parameters and the full flux-surface + quantities alongside the safety factor array [-]. + + Returns + ------- + np.ndarray or dict of str to Any + If ``return_details`` is False, returns a 1D array of safety factor values ``q`` + on ``levels_norm`` [-]. + If ``return_details`` is True, returns a dict with keys ``"q"``, ``"levels_norm"``, + ``"psi_levels"``, ``"q_axis"``, ``"q_95"``, and ``"surfaces"`` [-]. + + Raises + ------ + ValueError + The flux map is not shaped to ``(len(R), len(Z))``, ``psi_axis`` or + ``psi_boundary`` cannot be determined or are equal, or ``f_profile`` + format is invalid. + + Convention + ---------- + In Sauter and Medvedev (2013), the safety factor is defined as: + $$q(\\psi) = \\frac{F(\\psi)}{2\\pi} \\oint \\frac{dl_p}{R^2 B_p} + = \\frac{F(\\psi)}{2\\pi} (2\\pi)^{e_{B_p}} \\oint \\frac{dl_p}{R |\\nabla\\psi|}$$ + Expressed in terms of the geometric flux-surface averages computed by + :func:`flux_surface_quantities` (which operates in Wb/rad, $e_{B_p} = 0$, $dV/d\\psi$ per radian): + $$q(\\psi) = \\sigma \\cdot \\frac{F(\\psi)}{(2\\pi)^2} \\left\\langle \\frac{1}{R^2} \\right\\rangle \\frac{dV}{d\\psi}$$ + where $\\langle 1/R^2 \\rangle$ is ``gm1``, $dV/d\\psi$ is ``dvolume_dpsi``, and $\\sigma$ is + the orientation sign from Sauter Eq. 23: $\\sigma_q = \\sigma_{Ip} \\sigma_{B0} \\sigma_{\\rho\\theta\\varphi}$. + When ``cocos`` indicates full Weber storage ($e_{B_p} = 1$, COCOS 11-18), the input fluxes are + divided by $2\\pi$ to enter the contour tracing engine, and the resulting profile carries the + exact COCOS-mandated sign. + + Processing steps + ---------------- + 1. Validate grid dimensions and identify the flux conversion factor from ``cocos``. + 2. Normalize or deduce ``psi_axis`` and ``psi_boundary``, converting to Wb/rad if necessary. + 3. Evaluate and interpolate ``f_profile`` onto the requested ``levels_norm``. + 4. Call :func:`flux_surface_quantities` to trace contours and compute geometric averages. + 5. Apply on-axis quadratic extrapolation in radius (linear in normalized flux $\\psi_N$) to + resolve $q_0$. + 6. Attach the COCOS sign factor $\\sigma_q$ and package the output. + + Defaults + -------- + ``levels_norm`` defaults to 65 points from 0 to 1 (numerical convenience). + ``cocos=11`` matches the IMAS standard data dictionary convention. + + Applicability + ------------- + Machine-independent. Works for any 2D tokamak poloidal flux map. + + Provenance + ---------- + .. [1] O. Sauter and S. Yu. Medvedev, Comput. Phys. Commun. 184 (2013) 293, + Eq. (17) and Table I for COCOS relations and safety factor definitions. + .. [2] J. Wesson, *Tokamaks*, 4th ed., Oxford University Press (2011), Sec. 3.4. + """ + from vaft.data.cocos import cocos_spec + + if f_profile is None: + raise ValueError("f_profile must be provided and cannot be None.") + + R_arr = np.asarray(R, dtype=float).reshape(-1) + Z_arr = np.asarray(Z, dtype=float).reshape(-1) + psi_arr = np.asarray(psi_grid, dtype=float) + if psi_arr.shape != (R_arr.size, Z_arr.size): + raise ValueError( + f"psi_grid shape {psi_arr.shape} must equal (len(R), len(Z)) = {(R_arr.size, Z_arr.size)}." + ) + + if levels_norm is None: + levels = np.linspace(0.0, 1.0, 65) + else: + levels = np.asarray(levels_norm, dtype=float).reshape(-1) + + spline = None + if psi_axis is None: + if axis_rz is not None: + spline = RectBivariateSpline(R_arr, Z_arr, psi_arr) + psi_axis = float(spline.ev(axis_rz[0], axis_rz[1])) + else: + edge_val = float( + np.mean( + [ + psi_arr[0, :], + psi_arr[-1, :], + psi_arr[:, 0], + psi_arr[:, -1], + ] + ) + ) + diff = psi_arr - edge_val + idx_max = np.unravel_index(np.argmax(np.abs(diff)), psi_arr.shape) + psi_axis = float(psi_arr[idx_max]) + if axis_rz is None: + axis_rz = (float(R_arr[idx_max[0]]), float(Z_arr[idx_max[1]])) + + if psi_boundary is None: + if boundary is not None: + if spline is None: + spline = RectBivariateSpline(R_arr, Z_arr, psi_arr) + r_b = np.asarray(boundary[0], dtype=float).reshape(-1) + z_b = np.asarray(boundary[1], dtype=float).reshape(-1) + psi_boundary = float(np.mean(spline.ev(r_b, z_b))) + else: + raise ValueError( + "psi_boundary could not be deduced; provide psi_boundary or boundary outline." + ) + + if psi_axis == psi_boundary: + raise ValueError("psi_axis and psi_boundary must differ to define normalized flux.") + + psi_levels = psi_axis + levels * (psi_boundary - psi_axis) + + if callable(f_profile): + f_values = np.asarray([float(f_profile(p)) for p in psi_levels], dtype=float) + elif isinstance(f_profile, (tuple, list)) and len(f_profile) == 2: + psi_f = np.asarray(f_profile[0], dtype=float).reshape(-1) + f_raw = np.asarray(f_profile[1], dtype=float).reshape(-1) + if psi_f.size != f_raw.size: + raise ValueError("f_profile tuple (psi, f) must have equal-length arrays.") + order = np.argsort(psi_f) + f_values = np.interp(psi_levels, psi_f[order], f_raw[order]) + elif np.ndim(f_profile) == 0: + f_values = np.full(levels.size, float(f_profile)) + else: + f_arr = np.asarray(f_profile, dtype=float).reshape(-1) + if f_arr.size != levels.size: + raise ValueError( + f"f_profile length ({f_arr.size}) must match levels_norm length ({levels.size})." + ) + f_values = f_arr + + spec = cocos_spec(cocos) + scale_to_wb_per_rad = (1.0 / (2.0 * np.pi)) if spec.exp_bp == 1 else 1.0 + + psi_grid_rad = psi_arr * scale_to_wb_per_rad + psi_axis_rad = float(psi_axis) * scale_to_wb_per_rad + psi_boundary_rad = float(psi_boundary) * scale_to_wb_per_rad + + surfaces = flux_surface_quantities( + psi_grid=psi_grid_rad, + R=R_arr, + Z=Z_arr, + psi_axis=psi_axis_rad, + psi_boundary=psi_boundary_rad, + levels_norm=levels, + f_profile=f_values, + axis_rz=axis_rz, + boundary=boundary, + min_points=min_points, + ) + + f_mean = float(np.nanmean(f_values)) if f_values.size else 0.0 + if sigma_b0 == 1 and f_mean < -1e-12: + effective_sigma_b0 = -1 + else: + effective_sigma_b0 = sigma_b0 + + target_sign = spec.expected_sign("q", sigma_ip=sigma_ip, sigma_b0=effective_sigma_b0) + q_final = target_sign * np.abs(surfaces["q"]) + surfaces["q"] = q_final + + if return_details: + order_lvl = np.argsort(levels) + lvl_sorted = levels[order_lvl] + q_sorted = q_final[order_lvl] + if (levels == 0.0).any(): + q_axis = float(q_final[levels == 0.0][0]) + elif len(lvl_sorted) >= 2: + p1, p2 = lvl_sorted[0], lvl_sorted[1] + q1, q2 = q_sorted[0], q_sorted[1] + q_axis = float(q1 - (q2 - q1) / (p2 - p1) * p1) if p2 != p1 else float(q1) + else: + q_axis = float(q_sorted[0]) + q_95 = float(np.interp(0.95, lvl_sorted, q_sorted)) + return { + "q": q_final, + "levels_norm": levels, + "psi_levels": psi_levels, + "q_axis": q_axis, + "q_95": q_95, + "surfaces": surfaces, + } + return q_final + + def equilibrium_field_on_grid( R_grid_1d: np.ndarray, Z_grid_1d: np.ndarray, @@ -3098,6 +3368,7 @@ def _cumulative_arc_length(phi: list[float], R: list[float], Z: list[float]) -> except ImportError: # direct ``spec_from_file_location`` loading from vaft.process._equilibrium_parametric import * # noqa: E402,F401,F403 + @dataclass(frozen=True) class ParallelCurrentResult: """Parallel current density derived from an enclosed toroidal current.