Merge pull request #1769 from Mani212005:feature/low-temperature-exit

PiperOrigin-RevId: 852282855

Author: Torax team

.gitignore

@@ -29,3 +29,7 @@ _build
 
 # venv
 venv
+.venv
+
+# Python version file
+.python-version

torax/_src/config/numerics.py

@@ -126,6 +126,8 @@ class Numerics(torax_pydantic.BaseModelFrozen):
   adaptive_T_source_prefactor: pydantic.PositiveFloat = 2.0e10
   adaptive_n_source_prefactor: pydantic.PositiveFloat = 2.0e8
 
+  T_minimum_eV: pydantic.PositiveFloat = 5.0
+
   @pydantic.model_validator(mode='after')
   def model_validation(self) -> Self:
     if self.t_initial > self.t_final:

torax/_src/orchestration/step_function.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Logic which controls the stepping over time of the simulation."""
+
 import dataclasses
 import functools
 
@@ -151,6 +152,13 @@ def check_for_errors(
             < 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(
+        self._runtime_params_provider.numerics.T_minimum_eV
+    ):
+      return state.SimError.LOW_TEMPERATURE_COLLAPSE
+
     state_error = output_state.check_for_errors()
     if state_error != state.SimError.NO_ERROR:
       return state_error

torax/_src/state.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Classes defining the TORAX state that evolves over time."""
+
 import dataclasses
 import enum
 import functools
@@ -103,9 +104,7 @@ class CoreProfiles:
   def pressure_thermal_e(self) -> cell_variable.CellVariable:
     """Electron thermal pressure [Pa]."""
     return cell_variable.CellVariable(
-        value=self.n_e.value
-        * self.T_e.value
-        * constants.CONSTANTS.keV_to_J,
+        value=self.n_e.value * self.T_e.value * constants.CONSTANTS.keV_to_J,
         dr=self.n_e.dr,
         right_face_constraint=self.n_e.right_face_constraint
         * self.T_e.right_face_constraint
@@ -167,6 +166,17 @@ def negative_temperature_or_density(self) -> jax.Array:
         ])
     )
 
+  def below_minimum_temperature(self, T_minimum_eV: float) -> bool:
+    """Return True if T_e or T_i is below the minimum temperature threshold."""
+    # Convert eV -> keV since internal storage is keV
+    T_minimum_keV = T_minimum_eV / 1000.0
+
+    is_low_te = jnp.any(self.T_e.value < T_minimum_keV)
+    is_low_ti = jnp.any(self.T_i.value < T_minimum_keV)
+
+    # Use .item() to return a concrete Python boolean
+    return (is_low_te | is_low_ti).item()
+
   def __str__(self) -> str:
     return f"""
       CoreProfiles(
@@ -292,6 +302,7 @@ class SimError(enum.Enum):
   QUASINEUTRALITY_BROKEN = 2
   NEGATIVE_CORE_PROFILES = 3
   REACHED_MIN_DT = 4
+  LOW_TEMPERATURE_COLLAPSE = 5
 
   def log_error(self):
     match self:
@@ -320,6 +331,12 @@ def log_error(self):
             quasineutrality. Check the output file for near-zero temperatures or
             densities at the last valid step.
             """)
+      case SimError.LOW_TEMPERATURE_COLLAPSE:
+        logging.error("""
+          Simulation stopped because ion or electron temperature fell below the
+          configured minimum threshold. This is usually caused by radiative
+          collapse. Output file contains all profiles up to the last valid step.
+          """)
       case SimError.NO_ERROR:
         pass
       case _:

torax/_src/tests/state_test.py

@@ -127,6 +127,54 @@ def test_core_profiles_negative_values_check(self):
       self.assertFalse(new_core_profiles.negative_temperature_or_density())
 
 
