Speeding up tests

- Avoid using newton-raphson when not needed - aka in update tests
- Move slow gradient tests to their own file
- Remove the transport model from the radiation collapse to streamline

PiperOrigin-RevId: 853666131

Author: tamaranorman

torax/_src/orchestration/step_function.py

@@ -142,16 +142,6 @@ def check_for_errors(
       post_processed_outputs: post_processing.PostProcessedOutputs,
   ) -> state.SimError:
     """Checks for errors in the simulation state."""
-    if self._runtime_params_provider.numerics.adaptive_dt:
-      if output_state.solver_numeric_outputs.solver_error_state == 1:
-        # Only check for min dt if the solver did not converge. Else we may have
-        # converged at a dt > min_dt just before we reach min_dt.
-        if (
-            output_state.dt
-            / self._runtime_params_provider.numerics.dt_reduction_factor
-            < self._runtime_params_provider.numerics.min_dt
-        ):
-          return state.SimError.REACHED_MIN_DT
 
     # Low-temperature collapse check
     if output_state.core_profiles.below_minimum_temperature(
@@ -162,8 +152,20 @@ def check_for_errors(
     state_error = output_state.check_for_errors()
     if state_error != state.SimError.NO_ERROR:
       return state_error
-    else:
-      return post_processed_outputs.check_for_errors()
+
+    post_processed_error = post_processed_outputs.check_for_errors()
+    if post_processed_error != state.SimError.NO_ERROR:
+      return post_processed_error
+
+    # Check if reached the minimum time step last - this is often caused by
+    # other errors so check those first to give more informative error messages.
+    if self._runtime_params_provider.numerics.adaptive_dt:
+      if output_state.solver_numeric_outputs.solver_error_state == 1:
+        # If using adaptive stepping and the solver did not converge we must
+        # have reached the minimum time step, so we can exit the simulation.
+        return state.SimError.REACHED_MIN_DT
+
+    return state.SimError.NO_ERROR
 
   @jax.jit
   def __call__(
@@ -298,8 +300,13 @@ def fixed_time_step(
     remaining_dt = dt
 
     def cond(args):
-      remaining_dt, _, _ = args
-      return remaining_dt > constants.CONSTANTS.eps
+      remaining_dt, prev_state, _ = args
+      if self.runtime_params_provider.numerics.adaptive_dt:
+        exit_min_dt = prev_state.solver_numeric_outputs.solver_error_state == 1
+      else:
+        exit_min_dt = False
+      return jnp.logical_and(
+          remaining_dt > constants.CONSTANTS.eps, ~exit_min_dt)
 
     def body(args):
       remaining_dt, prev_state, prev_post_processed = args

torax/_src/orchestration/tests/step_function_jax_transforms_test.py

@@ -0,0 +1,68 @@
+# Copyright 2024 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from absl import logging
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax
+import jax.test_util as jtu
+from torax._src.config import config_loader
+from torax._src.orchestration import run_simulation
+from torax._src.torax_pydantic import interpolated_param_1d
+
+
+class StepFunctionTest(parameterized.TestCase):
+
+  @parameterized.parameters([
+      'basic_config',
+      'iterhybrid_predictor_corrector',
+  ])
+  def test_step_function_grad(self, config_name_no_py):
+    example_config_paths = config_loader.example_config_paths()
+    example_config_path = example_config_paths[config_name_no_py]
+    logging.info('a')
+    cfg = config_loader.build_torax_config_from_file(example_config_path)
+    logging.info('b')
+    (
+        sim_state,
+        post_processed_outputs,
+        step_fn,
+    ) = run_simulation.prepare_simulation(cfg)
+    logging.info('c')
+    params_provider = step_fn.runtime_params_provider
+    input_value = params_provider.profile_conditions.Ip.value
+    logging.info('d')
+
+    @jax.jit
+    def f(override_value):
+      ip_update = interpolated_param_1d.TimeVaryingScalarUpdate(
+          value=override_value
+      )
+      runtime_params_overrides = params_provider.update_provider(
+          lambda x: (x.profile_conditions.Ip,),
+          (ip_update,),
+      )
+      _, new_post_processed_outputs = step_fn(
+          sim_state,
+          post_processed_outputs,
+          runtime_params_overrides=runtime_params_overrides,
+      )
+      return new_post_processed_outputs.Q_fusion
+
+    logging.info('e')
+    jtu.check_grads(f, (input_value,), order=1, modes=('rev',))
+    logging.info('f')
+
+
+if __name__ == '__main__':
+  absltest.main()

torax/_src/orchestration/tests/step_function_test.py

@@ -12,14 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import dataclasses
+import os
 from typing import Any
 
 from absl.testing import absltest
 from absl.testing import parameterized
 import chex
-import jax
 import jax.numpy as jnp
-import jax.test_util as jtu
 import numpy as np
 from torax._src import state
 from torax._src.config import config_loader
@@ -28,6 +27,7 @@
 from torax._src.orchestration import step_function
 from torax._src.output_tools import post_processing
 from torax._src.test_utils import default_configs
+from torax._src.test_utils import paths
 from torax._src.torax_pydantic import interpolated_param_1d
 from torax._src.torax_pydantic import model_config
 
@@ -266,6 +266,24 @@ def test_adaptive_step_with_smaller_passed_max_dt(self):
     )
     self.assertTrue(np.less_equal(output_state.dt, passed_max_dt))
 
+  def test_fixed_step_with_high_density_errors_and_does_not_hang(self):
+    # This test enforces that we exit the fixed step function early if we hit
+    # min_dt. If we don't do this then we risk hanging for a very long time as
+    # we stay at min_dt and the step never seems to make progress. This test
+    # ensures that we don't hang and instead fail early.
+    test_data_dir = paths.test_data_dir()
+    torax_config = config_loader.build_torax_config_from_file(
+        os.path.join(test_data_dir, 'test_iterhybrid_radiation_collapse.py')
+    )
+    sim_state, post_processed_outputs, step_fn = (
+        run_simulation.prepare_simulation(torax_config)
+    )
+    sim_state, post_processed_outputs = step_fn.fixed_time_step(
+        np.array(1.), sim_state, post_processed_outputs)
+
+    sim_error = step_fn.check_for_errors(sim_state, post_processed_outputs)
+    self.assertEqual(sim_error, state.SimError.NAN_DETECTED)
+
   def test_call_with_sawtooth_solver_smoke_test(self):
     """Smoke test for the boolean logic around the sawtooth solver.
 
@@ -338,49 +356,11 @@ def test_fixed_time_step_t_less_than_min_dt(self):
     )
     np.testing.assert_allclose(output_state.dt, 0.01, atol=1e-7)
 
-  @parameterized.parameters([
-      'basic_config',
-      'iterhybrid_predictor_corrector',
-  ])
-  def test_step_function_grad(self, config_name_no_py):
-    example_config_paths = config_loader.example_config_paths()
-    example_config_path = example_config_paths[config_name_no_py]
-    cfg = config_loader.build_torax_config_from_file(example_config_path)
-    (
-        sim_state,
-        post_processed_outputs,
-        step_fn,
-    ) = run_simulation.prepare_simulation(cfg)
-    params_provider = step_fn.runtime_params_provider
-    input_value = params_provider.profile_conditions.Ip.value
-
-    @jax.jit
-    def f(override_value):
-      ip_update = interpolated_param_1d.TimeVaryingScalarUpdate(
-          value=override_value
-      )
-      runtime_params_overrides = params_provider.update_provider(
-          lambda x: (x.profile_conditions.Ip,),
-          (ip_update,),
-      )
-      _, new_post_processed_outputs = step_fn(
-          sim_state,
-          post_processed_outputs,
-          runtime_params_overrides=runtime_params_overrides,
-      )
-      return new_post_processed_outputs.Q_fusion
-
-    jtu.check_grads(f, (input_value,), order=1, modes=('rev',))
-
-  @parameterized.parameters([
-      'iterhybrid_predictor_corrector',
-      'iterhybrid_rampup',
-  ])
-  def test_step_function_overrides(self, config_name_no_py):
-    example_config_paths = config_loader.example_config_paths()
-    example_config_path = example_config_paths[config_name_no_py]
-    raw_config = config_loader.import_module(example_config_path)['CONFIG']
-    cfg = config_loader.build_torax_config_from_file(example_config_path)
+  def test_step_function_overrides(self):
+    original_ip = 15e6
+    config_dict = default_configs.get_default_config_dict()
+    config_dict['profile_conditions']['Ip'] = original_ip
+    cfg = model_config.ToraxConfig.from_dict(config_dict)
     (
         sim_state,
         post_processed_outputs,
@@ -403,10 +383,7 @@ def test_step_function_overrides(self, config_name_no_py):
     )
 
     # Update the config itself and re-run the step.
-    doubled_ip = jax.tree_util.tree_map(
-        lambda x: x * 2.0, raw_config['profile_conditions']['Ip']
-    )
-    cfg.update_fields({'profile_conditions.Ip': doubled_ip})
+    cfg.update_fields({'profile_conditions.Ip': original_ip * 2.0})
     step_fn = run_simulation.make_step_fn(cfg)
     ref_state, ref_post_processed_outputs = step_fn(
         # Use original state and post-processed outputs as the initial value.
@@ -419,13 +396,9 @@ def test_step_function_overrides(self, config_name_no_py):
         override_post_processed_outputs, ref_post_processed_outputs
     )
 
-  @parameterized.parameters([
-      ('iterhybrid_rampup',),
-  ])
-  def test_step_function_geo_overrides(self, config_name_no_py):
-    example_config_paths = config_loader.example_config_paths()
-    example_config_path = example_config_paths[config_name_no_py]
-    cfg = config_loader.build_torax_config_from_file(example_config_path)
+  def test_step_function_geo_overrides(self):
+    config_dict = default_configs.get_default_config_dict()
+    cfg = model_config.ToraxConfig.from_dict(config_dict)
     (
         sim_state,
         post_processed_outputs,

torax/tests/test_data/test_iterhybrid_radiation_collapse.py

@@ -31,6 +31,10 @@
     'W': W_frac,
 }
 CONFIG['plasma_composition']['Z_eff'] = 3.0
+
+# Remove QLKNN transport model to simplify step and avoid QLKNN load.
+CONFIG['transport'] = {}
+
 CONFIG['sources']['impurity_radiation'] = {
     'model_name': 'mavrin_fit',
 }