diff --git a/torax/_src/output_tools/safety_factor_fit.py b/torax/_src/output_tools/safety_factor_fit.py index dce470d6e..b74517e88 100644 --- a/torax/_src/output_tools/safety_factor_fit.py +++ b/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( diff --git a/torax/_src/output_tools/tests/safety_factor_fit_test.py b/torax/_src/output_tools/tests/safety_factor_fit_test.py index 1c6811c84..207256737 100644 --- a/torax/_src/output_tools/tests/safety_factor_fit_test.py +++ b/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__": diff --git a/torax/tests/test_data/test_all_transport_fusion_qlknn.nc b/torax/tests/test_data/test_all_transport_fusion_qlknn.nc index 2c98c407f..e47bdc2f1 100644 Binary files a/torax/tests/test_data/test_all_transport_fusion_qlknn.nc and b/torax/tests/test_data/test_all_transport_fusion_qlknn.nc differ diff --git a/torax/tests/test_data/test_bohmgyrobohm_all.nc b/torax/tests/test_data/test_bohmgyrobohm_all.nc index f353ba666..6dcfefbd1 100644 Binary files a/torax/tests/test_data/test_bohmgyrobohm_all.nc and b/torax/tests/test_data/test_bohmgyrobohm_all.nc differ diff --git a/torax/tests/test_data/test_bremsstrahlung_time_dependent_Zimp.nc b/torax/tests/test_data/test_bremsstrahlung_time_dependent_Zimp.nc index 1e63d69e0..a1c123480 100644 Binary files a/torax/tests/test_data/test_bremsstrahlung_time_dependent_Zimp.nc and b/torax/tests/test_data/test_bremsstrahlung_time_dependent_Zimp.nc differ diff --git a/torax/tests/test_data/test_changing_config_after.nc b/torax/tests/test_data/test_changing_config_after.nc index cedf2ea08..5a9f01930 100644 Binary files a/torax/tests/test_data/test_changing_config_after.nc and b/torax/tests/test_data/test_changing_config_after.nc differ diff --git a/torax/tests/test_data/test_changing_config_before.nc b/torax/tests/test_data/test_changing_config_before.nc index 2b1952d37..f2f8617f9 100644 Binary files a/torax/tests/test_data/test_changing_config_before.nc and b/torax/tests/test_data/test_changing_config_before.nc differ diff --git a/torax/tests/test_data/test_chease.nc b/torax/tests/test_data/test_chease.nc index 5ab22c05b..c57e0d313 100644 Binary files a/torax/tests/test_data/test_chease.nc and b/torax/tests/test_data/test_chease.nc differ diff --git a/torax/tests/test_data/test_combined_transport.nc b/torax/tests/test_data/test_combined_transport.nc index bd2c44d72..73e186fdc 100644 Binary files a/torax/tests/test_data/test_combined_transport.nc and b/torax/tests/test_data/test_combined_transport.nc differ diff --git a/torax/tests/test_data/test_fixed_dt.nc b/torax/tests/test_data/test_fixed_dt.nc index c6bcd474d..a3ce87f6b 100644 Binary files a/torax/tests/test_data/test_fixed_dt.nc and b/torax/tests/test_data/test_fixed_dt.nc differ diff --git a/torax/tests/test_data/test_imas_profiles_and_geo.nc b/torax/tests/test_data/test_imas_profiles_and_geo.nc index 511ecc8dc..e904b995f 100644 Binary files a/torax/tests/test_data/test_imas_profiles_and_geo.nc and b/torax/tests/test_data/test_imas_profiles_and_geo.nc differ diff --git a/torax/tests/test_data/test_implicit.nc b/torax/tests/test_data/test_implicit.nc index 0f731814a..d037e13b2 100644 Binary files a/torax/tests/test_data/test_implicit.nc and b/torax/tests/test_data/test_implicit.nc differ diff --git a/torax/tests/test_data/test_implicit_short_optimizer.nc b/torax/tests/test_data/test_implicit_short_optimizer.nc index bb2b24b17..1f94be3b9 100644 Binary files a/torax/tests/test_data/test_implicit_short_optimizer.nc and b/torax/tests/test_data/test_implicit_short_optimizer.nc differ diff --git a/torax/tests/test_data/test_iterbaseline_mockup.nc b/torax/tests/test_data/test_iterbaseline_mockup.nc index 042b70100..2b93d99f5 100644 Binary files a/torax/tests/test_data/test_iterbaseline_mockup.nc and b/torax/tests/test_data/test_iterbaseline_mockup.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_lh_transition.nc b/torax/tests/test_data/test_iterhybrid_lh_transition.nc index 4699197d7..e6bd5371b 100644 Binary files a/torax/tests/test_data/test_iterhybrid_lh_transition.nc and b/torax/tests/test_data/test_iterhybrid_lh_transition.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_lh_transition_internal_boundary_condition.nc b/torax/tests/test_data/test_iterhybrid_lh_transition_internal_boundary_condition.nc index 677a4044c..7707eb407 100644 Binary files a/torax/tests/test_data/test_iterhybrid_lh_transition_internal_boundary_condition.nc and b/torax/tests/test_data/test_iterhybrid_lh_transition_internal_boundary_condition.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_mockup.nc b/torax/tests/test_data/test_iterhybrid_mockup.nc index 276aa40cd..741de19a8 100644 Binary files a/torax/tests/test_data/test_iterhybrid_mockup.nc and b/torax/tests/test_data/test_iterhybrid_mockup.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc index b0b3e6cd9..093d8d9be 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_Lmode_combined.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_Lmode_combined.nc index c6ccf32be..6e9c2e079 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_Lmode_combined.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_Lmode_combined.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc index d03666daf..15b0e33d7 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc index 168934925..69d0de570 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc index 915a9f82d..38cb822a3 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc index 9dcf3582b..018a666ca 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc index 9297e84f9..ac0d32bb2 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_imas.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_imas.nc index 54caaad52..19ef6f129 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_imas.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_imas.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc index fccb23103..bd3907b66 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios.nc index b3df86206..a3256a35d 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.nc index cf96c96b6..72faad473 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_z_eff.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_z_eff.nc index 09f5087e1..dda74bba0 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_z_eff.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_z_eff.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_neoclassical.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_neoclassical.nc index 23c5e767e..6828b1648 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_neoclassical.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_neoclassical.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_rotation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_rotation.nc index 9e04081eb..ced857f4f 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_rotation.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_rotation.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc index 6cdeae021..d62aa6542 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea.nc index fdb81ff8d..3f40dc7fa 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea_rotation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea_rotation.nc index aba1d9316..6b8b2ddb2 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea_rotation.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tglfnn_ukaea_rotation.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc index b35d29d23..0df2513ef 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc index cf240b788..be10f01c9 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc index 754b4a68c..8b662818f 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_radiation_collapse.nc b/torax/tests/test_data/test_iterhybrid_radiation_collapse.nc index 8bbaba315..55e12f168 100644 Binary files a/torax/tests/test_data/test_iterhybrid_radiation_collapse.nc and b/torax/tests/test_data/test_iterhybrid_radiation_collapse.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_rampup.nc b/torax/tests/test_data/test_iterhybrid_rampup.nc index dfa8c5f7b..92ceb8cd7 100644 Binary files a/torax/tests/test_data/test_iterhybrid_rampup.nc and b/torax/tests/test_data/test_iterhybrid_rampup.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_rampup_sawtooth.nc b/torax/tests/test_data/test_iterhybrid_rampup_sawtooth.nc index 547c98783..df0ba3562 100644 Binary files a/torax/tests/test_data/test_iterhybrid_rampup_sawtooth.nc and b/torax/tests/test_data/test_iterhybrid_rampup_sawtooth.nc differ diff --git a/torax/tests/test_data/test_ne_qlknn_deff_veff.nc b/torax/tests/test_data/test_ne_qlknn_deff_veff.nc index f77117855..89e6415bd 100644 Binary files a/torax/tests/test_data/test_ne_qlknn_deff_veff.nc and b/torax/tests/test_data/test_ne_qlknn_deff_veff.nc differ diff --git a/torax/tests/test_data/test_ne_qlknn_defromchie.nc b/torax/tests/test_data/test_ne_qlknn_defromchie.nc index 362dcaebd..0bd6ee1db 100644 Binary files a/torax/tests/test_data/test_ne_qlknn_defromchie.nc and b/torax/tests/test_data/test_ne_qlknn_defromchie.nc differ diff --git a/torax/tests/test_data/test_particle_sources_cgm.nc b/torax/tests/test_data/test_particle_sources_cgm.nc index 0926648bf..cde62c845 100644 Binary files a/torax/tests/test_data/test_particle_sources_cgm.nc and b/torax/tests/test_data/test_particle_sources_cgm.nc differ diff --git a/torax/tests/test_data/test_prescribed_generic_current_source.nc b/torax/tests/test_data/test_prescribed_generic_current_source.nc index 1897bd089..0347c953b 100644 Binary files a/torax/tests/test_data/test_prescribed_generic_current_source.nc and b/torax/tests/test_data/test_prescribed_generic_current_source.nc differ diff --git a/torax/tests/test_data/test_prescribed_timedependent_ne.nc b/torax/tests/test_data/test_prescribed_timedependent_ne.nc index 2b4822990..33d2ac00d 100644 Binary files a/torax/tests/test_data/test_prescribed_timedependent_ne.nc and b/torax/tests/test_data/test_prescribed_timedependent_ne.nc differ diff --git a/torax/tests/test_data/test_prescribed_transport.nc b/torax/tests/test_data/test_prescribed_transport.nc index 6ada7612b..b7b299afd 100644 Binary files a/torax/tests/test_data/test_prescribed_transport.nc and b/torax/tests/test_data/test_prescribed_transport.nc differ diff --git a/torax/tests/test_data/test_psi_and_heat.nc b/torax/tests/test_data/test_psi_and_heat.nc index 66aee3e9d..da815df3d 100644 Binary files a/torax/tests/test_data/test_psi_and_heat.nc and b/torax/tests/test_data/test_psi_and_heat.nc differ diff --git a/torax/tests/test_data/test_psi_heat_dens.nc b/torax/tests/test_data/test_psi_heat_dens.nc index e5a40be4d..ec7eca76c 100644 Binary files a/torax/tests/test_data/test_psi_heat_dens.nc and b/torax/tests/test_data/test_psi_heat_dens.nc differ diff --git a/torax/tests/test_data/test_psichease_ip_chease_vloop.nc b/torax/tests/test_data/test_psichease_ip_chease_vloop.nc index ece4138e5..2e2fe6522 100644 Binary files a/torax/tests/test_data/test_psichease_ip_chease_vloop.nc and b/torax/tests/test_data/test_psichease_ip_chease_vloop.nc differ diff --git a/torax/tests/test_data/test_psichease_ip_parameters_vloop_varying.nc b/torax/tests/test_data/test_psichease_ip_parameters_vloop_varying.nc index b73bce961..5d86beb93 100644 Binary files a/torax/tests/test_data/test_psichease_ip_parameters_vloop_varying.nc and b/torax/tests/test_data/test_psichease_ip_parameters_vloop_varying.nc differ diff --git a/torax/tests/test_data/test_psichease_prescribed_johm.nc b/torax/tests/test_data/test_psichease_prescribed_johm.nc index 2aae9adc6..ab3ddc71b 100644 Binary files a/torax/tests/test_data/test_psichease_prescribed_johm.nc and b/torax/tests/test_data/test_psichease_prescribed_johm.nc differ diff --git a/torax/tests/test_data/test_psichease_prescribed_jtot.nc b/torax/tests/test_data/test_psichease_prescribed_jtot.nc index f50bcb4da..37b303244 100644 Binary files a/torax/tests/test_data/test_psichease_prescribed_jtot.nc and b/torax/tests/test_data/test_psichease_prescribed_jtot.nc differ diff --git a/torax/tests/test_data/test_psichease_prescribed_jtot_vloop.nc b/torax/tests/test_data/test_psichease_prescribed_jtot_vloop.nc index 050d6cc29..ad4c9f665 100644 Binary files a/torax/tests/test_data/test_psichease_prescribed_jtot_vloop.nc and b/torax/tests/test_data/test_psichease_prescribed_jtot_vloop.nc differ diff --git a/torax/tests/test_data/test_semiimplicit_convection.nc b/torax/tests/test_data/test_semiimplicit_convection.nc index cfefb75d2..429c1d432 100644 Binary files a/torax/tests/test_data/test_semiimplicit_convection.nc and b/torax/tests/test_data/test_semiimplicit_convection.nc differ diff --git a/torax/tests/test_data/test_step_flattop_bgb.nc b/torax/tests/test_data/test_step_flattop_bgb.nc index 7d8519b3f..352ff65ed 100644 Binary files a/torax/tests/test_data/test_step_flattop_bgb.nc and b/torax/tests/test_data/test_step_flattop_bgb.nc differ diff --git a/torax/tests/test_data/test_timedependence.nc b/torax/tests/test_data/test_timedependence.nc index 7e73a0066..6c990fd0b 100644 Binary files a/torax/tests/test_data/test_timedependence.nc and b/torax/tests/test_data/test_timedependence.nc differ