Simplify safety factor statistics calculation.

Replace the error prone quadratic fitting routine with simpler logic.

For q min we simply return the min value.
For the intercepts we do linear interpolation of each interval and look for potential intercepts in each interval.

This fits better with making fewer assumptions on the underlying profile and fixes the buggy values returned by quadratic fits.

PiperOrigin-RevId: 949037347

Author: Nush395

torax/_src/output_tools/safety_factor_fit.py

@@ -13,9 +13,7 @@
 # limitations under the License.
 """Helper for finding safety factor outputs."""
 import dataclasses
-import functools
 
-import chex
 import jax
 from jax import numpy as jnp
 from torax._src import array_typing
@@ -24,7 +22,7 @@
 @jax.tree_util.register_dataclass
 @dataclasses.dataclass(frozen=True)
 class SafetyFactorFit:
-  """Collection of outputs calculated after each simulation step.
+  """Collection of statistics of safety factor calculated each simulation step.
 
   Attributes:
     rho_q_min: rho_norm at the minimum q.
@@ -53,115 +51,41 @@ class SafetyFactorFit:
   rho_q_3_1_second: array_typing.FloatScalar
 
 
-def _sliding_window_of_three(flat_array: jax.Array) -> jax.Array:
-  """Sliding window equivalent to numpy.lib.stride_tricks.sliding_window."""
-  window_size = 3
-  starts = jnp.arange(len(flat_array) - window_size + 1)
-  return jax.vmap(
-      lambda start: jax.lax.dynamic_slice(flat_array, (start,), (window_size,))
-  )(starts)
-
+def _linear_intercepts(
+    rho_norm: jax.Array, q_face: jax.Array, q_target: float
+) -> jax.Array:
+  """Finds intercepts of q_face with q_target using linear interpolation.
 
-def _fit_polynomial_to_intervals_of_three(
-    rho_norm: jax.Array, q_face: jax.Array
-) -> tuple[jax.Array, jax.Array, jax.Array]:
-  """Fit polynomial to every set of three points in given arrays."""
-  q_face_intervals = _sliding_window_of_three(
-      q_face,
-  )
-  rho_norm_intervals = _sliding_window_of_three(
-      rho_norm,
-  )
+  Args:
+    rho_norm: Array of rho_norm values on face grid.
+    q_face: Array of q values on face grid.
+    q_target: q value to find intercepts for.
 
-  @jax.vmap
-  def batch_polyfit(
-      q_face_interval: jax.Array, rho_norm_interval: jax.Array
-  ) -> jax.Array:
-    """Solve Ax=b to get coefficients for a quadratic polynomial."""
-    chex.assert_shape(q_face_interval, (3,))
-    chex.assert_shape(rho_norm_interval, (3,))
-    rho_norm_squared = rho_norm_interval**2
-    A = jnp.array([  # pylint: disable=invalid-name
-        [rho_norm_squared[0], rho_norm_interval[0], 1],
-        [rho_norm_squared[1], rho_norm_interval[1], 1],
-        [rho_norm_squared[2], rho_norm_interval[2], 1],
-    ])
-    b = jnp.array([q_face_interval[0], q_face_interval[1], q_face_interval[2]])
-    coeffs = jnp.linalg.solve(A, b)
-    return coeffs
-
-  return (
-      batch_polyfit(q_face_intervals, rho_norm_intervals),
-      rho_norm_intervals,
-      q_face_intervals,
-  )
+  Returns:
+    Array of rho_norm values of (potential) intercepts in each interval. If no
+    intercept is found in an interval, -jnp.inf is placed in the array.
+  """
+  q_diff = q_face[1:] - q_face[:-1]
+  safe_q_diff = jnp.where(q_diff == 0.0, 1.0, q_diff)
+  t = (q_target - q_face[:-1]) / safe_q_diff
 
+  # If q_diff is 0 and q_face[:-1] == q_target, we force t = 1.0 (intercept at
+  # rho_norm[i+1])
+  t = jnp.where((q_diff == 0.0) & (q_face[:-1] == q_target), 1.0, t)
 
-@jax.vmap
-def _minimum_location_value_in_interval(
-    coeffs: jax.Array, rho_norm_interval: jax.Array, q_interval: jax.Array
-) -> tuple[jax.Array, jax.Array]:
-  """Returns the minimum value and location of the fit quadratic in the interval."""
-  min_interval, max_interval = rho_norm_interval[0], rho_norm_interval[1]
-  q_min_interval, q_max_interval = (
-      q_interval[0],
-      q_interval[1],
-  )
-  a, b = coeffs[0], coeffs[1]
-  extremum_location = -b / (2 * a)
-  extremum_in_interval = jnp.greater(
-      extremum_location, min_interval
-  ) & jnp.less(extremum_location, max_interval)
-  extremum_value = jax.lax.cond(
-      extremum_in_interval,
-      lambda x: jnp.polyval(coeffs, x),
-      lambda x: jnp.inf,
-      extremum_location,
-  )
+  is_intercept = (t > 0.0) & (t <= 1.0)
+  # Mask out cases where q_diff == 0 and q_face[:-1] != q_target
+  is_intercept = is_intercept & ~((q_diff == 0.0) & (q_face[:-1] != q_target))
 
