Simplify safety factor statistics calculation.

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__":