+class CoreProfilesTemperatureCheckTest(parameterized.TestCase):
+  """Tests for the below_minimum_temperature method in CoreProfiles."""
+
+  def setUp(self):
+    super().setUp()
+    self.geo = circular_geometry.CircularConfig(n_rho=5).build_geometry()
+    self.base_profiles = core_profile_helpers.make_zero_core_profiles(self.geo)
+
+  @parameterized.named_parameters(
+      ('all_above', 0.2, 0.2, 100.0, False),  # 0.2 keV = 200 eV > 100 eV
+      ('te_below', 0.05, 0.2, 100.0, True),  # 0.05 keV = 50 eV < 100 eV
+      ('ti_below', 0.2, 0.05, 100.0, True),
+      ('both_below', 0.05, 0.05, 100.0, True),
+      # 0.1 keV = 100 eV. Logic is strictly <, so this should pass (False).
+      ('exact_boundary', 0.1, 0.1, 100.0, False),
+  )
+  def test_below_minimum_temperature(
+      self, te_val, ti_val, threshold_ev, expected
+  ):
+    """Verifies below_minimum_temperature returns correct boolean flag."""
+    # Create profiles with constant values across the radius
+    core_profiles = dataclasses.replace(
+        self.base_profiles,
+        T_e=core_profile_helpers.make_constant_core_profile(self.geo, te_val),
+        T_i=core_profile_helpers.make_constant_core_profile(self.geo, ti_val),
+    )
+
+    result = core_profiles.below_minimum_temperature(threshold_ev)
+
+    self.assertIsInstance(result, bool)
+    self.assertEqual(result, expected)
+
+  def test_below_minimum_temperature_mixed_profile(self):
+    """Tests detection when only part of the profile is below threshold."""
+    threshold_ev = 100.0  # 0.1 keV
+
+    te_values = jnp.array([0.2, 0.2, 0.05, 0.2, 0.2])  # 0.05 is below threshold
+    ti_values = jnp.array([0.2, 0.2, 0.2, 0.2, 0.2])
+
+    core_profiles = dataclasses.replace(
+        self.base_profiles,
+        T_e=dataclasses.replace(self.base_profiles.T_e, value=te_values),
+        T_i=dataclasses.replace(self.base_profiles.T_i, value=ti_values),
+    )
+
+    self.assertTrue(core_profiles.below_minimum_temperature(threshold_ev))
+
+
 class ImpurityFractionsTest(parameterized.TestCase):
   """Tests for the impurity_fractions attribute in CoreProfiles."""
 

torax/tests/sim_test.py

@@ -590,6 +590,18 @@ def test_nans_trigger_error(self):
     self.assertEqual(state_history.sim_error, state.SimError.NAN_DETECTED)
     self.assertLess(state_history.times[-1], torax_config.numerics.t_final)
 
+  def test_low_temperature_error(self):
+    """Verify that a config with radiation collapse triggers early stopping and an error."""
+    torax_config = self._get_torax_config(
+        'test_iterhybrid_radiation_collapse.py'
+    )
+    _, state_history = run_simulation.run_simulation(torax_config)
+
+    self.assertEqual(
+        state_history.sim_error, state.SimError.LOW_TEMPERATURE_COLLAPSE
+    )
+    self.assertLess(state_history.times[-1], torax_config.numerics.t_final)
+
   def test_full_output_matches_reference(self):
     """Check for complete output match with reference."""
     torax_config = self._get_torax_config('test_iterhybrid_rampup.py')

torax/tests/test_data/test_iterhybrid_radiation_collapse.py

@@ -0,0 +1,53 @@
+# 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.
+
+"""Based on test_iterhybrid_predictor_corrector but with radiation collapse."""
+
+# pylint: disable=invalid-name
+
+import copy
+
+from torax.tests.test_data import test_iterhybrid_predictor_corrector
+
+W_frac = 1e-3
+CONFIG = copy.deepcopy(test_iterhybrid_predictor_corrector.CONFIG)
+
+assert isinstance(CONFIG['plasma_composition'], dict)
+# increasing density to eventually trigger radiation collapse
+CONFIG['profile_conditions']['nbar'] = {0: 0.8, 5.0: 1.2}
+CONFIG['plasma_composition']['impurity'] = {
+    'Ne': 1 - W_frac,
+    'W': W_frac,
+}
+CONFIG['plasma_composition']['Z_eff'] = 3.0
+CONFIG['sources']['impurity_radiation'] = {
+    'model_name': 'mavrin_fit',
+}
+# Remove fusion source to make collapse easier to attain
+CONFIG['sources'].pop('fusion')
+# Reduce the pedestal temperature to make collapse easier to attain
+CONFIG['pedestal']['T_i_ped'] = 1.0
+CONFIG['pedestal']['T_e_ped'] = 1.0
+# nonlinear solver to enable adaptive dt such that small temperature is reached
+# during the collapse.
+CONFIG['solver'] = (
+    {
+        'solver_type': 'newton_raphson',
+        'use_predictor_corrector': True,
+        'n_corrector_steps': 5,
+        'chi_pereverzev': 30,
+        'D_pereverzev': 15,
+        'use_pereverzev': True,
+    }
+)