-  interval_minimum_location, interval_minimum_value = jax.lax.cond(
-      jnp.less(q_min_interval, q_max_interval),
-      lambda: (min_interval, q_min_interval),
-      lambda: (max_interval, q_max_interval),
-  )
-  overall_minimum_location, overall_minimum_value = jax.lax.cond(
-      jnp.less(interval_minimum_value, extremum_value),
-      lambda: (interval_minimum_location, interval_minimum_value),
-      lambda: (extremum_location, extremum_value),
-  )
-  return overall_minimum_location, overall_minimum_value
-
-
-def _find_roots_quadratic(coeffs: jax.Array) -> jax.Array:
-  """Finds the roots of a quadratic equation."""
-  a, b, c = coeffs[0], coeffs[1], coeffs[2]
-  determinant = b**2 - 4.0 * a * c
-  roots_exist = jnp.greater(determinant, 0)
-  plus_root = jax.lax.cond(
-      roots_exist,
-      lambda: (-b + jnp.sqrt(determinant)) / (2.0 * a),
-      lambda: -jnp.inf,
-  )
-  minus_root = jax.lax.cond(
-      roots_exist,
-      lambda: (-b - jnp.sqrt(determinant)) / (2.0 * a),
-      lambda: -jnp.inf,
-  )
-  return jnp.array([plus_root, minus_root])
+  intercepts = rho_norm[:-1] + t * (rho_norm[1:] - rho_norm[:-1])
 
+  all_intercepts = jnp.where(is_intercept, intercepts, -jnp.inf)
 
