diff --git a/notebooks/analytic_solovev_equilibrium.ipynb b/notebooks/analytic_solovev_equilibrium.ipynb index 313ca0ac..c248cc2f 100644 --- a/notebooks/analytic_solovev_equilibrium.ipynb +++ b/notebooks/analytic_solovev_equilibrium.ipynb @@ -86,7 +86,7 @@ "# 2*pi to recover the per-radian fields evaluate_solovev works with.\n", "dpsi_dr = np.gradient(equilibrium.psi, r, axis=0, edge_order=2) / (2 * np.pi)\n", "dpsi_dz = np.gradient(equilibrium.psi, z, axis=1, edge_order=2) / (2 * np.pi)\n", - "br_grid = -dpsi_dz/rm; bz_grid = dpsi_dr/rm\n", + "br_grid = dpsi_dz/rm; bz_grid = -dpsi_dr/rm\n", "np.testing.assert_allclose(br_grid[3:-3, 3:-3], analytic[\"b_r\"][3:-3, 3:-3], rtol=1e-2, atol=3e-2)\n", "np.testing.assert_allclose(bz_grid[3:-3, 3:-3], analytic[\"b_z\"][3:-3, 3:-3], rtol=1e-2, atol=3e-2)\n", "print(\"Analytic and discretized poloidal fields agree within the grid tolerance.\")\n" diff --git a/test/test_parametric_equilibrium.py b/test/test_parametric_equilibrium.py index 1a25f7fe..d475ea15 100644 --- a/test/test_parametric_equilibrium.py +++ b/test/test_parametric_equilibrium.py @@ -321,17 +321,111 @@ def test_solovev_export_rejects_an_open_boundary_contour(): solovev_to_equilibrium(model, r, z, magnetic_axis=(1.0, 0.0)) +def test_evaluate_solovev_orientation_families_and_convention_alias(): + """evaluate_solovev poloidal field orientation matches COCOS definitions.""" + from vaft.data.cocos import COCOS_INDICES, cocos_spec + from vaft.process.equilibrium import evaluate_solovev + + model = SolovevEquilibrium(np.array([0.03, -0.02, 0.015, 0.004, -0.001]), -1200.0, 0.08, 1.0) + r = np.linspace(0.7, 1.3, 21); z = np.linspace(-0.4, 0.4, 19) + rm, zm = np.meshgrid(r, z, indexing="ij") + + # Default is COCOS 11 (sigma = +1) + default_vals = evaluate_solovev(model, rm, zm) + cocos11_vals = evaluate_solovev(model, rm, zm, cocos=11) + np.testing.assert_allclose(default_vals["b_r"], cocos11_vals["b_r"]) + np.testing.assert_allclose(default_vals["b_z"], cocos11_vals["b_z"]) + + # Fallback cocos=None matches historical orientation (sigma = -1) + hist_vals = evaluate_solovev(model, rm, zm, cocos=None) + np.testing.assert_allclose(default_vals["b_r"], -hist_vals["b_r"]) + np.testing.assert_allclose(default_vals["b_z"], -hist_vals["b_z"]) + + # Alias convention=... + alias_vals = evaluate_solovev(model, rm, zm, convention=11) + np.testing.assert_allclose(default_vals["b_r"], alias_vals["b_r"]) + + # Conflicting arguments + with pytest.raises(ValueError, match="conflicting"): + evaluate_solovev(model, rm, zm, cocos=2, convention=12) + + # Invalid COCOS indices + for invalid in (0, 9, 10, 19, -1): + with pytest.raises(ValueError, match="valid COCOS"): + evaluate_solovev(model, rm, zm, cocos=invalid) + + # Check all 16 COCOS indices match their expected orientation sign + for index in COCOS_INDICES: + vals = evaluate_solovev(model, rm, zm, cocos=index) + spec = cocos_spec(index) + expected_sign = spec.sigma_rpz * spec.sigma_bp + if expected_sign > 0: + np.testing.assert_allclose(vals["b_r"], default_vals["b_r"]) + np.testing.assert_allclose(vals["b_z"], default_vals["b_z"]) + else: + np.testing.assert_allclose(vals["b_r"], hist_vals["b_r"]) + np.testing.assert_allclose(vals["b_z"], hist_vals["b_z"]) + + def test_solovev_export_honors_the_declared_psi_convention(): """The export must be self-consistent with its COCOS tag: full-weber conventions carry 2*pi*psi, and derived fields/descriptors agree with the per-radian export.""" - from vaft.process.equilibrium import solovev_to_equilibrium + from vaft.formula.constants import MU0 + from vaft.process.equilibrium import ( + evaluate_solovev, + solovev_to_equilibrium, + solve_solovev_constraints, + ) + from vaft.process._equilibrium_parametric import _grid_fields # c2 > 0 makes the axis a saddle in Z, so no magnetic axis exists. saddle = SolovevEquilibrium(np.array([0.0, -0.02, 0.015, 0.0, 0.0]), -1.0e5, 0.0, 1.0) with pytest.raises(ValueError, match="magnetic axis"): solovev_to_equilibrium(saddle, np.linspace(0.5, 1.5, 121), np.linspace(-0.7, 0.7, 121)) + # Invalid conventions + with pytest.raises(ValueError, match="convention must be a COCOS index"): + solovev_to_equilibrium(saddle, np.linspace(0.5, 1.5, 121), np.linspace(-0.7, 0.7, 121), convention=0) + with pytest.raises(ValueError, match="convention must be a COCOS index"): + solovev_to_equilibrium(saddle, np.linspace(0.5, 1.5, 121), np.linspace(-0.7, 0.7, 121), convention=9) + + # Well-formed equilibrium across conventions + R0, kappa, psi_b = 1.0, 1.4, 0.08 + pprime = -8.0 * ((1.0 + 1.0 / kappa**2) / 4.0) / MU0 + r_out = np.sqrt(R0**2 + 2 * np.sqrt(psi_b)) + r_in = np.sqrt(R0**2 - 2 * np.sqrt(psi_b)) + z_top = kappa * np.sqrt(psi_b) / R0 + constraints = [ + SolovevConstraint(r_out, 0.0, "psi", psi_b), + SolovevConstraint(r_in, 0.0, "psi", psi_b), + SolovevConstraint(R0, z_top, "psi", psi_b), + SolovevConstraint(R0, 0.0, "psi", 0.0), + SolovevConstraint(R0, 0.0, "dpsi_dr", 0.0), + ] + model = solve_solovev_constraints( + constraints, pprime=pprime, ffprime=0.0, rref=R0, psi_boundary=psi_b, f_boundary=1.4 + ) + r = np.linspace(0.5, 1.5, 151) + z = np.linspace(-0.7, 0.7, 151) + rm, zm = np.meshgrid(r, z, indexing="ij") + + # Fields on grid derived through _grid_fields must match evaluate_solovev in COCOS 11 + eq11 = solovev_to_equilibrium(model, r, z, convention=11) + _, _, br11, bz11, _ = _grid_fields(eq11) + analytic11 = evaluate_solovev(model, rm, zm, cocos=11) + np.testing.assert_allclose(br11[5:-5, 5:-5], analytic11["b_r"][5:-5, 5:-5], rtol=1e-3, atol=1e-3) + np.testing.assert_allclose(bz11[5:-5, 5:-5], analytic11["b_z"][5:-5, 5:-5], rtol=1e-3, atol=1e-3) + + # Check both orientation families (e.g. COCOS 1, 2, 3, 11, 12, 13) + for c in (1, 2, 3, 11, 12, 13): + eq_c = solovev_to_equilibrium(model, r, z, convention=c) + assert eq_c.convention.cocos == c + _, _, br_c, bz_c, bp_c = _grid_fields(eq_c) + # Physical field components in (R, Z) match the analytic fields for all conventions + np.testing.assert_allclose(br_c[5:-5, 5:-5], analytic11["b_r"][5:-5, 5:-5], rtol=1e-3, atol=1e-3) + np.testing.assert_allclose(bz_c[5:-5, 5:-5], analytic11["b_z"][5:-5, 5:-5], rtol=1e-3, atol=1e-3) + def test_d_r_sep_uses_the_midplane_and_reports_asymmetry(): balanced = derive_boundary_representation(_analytic_equilibrium("double")) diff --git a/vaft/process/_equilibrium_parametric.py b/vaft/process/_equilibrium_parametric.py index 16337a87..d687fc66 100644 --- a/vaft/process/_equilibrium_parametric.py +++ b/vaft/process/_equilibrium_parametric.py @@ -17,6 +17,7 @@ from scipy.optimize import least_squares, root from scipy.spatial import cKDTree +from vaft.data.cocos import VAFT_INTERNAL_COCOS, cocos_spec from vaft.data.equilibrium import ( BoundaryRepresentation, Contour, @@ -1368,7 +1369,14 @@ def _solovev_components(model: SolovevEquilibrium, r: Any, z: Any) -> tuple[np.n return psi, dpsi_dr, dpsi_dz, basis -def evaluate_solovev(model: SolovevEquilibrium, r: Any, z: Any) -> Mapping[str, np.ndarray]: +def evaluate_solovev( + model: SolovevEquilibrium, + r: Any, + z: Any, + *, + cocos: int | None = VAFT_INTERNAL_COCOS, + convention: int | None = None, +) -> Mapping[str, np.ndarray]: """Evaluate an analytic Solov'ev equilibrium on a set of points. The Solov'ev solution is the Grad-Shafranov equation's closed form for a @@ -1384,6 +1392,12 @@ def evaluate_solovev(model: SolovevEquilibrium, r: Any, z: Any) -> Mapping[str, Major radius, strictly positive [m]. z : array_like Height [m]. + cocos : int or None, optional + COCOS coordinate convention index (1-8, 11-18) setting the poloidal + field orientation factor sigma = sigma_RphiZ * sigma_Bp; None selects + the historical orientation fallback (k = -1) [-]. + convention : int or None, optional + Alias for *cocos* [-]. Returns ------- @@ -1396,18 +1410,22 @@ def evaluate_solovev(model: SolovevEquilibrium, r: Any, z: Any) -> Mapping[str, Raises ------ ValueError - Any point has a non-positive major radius. + Any point has a non-positive major radius, or an invalid COCOS index is + provided. Convention ---------- Works in poloidal flux per radian, with the poloidal field taken as - ``B_R = -dpsi/dz / R`` and ``B_Z = +dpsi/dR / R``. That is the orientation - family with a positive product of the toroidal and poloidal sign conventions, - which is *not* the family the rest of this module falls back to for an - unidentified equilibrium; the mismatch is tracked in #600. + ``B_R = (sigma / R) dpsi/dz`` and ``B_Z = -(sigma / R) dpsi/dR``, where + ``sigma = sigma_RphiZ * sigma_Bp`` is determined by *cocos*. The default + is :data:`~vaft.data.cocos.VAFT_INTERNAL_COCOS` (11), which has ``sigma = +1``; + passing ``cocos=None`` retains the historical fallback ``sigma = -1`` + (the COCOS 2/3/6/7 orientation family). :func:`solovev_to_equilibrium` is where this is reconciled with a declared - convention. Pressure and the squared poloidal current are linear in psi by - construction, so their gradients are the model's own constants. + convention, scaling the flux by 2*pi for full-weber conventions and + transforming signs appropriately. Pressure and the squared poloidal + current are linear in psi by construction, so their gradients are the + model's own constants. Applicability ------------- @@ -1427,6 +1445,20 @@ def evaluate_solovev(model: SolovevEquilibrium, r: Any, z: Any) -> Mapping[str, .. [2] The five-term homogeneous basis and the particular integral follow the standard polynomial construction for that solution. """ + if convention is not None: + if cocos != VAFT_INTERNAL_COCOS and cocos != convention: + raise ValueError("both cocos and convention were provided with conflicting values") + resolved_cocos = convention + else: + resolved_cocos = cocos + + if resolved_cocos is None: + k_sign = -1.0 + else: + if resolved_cocos not in range(1, 19) or resolved_cocos in (9, 10): + raise ValueError(f"cocos must be a valid COCOS index in 1..8 or 11..18, got {resolved_cocos}") + spec = cocos_spec(int(resolved_cocos)) + k_sign = float(spec.sigma_rpz * spec.sigma_bp) psi, dpsi_dr, dpsi_dz, _ = _solovev_components(model, r, z) rr = np.asarray(r, dtype=float) @@ -1436,7 +1468,7 @@ def evaluate_solovev(model: SolovevEquilibrium, r: Any, z: Any) -> Mapping[str, f = model.f_sign*np.sqrt(np.clip(f_squared, 0, None)) return { "psi": psi, "dpsi_dr": dpsi_dr, "dpsi_dz": dpsi_dz, - "b_r": -dpsi_dz/rr, "b_z": dpsi_dr/rr, "b_phi": f/rr, + "b_r": k_sign*dpsi_dz/rr, "b_z": -k_sign*dpsi_dr/rr, "b_phi": f/rr, "pressure": pressure, "f": f, "j_phi": rr*model.pprime + model.ffprime/(MU0*rr), "grad_shafranov_source": -MU0*rr**2*model.pprime-model.ffprime, @@ -1680,18 +1712,12 @@ def solovev_to_equilibrium( Convention ---------- - :func:`evaluate_solovev` works in weber per radian, taking the poloidal field - as the flux gradient over the major radius. The exported flux quantities are - scaled to honour the declared *convention*: a full-weber index, 11 to 18, - multiplies psi by ``2*pi`` so that descriptor derivation, which divides a - full-weber flux gradient by ``2*pi``, recovers the analytic fields exactly. - The two source gradients are rescaled inversely, because a derivative against - flux scales by the inverse of what flux does. The plasma current and the - current density are deliberately *not* scaled: they are physical fields that - the storage choice does not change. Only the ``2*pi`` half of the - reconciliation is handled here; the orientation sign is discussed at - :func:`evaluate_solovev` and tracked in #600. The unscaled quantities are - noted in #608. + :func:`evaluate_solovev` works in weber per radian under the requested + *convention*'s orientation. The exported equilibrium is natively constructed + in :data:`~vaft.data.cocos.VAFT_INTERNAL_COCOS` (11, full weber) and converted + to *convention* via :func:`convert_cocos` if different, reconciling both the + 2*pi normalisation and the orientation sign across all COCOS conventions. + Pressure and the magnetic axis are invariant under convention transforms. Applicability ------------- @@ -1710,12 +1736,15 @@ def solovev_to_equilibrium( .. [1] Solov'ev (1968) through :func:`evaluate_solovev`, which supplies every field this exports. """ + if convention not in range(1, 19) or convention in (9, 10): + raise ValueError("convention must be a COCOS index in the range 1..18 (excluding 9 and 10)") r = np.asarray(r, dtype=float).reshape(-1); z = np.asarray(z, dtype=float).reshape(-1) rm, zm = np.meshgrid(r, z, indexing="ij") - values = evaluate_solovev(model, rm, zm); psi = values["psi"] + values = evaluate_solovev(model, rm, zm, cocos=11) + psi = values["psi"] if magnetic_axis is None: magnetic_axis = _locate_solovev_axis(model, r, z, psi) - psi_axis = float(evaluate_solovev(model, *magnetic_axis)["psi"]) + psi_axis = float(evaluate_solovev(model, *magnetic_axis, cocos=11)["psi"]) temp = EquilibriumData(r=r, z=z, psi=psi, psi_axis=psi_axis, psi_boundary=model.psi_boundary, magnetic_axis=magnetic_axis) lcfs, lcfs_level = _closed_boundary_contour(temp, magnetic_axis) if lcfs is None: @@ -1728,12 +1757,10 @@ def solovev_to_equilibrium( f = model.f_sign*np.sqrt(np.clip(model.f_boundary**2+2*model.ffprime*(psi_1d-model.psi_boundary), 0, None)) mask = _mpl_path(lcfs.points).contains_points(np.column_stack((rm.ravel(), zm.ravel()))).reshape(rm.shape) ip = float(np.sum(values["j_phi"]*np.gradient(r)[:, None]*np.gradient(z)[None, :]*mask)) - if convention not in range(1, 19): - raise ValueError("convention must be a COCOS index in the range 1..18") - psi_factor = 2.0*np.pi if convention >= 11 else 1.0 - conv = _detect_convention(explicit=convention, bt0=float(model.f_boundary/model.rref), ip=ip, q=None, psi_1d=psi_1d*psi_factor, source="analytic Solovev") + psi_factor = 2.0*np.pi + conv_11 = _detect_convention(explicit=11, bt0=float(model.f_boundary/model.rref), ip=ip, q=None, psi_1d=psi_1d*psi_factor, source="analytic Solovev") # Keyword arguments throughout: the field order is not part of the contract. - return EquilibriumData( + eq_11 = EquilibriumData( r=r, z=z, psi=psi*psi_factor, psi_axis=psi_axis*psi_factor, psi_boundary=model.psi_boundary*psi_factor, magnetic_axis=magnetic_axis, lcfs=lcfs, limiter=limiter, @@ -1743,10 +1770,13 @@ def solovev_to_equilibrium( pprime=np.full(psi_1d.size, model.pprime/psi_factor), ffprime=np.full(psi_1d.size, model.ffprime/psi_factor), ip=ip, bt0=float(model.f_boundary/model.rref), r0=model.rref, - time=None, convention=conv, + time=None, convention=conv_11, metadata={"source_type": "solovev", "model": model, "lcfs_psi_n": lcfs_level, "topology_assumptions": "axisymmetric limited or upper/lower-null"}, ) + if convention == 11: + return eq_11 + return convert_cocos(eq_11, convention) def _grid_spacing(eq: EquilibriumData) -> float: