Centralize sign convention enforcement in StandardGeometryIntermediates.

This change removes np.abs() calls from geometry loading functions (fbt.py, equilibrium.py) and instead enforces sign conventions for psi and Ip_profile within the StandardGeometryIntermediates.__post_init__ method. This ensures that psi increases with radius and Ip_profile is positive, regardless of the input data's sign conventions. Tests are added to verify the sign flipping behavior.

This also fixes a bug where any initial psi from geo that was not strictly positive or strictly negative would get a "kink" due to the abs. We now ensure correct trend of psi, and do not enforce an abs

PiperOrigin-RevId: 883157580

Author: jcitrin

torax/_src/geometry/fbt.py

@@ -505,22 +505,22 @@ def _get_val(key):
       B_0=B_0,
       psi=psi,
       Phi=Phi,
-      Ip_profile=np.abs(LY['ItQ']),
+      Ip_profile=LY['ItQ'],
       R_in=LY['rgeom'] - LY['aminor'],
       R_out=LY['rgeom'] + LY['aminor'],
-      F=np.abs(LY['TQ']),
+      F=LY['TQ'],
       int_dl_over_Bp=1 / LY_Q1Q,
       flux_surf_avg_1_over_R=LY['Q0Q'],
       flux_surf_avg_1_over_R2=LY['Q2Q'],
-      flux_surf_avg_grad_psi2_over_R2=np.abs(LY['Q3Q']),
-      flux_surf_avg_grad_psi=np.abs(LY['Q5Q']),
-      flux_surf_avg_grad_psi2=np.abs(LY['Q4Q']),
+      flux_surf_avg_grad_psi2_over_R2=LY['Q3Q'],
+      flux_surf_avg_grad_psi=LY['Q5Q'],
+      flux_surf_avg_grad_psi2=LY['Q4Q'],
       flux_surf_avg_B2=flux_surf_avg_B2,
       flux_surf_avg_1_over_B2=flux_surf_avg_1_over_B2,
       delta_upper_face=LY['deltau'],
       delta_lower_face=LY['deltal'],
       elongation=LY['kappa'],
-      vpr=4 * np.pi * Phi[-1] * rhon / (np.abs(LY['TQ']) * LY['Q2Q']),
+      vpr=4 * np.pi * Phi[-1] * rhon / (LY['TQ'] * LY['Q2Q']),
       face_centers=face_centers,
       hires_factor=hires_factor,
       diverted=diverted,

torax/_src/geometry/standard_geometry.py

@@ -251,13 +251,44 @@ class StandardGeometryIntermediates:
   B_pol_OMP: array_typing.FloatScalar | None
 
   def __post_init__(self):
