Align evaluate_solovev poloidal field orientation with COCOS conventions (#600)

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"

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"))

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: