Exit fixed_step_fn if reach a min_dt when adaptive stepping to avoid hanging

If a min_dt was reached computations could hang as too many steps were performed and the end state was never reached
Other errors such as NAN etc generally cause a MIN_DT on the next step so this is a bit of a catch-all for all the errors

Check for the other errors first as MIN_DT is the least informative and can often be caused by other errors

PiperOrigin-RevId: 852834152

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_test.py

@@ -12,6 +12,7 @@
 # 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
@@ -28,6 +29,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 +268,20 @@ 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):
+    test_data_dir = paths.test_data_dir()
+    config_dict = 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(config_dict)
+    )
+    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.