Add enable_fast_ions toggle to numerics config

PiperOrigin-RevId: 879030506

Author: hamelphi

torax/_src/config/numerics.py

@@ -53,6 +53,7 @@ class RuntimeParams:
   evolve_density: bool = dataclasses.field(metadata={'static': True})
   exact_t_final: bool = dataclasses.field(metadata={'static': True})
   adaptive_dt: bool = dataclasses.field(metadata={'static': True})
+  enable_fast_ions: bool = dataclasses.field(metadata={'static': True})
 
   @functools.cached_property
   def evolving_names(self) -> tuple[str, ...]:
@@ -128,6 +129,7 @@ class Numerics(torax_pydantic.BaseModelFrozen):
   evolve_electron_heat: Annotated[bool, torax_pydantic.JAX_STATIC] = True
   evolve_current: Annotated[bool, torax_pydantic.JAX_STATIC] = False
   evolve_density: Annotated[bool, torax_pydantic.JAX_STATIC] = False
+  enable_fast_ions: Annotated[bool, torax_pydantic.JAX_STATIC] = False
   resistivity_multiplier: torax_pydantic.TimeVaryingScalar = (
       torax_pydantic.ValidatedDefault(1.0)
   )
@@ -188,4 +190,5 @@ def build_runtime_params(self, t: chex.Numeric) -> RuntimeParams:
         evolve_density=self.evolve_density,
         exact_t_final=self.exact_t_final,
         adaptive_dt=self.adaptive_dt,
+        enable_fast_ions=self.enable_fast_ions,
     )

torax/_src/core_profiles/initialization.py

@@ -99,8 +99,9 @@ def initial_core_profiles(
   )
 
   fast_ions_list = []
-  for s in source_models.standard_sources.values():
-    fast_ions_list.extend(s.zero_fast_ions(geo))
+  if runtime_params.numerics.enable_fast_ions:
+    for s in source_models.standard_sources.values():
+      fast_ions_list.extend(s.zero_fast_ions(geo))
 
   core_profiles = state.CoreProfiles(
       T_i=T_i,

torax/_src/sources/source.py

@@ -222,34 +222,55 @@ def get_value(
             conductivity,
         )
       case sources_runtime_params_lib.Mode.PRESCRIBED:
-        if len(self.affected_core_profiles) != len(
-            source_params.prescribed_values
+        expected_len = len(self.affected_core_profiles)
+        prescribed_len = len(source_params.prescribed_values)
+        if (
+            AffectedCoreProfile.FAST_IONS in self.affected_core_profiles
+            and not runtime_params.numerics.enable_fast_ions
+            and prescribed_len == expected_len - 1
         ):
+          fast_ions_idx = self.affected_core_profiles.index(
+              AffectedCoreProfile.FAST_IONS
+          )
+          res_list = list(source_params.prescribed_values)
+          res_list.insert(fast_ions_idx, ())
+          res = tuple(res_list)
+        elif prescribed_len != expected_len:
           raise ValueError(
               'When using PRESCRIBED mode, the number of prescribed values must'
               ' match the number of affected core profiles. Was: '
               f'{len(source_params.prescribed_values)} '
               f' Expected: {len(self.affected_core_profiles)}.'
           )
-        res = source_params.prescribed_values
+        else:
+          res = source_params.prescribed_values
       case sources_runtime_params_lib.Mode.ZERO:
         zeros = jnp.zeros(geo.rho_norm.shape)
-        res = tuple(
-            self.zero_fast_ions(geo)
-            if acp == AffectedCoreProfile.FAST_IONS
-            else zeros
-            for acp in self.affected_core_profiles
-        )
+        res_list = []
+        for affected_core_profile in self.affected_core_profiles:
+          if affected_core_profile == AffectedCoreProfile.FAST_IONS:
+            if runtime_params.numerics.enable_fast_ions:
+              res_list.append(self.zero_fast_ions(geo))
+            else:
+              res_list.append(())
+          else:
+            res_list.append(zeros)
+        res = tuple(res_list)
       case _:
         raise ValueError(f'Unknown mode: {mode}')
 
     if AffectedCoreProfile.FAST_IONS in self.affected_core_profiles:
       fast_ions_idx = self.affected_core_profiles.index(
           AffectedCoreProfile.FAST_IONS
       )
-      self._validate_fast_ions(
-          res[fast_ions_idx],
-          geo,
-      )
+      if runtime_params.numerics.enable_fast_ions:
+        self._validate_fast_ions(
+            res[fast_ions_idx],
+            geo,
+        )
+      elif res[fast_ions_idx]:
+        res_list = list(res)
+        res_list[fast_ions_idx] = ()
+        res = tuple(res_list)
 
     return res

torax/_src/sources/tests/source_test.py

@@ -18,6 +18,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import numpy as np
+from torax._src.config import numerics
 from torax._src.config import runtime_params as runtime_params_lib
 from torax._src.geometry import geometry
 from torax._src.sources import electron_cyclotron_source
@@ -97,6 +98,9 @@ def test_correct_mode_called(
     dynamic_slice = mock.create_autospec(
         runtime_params_lib.RuntimeParams,
         sources=dynamic_source_params,
+        numerics=mock.create_autospec(
+            numerics.RuntimeParams, enable_fast_ions=True
+        ),
     )
     # Make a geo with rho_norm as we need it for the zero profile shape.
     geo = mock.create_autospec(
@@ -130,6 +134,9 @@ def test_prescribed_values_for_multiple_affected_profiles(self):
     dynamic_slice = mock.create_autospec(
         runtime_params_lib.RuntimeParams,
         sources=dynamic_source_params,
+        numerics=mock.create_autospec(
+            numerics.RuntimeParams, enable_fast_ions=True
+        ),
     )
     profile = source.get_value(
         runtime_params=dynamic_slice,
@@ -166,6 +173,9 @@ def test_source_with_mismatched_prescribed_values_raises_error(self):
     dynamic_slice = mock.create_autospec(
         runtime_params_lib.RuntimeParams,
         sources=dynamic_source_params,
+        numerics=mock.create_autospec(
+            numerics.RuntimeParams, enable_fast_ions=True
+        ),
     )
     with self.assertRaisesRegex(
         ValueError,