-@functools.partial(jax.vmap, in_axes=(0, 0, None))
-def _root_in_interval(
-    coeffs: jax.Array, interval: jax.Array, q_surface: float
-) -> jax.Array:
-  """Finds roots of quadratic(coeffs)=q_surface in the interval."""
-  intercept_coeffs = coeffs - jnp.array([0.0, 0.0, q_surface])
-  min_interval, max_interval = interval[0], interval[1]
-  root_values = _find_roots_quadratic(intercept_coeffs)
-  in_interval = jnp.greater(root_values, min_interval) & jnp.less(
-      root_values, max_interval
+  first_point_intercept = jnp.where(
+      q_face[0] == q_target, rho_norm[0], -jnp.inf
   )
-  return jnp.where(in_interval, root_values, -jnp.inf)
+
+  return jnp.concat([jnp.array([first_point_intercept]), all_intercepts])
 
 
 @jax.jit
@@ -170,33 +94,16 @@ def find_min_q_and_q_surface_intercepts(
 ) -> SafetyFactorFit:
   """Finds the minimum q and the q surface intercepts.
 
-  This method divides the input arrays into groups of three points, fits a
-  quadratic to each interval containing three points. It then uses the fit
-  quadratics to find the minimum q value as well as any intercepts of the q=3/2,
-  q=2/1, and q=3/1 planes.
-
-  As the quadratic fits could have overlapping definitions where the groups
-  overlap, we choose the quadratic for the domain [rho_norm[i-1], rho_norm[i]]
-  defined using the points (i, i-1, i-2).
-  (There is one exception which is the first group for which the domain for
-  points in [rho_norm[0], rho_norm[2]] are defined using points (0, 1, 2))
+  This method uses the input arrays to find the minimum q value and its
+  corresponding rho_norm. It then uses linear interpolation between adjacent
+  points to find any intercepts of the q=3/2, q=2/1, and q=3/1 planes.
 
   Args:
     rho_norm: Array of rho_norm values on face grid.
     q_face: Array of q values on face grid.
 
   Returns:
-    rho_q_min: rho_norm at the minimum q.
-    q_min: minimum q value.
-    outermost_rho_q_3_2: Array of two outermost rho_norm values that intercept
-    the q=3/2 plane. If fewer than two intercepts are found, entries will be set
-    to -inf.
-    outermost_rho_q_2_1: Array of 2 outermost rho_norm values that intercept the
-    q=2/1 plane. If fewer than two intercepts are found, entries will be set to
-    -inf.
-    outermost_rho_q_3_1: Array of 2 outermost rho_norm values that intercept the
-    q=3/1 plane. If fewer than two intercepts are found, entries will be set to
-    -inf.
+    Safety factor statistics.
   """
   if len(q_face) != len(rho_norm):
     raise ValueError(
@@ -209,36 +116,21 @@ def find_min_q_and_q_surface_intercepts(
   sorted_indices = jnp.argsort(rho_norm)
   rho_norm = rho_norm[sorted_indices]
   q_face = q_face[sorted_indices]
-  # Fit polynomial to every set of three points in given arrays.
-  poly_coeffs, rho_norm_3, q_face_3 = _fit_polynomial_to_intervals_of_three(
-      rho_norm, q_face
-  )
-  # To bridge across the entire span of rho_norm, and avoid overlapping
-  # interval domains, we define the first group from points (0, 2), and all
-  # subsequent groups from (i+1, i+2), from i=1 until i=len(rho_norm)-1.
-  first_rho_norm = jnp.expand_dims(
-      jnp.array([rho_norm[0], rho_norm[2]]), axis=0
-  )
-  first_q_face = jnp.expand_dims(jnp.array([q_face[0], q_face[2]]), axis=0)
-  rho_norms = jnp.concat([first_rho_norm, rho_norm_3[1:, 1:]], axis=0)
-  q_faces = jnp.concat([first_q_face, q_face_3[1:, 1:]], axis=0)
 
-  # Find the minimum q value and its location for each interval.
-  rho_q_min_intervals, q_min_intervals = _minimum_location_value_in_interval(
-      poly_coeffs, rho_norms, q_faces
-  )
-  # Find the minimum q value and its location out of all intervals.
-  arg_q_min = jnp.argmin(q_min_intervals)
-  rho_q_min = rho_q_min_intervals[arg_q_min]
-  q_min = q_min_intervals[arg_q_min]
+  # Find the minimum q value and its location.
+  idx_min = jnp.argmin(q_face)
+  q_min = q_face[idx_min]
+  rho_q_min = rho_norm[idx_min]
 
   # Find the outermost rho_norm values that intercept the q=3/2, q=2/1, and
   # q=3/1 planes. If none are found, fill from the left with -jnp.inf.
-  rho_q_3_2 = _root_in_interval(poly_coeffs, rho_norms, 1.5).flatten()
+  rho_q_3_2 = _linear_intercepts(rho_norm, q_face, 1.5)
   outermost_rho_q_3_2 = rho_q_3_2[jnp.argsort(rho_q_3_2)[-2:]]
-  rho_q_2_1 = _root_in_interval(poly_coeffs, rho_norms, 2.0).flatten()
+
+  rho_q_2_1 = _linear_intercepts(rho_norm, q_face, 2.0)
   outermost_rho_q_2_1 = rho_q_2_1[jnp.argsort(rho_q_2_1)[-2:]]
-  rho_q_3_1 = _root_in_interval(poly_coeffs, rho_norms, 3.0).flatten()
+
+  rho_q_3_1 = _linear_intercepts(rho_norm, q_face, 3.0)
   outermost_rho_q_3_1 = rho_q_3_1[jnp.argsort(rho_q_3_1)[-2:]]
 
   return SafetyFactorFit(

torax/_src/output_tools/tests/safety_factor_fit_test.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 from absl.testing import absltest
 from absl.testing import parameterized
-import chex
 import jax
 from jax import numpy as jnp
 import numpy as np
@@ -72,7 +71,7 @@ def tearDown(self):
   def test_find_min_q_and_q_surface_intercepts_finds_minimum_q(
       self, transform, expected_q_min, expected_rho_q_min
   ):
-    rho_norm_face = jnp.linspace(0.0, 1.0, 1000)
+    rho_norm_face = jnp.linspace(0.0, 1.0, 1001)
     q_face = transform(rho_norm_face)
     outputs = safety_factor_fit.find_min_q_and_q_surface_intercepts(
         rho_norm_face, q_face
@@ -90,26 +89,26 @@ def test_find_min_q_and_q_surface_intercepts_finds_minimum_q(
       dict(
           testcase_name="3x^2+0.5",
           transform=lambda x: 3 * x**2 + 0.5,
-          expected_q_3_2=[-np.inf, 1 / np.sqrt(3)],
-          expected_q_2_1=[-np.inf, 1 / np.sqrt(2)],
-          expected_q_3_1=[-np.inf, np.sqrt(5 / 6)],
+          expected_q_3_2=[-np.inf, 0.5773500721500722],
+          expected_q_2_1=[-np.inf, 0.7071067137809187],
+          expected_q_3_1=[-np.inf, 0.91287086758],
       ),
       dict(
           testcase_name="2x^2+1",
           transform=lambda x: 2 * x**2 + 1,
           expected_q_3_2=[-np.inf, 0.5],
-          expected_q_2_1=[-np.inf, 1 / np.sqrt(2)],
+          expected_q_2_1=[-np.inf, 0.7071067137809187],
           expected_q_3_1=[
               -np.inf,
-              -np.inf,
-          ],  # We don't find the root on the border.
+              1.0,
+          ],
       ),
       dict(
           testcase_name="20*(x-0.5)^2",
           transform=lambda x: 20 * (x - 0.5) ** 2,
-          expected_q_3_2=[(10 - np.sqrt(30)) / 20, (10 + np.sqrt(30)) / 20],
-          expected_q_2_1=[0.5 - 1 / np.sqrt(10), 0.5 + 1 / np.sqrt(10)],
-          expected_q_3_1=[(5 - np.sqrt(15)) / 10, (5 + np.sqrt(15)) / 10],
+          expected_q_3_2=[0.2261389396709324, 0.7738610603290677],
+          expected_q_2_1=[0.18377251184834123, 0.8162274881516588],
+          expected_q_3_1=[0.11270193548387099, 0.887298064516129],
       ),
       dict(
           testcase_name="-3x^2-0.5",
@@ -136,7 +135,7 @@ def test_find_min_q_and_q_surface_intercepts_finds_minimum_q(
   def test_find_min_q_and_q_surface_intercepts_finds_q_surface_intercepts(
       self, transform, expected_q_3_2, expected_q_2_1, expected_q_3_1
   ):
-    rho_norm_face = jnp.linspace(0.0, 1.0, 1000)
+    rho_norm_face = jnp.linspace(0.0, 1.0, 1001)
     q_face = transform(rho_norm_face)
     outputs = safety_factor_fit.find_min_q_and_q_surface_intercepts(
         rho_norm_face, q_face
@@ -151,53 +150,21 @@ def test_find_min_q_and_q_surface_intercepts_finds_q_surface_intercepts(
     with self.subTest("q_3_1"):
       np.testing.assert_allclose(q_3_1, np.array(expected_q_3_1))
 
-  def test_polynomial_fit_to_threes_on_exact_quadratic(self):
-    """Test to give us an idea of the error on the quadratic fit."""
-    rho_norm = jnp.array([0.0, 0.5, 1.0], dtype=jnp.float64)
-    random_coeffs = jax.random.normal(
-        key=jax.random.PRNGKey(0), shape=(3,), dtype=jnp.float64
-    )
-    q_face = (
-        random_coeffs[0] * rho_norm**2
-        + random_coeffs[1] * rho_norm
-        + random_coeffs[2]
-    )
-    polyfit, rho_norm_3, q_face_3 = (
-        safety_factor_fit._fit_polynomial_to_intervals_of_three(
-            rho_norm, q_face
-        )
-    )
-    chex.assert_shape(polyfit, (1, 3))
-    chex.assert_shape(rho_norm_3, (1, 3))
-    chex.assert_shape(q_face_3, (1, 3))
-    np.testing.assert_allclose(polyfit[0], random_coeffs)
+  def test_linear_intercepts(self):
+    rho_norm = jnp.array([0.0, 0.5, 1.0])
+    q_face = jnp.array([1.0, 2.0, 3.0])
 
-  def test_root_in_interval_finds_correct_roots(self):
-    """Test correct roots found in each interval.
+    with self.subTest("target_1_5"):
+      intercepts = safety_factor_fit._linear_intercepts(rho_norm, q_face, 1.5)
+      np.testing.assert_allclose(intercepts, np.array([-np.inf, 0.25, -np.inf]))
 
-    y=4      *
-      |
-      |
-    y=1  *       *
-      |
-      *--1---2---3----> x
-    """
-    # Given two polynomials that fit the above points.
-    a_1, b_1, c_1 = 1.0, 0.0, 0.0
-    a_2, b_2, c_2 = -3.0, 12.0, -8.0
-    coeffs = np.array([[a_1, b_1, c_1], [a_2, b_2, c_2]])
-    rho_norm = np.array([[0.0, 2.0], [2.0, 3.0]])
-    # When we calculate roots for intercepts of the q=3/2 plane.
-    intercept = 1.5
-    intercepts = safety_factor_fit._root_in_interval(
-        coeffs, rho_norm, intercept
-    )
-    # Then expect only one root to be found in first interval.
-    np.testing.assert_allclose(
-        intercepts[0], np.array([np.sqrt(intercept), -np.inf])
-    )
-    # And expect only one root in second interval (as only over [2.0, 3.0]]).
-    np.testing.assert_allclose(intercepts[1], np.array([-np.inf, 2.912871]))
+    with self.subTest("target_2_0"):
+      intercepts = safety_factor_fit._linear_intercepts(rho_norm, q_face, 2.0)
+      np.testing.assert_allclose(intercepts, np.array([-np.inf, 0.5, -np.inf]))
+
+    with self.subTest("target_1_0"):
+      intercepts = safety_factor_fit._linear_intercepts(rho_norm, q_face, 1.0)
+      np.testing.assert_allclose(intercepts, np.array([0.0, -np.inf, -np.inf]))
 
 
 if __name__ == "__main__":