-    """Extrapolates edge values and smooths near-axis values.
-
+    """Enforces sign conventions, extrapolates edge, and smooths near-axis.
+
+    Sign conventions - TODO(b/335204606): Replace with proper COCOS handling:
+    - psi must grow with radius, i.e. psi(axis) < psi(separatrix).
+    - Total Ip must be positive.
+    - Positive-definite quantities are enforced via abs(): Phi, F,
+      int_dl_over_Bp, vpr, flux_surf_avg_grad_psi, flux_surf_avg_grad_psi2,
+      flux_surf_avg_grad_psi2_over_R2.
+    Then:
     - Edge extrapolation for a subset of attributes based on a Cubic spline fit.
     - Near-axis smoothing for a subset of attributes based on a Savitzky-Golay
       filter with an appropriate polynominal order based on the attribute.
     """
 
+    if self.psi[-1] < self.psi[0]:
+      object.__setattr__(self, 'psi', -self.psi)
+
+    if self.Ip_profile[-1] < 0:
+      object.__setattr__(self, 'Ip_profile', -self.Ip_profile)
+
+    object.__setattr__(self, 'Phi', np.abs(self.Phi))
+    object.__setattr__(self, 'F', np.abs(self.F))
+    object.__setattr__(
+        self, 'int_dl_over_Bp', np.abs(self.int_dl_over_Bp)
+    )
+    object.__setattr__(self, 'vpr', np.abs(self.vpr))
+    object.__setattr__(
+        self, 'flux_surf_avg_grad_psi', np.abs(self.flux_surf_avg_grad_psi)
+    )
+    object.__setattr__(
+        self, 'flux_surf_avg_grad_psi2', np.abs(self.flux_surf_avg_grad_psi2)
+    )
+    object.__setattr__(
+        self,
+        'flux_surf_avg_grad_psi2_over_R2',
+        np.abs(self.flux_surf_avg_grad_psi2_over_R2),
+    )
+
     # Check if last flux surface is diverted and correct via spline fit if so
     if self.diverted or self.flux_surf_avg_grad_psi2_over_R2[-1] < 1e-10:
       # Calculate rhon

torax/_src/geometry/tests/standard_geometry_test.py

@@ -146,6 +146,85 @@ def test_stack_geometries_standard_geometries(self):
         stacked_geo.g1_over_vpr2_face[:, 0], 1 / stacked_geo.rho_b**2
     )
 
+  def _make_intermediates(self, **overrides):
+    defaults = dict(
+        geometry_type=geometry.GeometryType.FBT,
+        Ip_from_parameters=True,
+        R_major=6.2,
+        a_minor=2.0,
+        B_0=5.3,
+        psi=np.linspace(0, 1.0, 100),
+        Ip_profile=np.linspace(0, 1e6, 100),
+        Phi=np.linspace(0, 1.0, 100),
+        R_in=np.linspace(4.0, 4.2, 100),
+        R_out=np.linspace(8.0, 8.4, 100),
+        F=np.linspace(30.0, 33.0, 100),
+        int_dl_over_Bp=np.linspace(0.01, 1.0, 100),
+        flux_surf_avg_1_over_R=np.linspace(0.1, 0.2, 100),
+        flux_surf_avg_1_over_R2=np.linspace(0.01, 0.04, 100),
+        flux_surf_avg_grad_psi2=np.linspace(0.01, 1.0, 100),
+        flux_surf_avg_grad_psi=np.linspace(0.01, 1.0, 100),
+        flux_surf_avg_grad_psi2_over_R2=np.linspace(0.01, 1.0, 100),
+        flux_surf_avg_B2=np.linspace(25.0, 30.0, 100),
+        flux_surf_avg_1_over_B2=np.linspace(0.03, 0.04, 100),
+        delta_upper_face=np.linspace(0.0, 0.3, 100),
+        delta_lower_face=np.linspace(0.0, 0.3, 100),
+        elongation=np.linspace(1.0, 1.7, 100),
+        vpr=np.linspace(0.01, 1.0, 100),
+        face_centers=interpolated_param_2d.get_face_centers(25),
+        hires_factor=4,
+        z_magnetic_axis=np.array(0.0),
+        diverted=None,
+        connection_length_target=None,
+        connection_length_divertor=None,
+        angle_of_incidence_target=None,
+        R_OMP=None,
+        R_target=None,
+        B_pol_OMP=None,
+    )
+    defaults.update(overrides)
+    return standard_geometry.StandardGeometryIntermediates(**defaults)
+
+  def test_post_init_flips_psi_when_decreasing(self):
+    psi_decreasing = np.linspace(1.0, 0.0, 100)
+    intermediates = self._make_intermediates(psi=psi_decreasing)
+    self.assertGreater(intermediates.psi[-1], intermediates.psi[0])
+
+  def test_post_init_preserves_psi_when_increasing(self):
+    psi_increasing = np.linspace(0.0, 1.0, 100)
+    intermediates = self._make_intermediates(psi=psi_increasing)
+    np.testing.assert_array_equal(intermediates.psi, psi_increasing)
+
+  def test_post_init_flips_negative_Ip(self):
+    Ip_negative = np.linspace(0, -1e6, 100)
+    intermediates = self._make_intermediates(Ip_profile=Ip_negative)
+    self.assertGreater(intermediates.Ip_profile[-1], 0)
+
+  def test_post_init_preserves_positive_Ip(self):
+    Ip_positive = np.linspace(0, 1e6, 100)
+    intermediates = self._make_intermediates(Ip_profile=Ip_positive)
+    np.testing.assert_array_equal(intermediates.Ip_profile, Ip_positive)
+
+  def test_post_init_enforces_positive_definite_quantities(self):
+    intermediates = self._make_intermediates(
+        Phi=-np.linspace(0, 1.0, 100),
+        F=-np.linspace(30.0, 33.0, 100),
+        int_dl_over_Bp=-np.linspace(0.01, 1.0, 100),
+        vpr=-np.linspace(0.01, 1.0, 100),
+        flux_surf_avg_grad_psi=-np.linspace(0.01, 1.0, 100),
+        flux_surf_avg_grad_psi2=-np.linspace(0.01, 1.0, 100),
+        flux_surf_avg_grad_psi2_over_R2=-np.linspace(0.01, 1.0, 100),
+    )
+    np.testing.assert_array_less(-1e-15, intermediates.Phi)
+    np.testing.assert_array_less(0, intermediates.F)
+    np.testing.assert_array_less(0, intermediates.int_dl_over_Bp)
+    np.testing.assert_array_less(0, intermediates.vpr)
+    np.testing.assert_array_less(0, intermediates.flux_surf_avg_grad_psi)
+    np.testing.assert_array_less(0, intermediates.flux_surf_avg_grad_psi2)
+    np.testing.assert_array_less(
+        0, intermediates.flux_surf_avg_grad_psi2_over_R2
+    )
+
 
 if __name__ == '__main__':
   absltest.main()

torax/_src/imas_tools/input/equilibrium.py

@@ -118,33 +118,38 @@ def geometry_from_IMAS(
         f"{len(equilibrium.time_slice)} time slices"
     )
   IMAS_data = equilibrium.time_slice[slice_index]
+  # IMAS python API returns custom primitive types (e.g. IDSFloat0D,
+  # IDSNumericArray) instead of standard python floats or numpy arrays. We must
+  # cast these to standard numpy arrays using np.asarray() (or via implicit
+  # numpy operations like np.abs), otherwise JAX JIT compilation will fail with
+  # TypeErrors during PyTree tracing.
   R_major = np.asarray(equilibrium.vacuum_toroidal_field.r0)
-  B_0 = np.asarray(np.abs(equilibrium.vacuum_toroidal_field.b0[0]))
+  B_0 = np.abs(equilibrium.vacuum_toroidal_field.b0[0])
 
   # Poloidal flux.
-  psi = np.abs(IMAS_data.profiles_1d.psi)
+  psi = np.asarray(IMAS_data.profiles_1d.psi)
 
   # Toroidal flux.
-  phi = np.abs(IMAS_data.profiles_1d.phi)
+  phi = np.asarray(IMAS_data.profiles_1d.phi)
 
   # Midplane radii.
   R_in = IMAS_data.profiles_1d.r_inboard
   R_out = IMAS_data.profiles_1d.r_outboard
   R_major_profile = (R_in + R_out) / 2.0
   # toroidal field flux function
-  F = np.abs(IMAS_data.profiles_1d.f)
+  F = np.asarray(IMAS_data.profiles_1d.f)
 
   # Flux surface integrals of various geometry quantities.
   # IDS Contour integrals.
   if IMAS_data.profiles_1d.dvolume_dpsi:
-    dvoldpsi = np.abs(IMAS_data.profiles_1d.dvolume_dpsi)
+    dvoldpsi = np.asarray(IMAS_data.profiles_1d.dvolume_dpsi)
   else:
     dvoldpsi = np.gradient(
         IMAS_data.profiles_1d.volume, IMAS_data.profiles_1d.psi
     )
   # dpsi_drho_tor
   if IMAS_data.profiles_1d.dpsi_drho_tor:
-    dpsidrhotor = np.abs(IMAS_data.profiles_1d.dpsi_drho_tor)
+    dpsidrhotor = np.asarray(IMAS_data.profiles_1d.dpsi_drho_tor)
   else:
     rho_tor = IMAS_data.profiles_1d.rho_tor
     if not rho_tor:

torax/tests/sim_test.py

@@ -296,6 +296,8 @@ class SimTest(sim_test_case.SimTestCase):
       (
           'test_iterhybrid_lh_transition',
           'test_iterhybrid_lh_transition.py',
+          _ALL_PROFILES,
+          1e-7,
       ),
       # Tests used for testing changing configs without recompiling.
       # Based on test_iterhybrid_predictor_corrector