Create Static container for numerics and package components of slice into it.

This is to package all the current components of `Numerics` that go into the static slice into a single object (that future attributes can be easily added to as well). This also keeps consistency with the other object patterns (e.g. `profile_conditions.StaticRuntimeParams`)

PiperOrigin-RevId: 775289301

Author: Nush395

torax/_src/config/build_runtime_params.py

@@ -42,13 +42,9 @@ def build_static_params_from_config(
       },
       torax_mesh=config.geometry.build_provider.torax_mesh,
       solver=config.solver.build_static_params(),
-      evolve_ion_heat=config.numerics.evolve_ion_heat,
-      evolve_electron_heat=config.numerics.evolve_electron_heat,
-      evolve_current=config.numerics.evolve_current,
-      evolve_density=config.numerics.evolve_density,
       main_ion_names=config.plasma_composition.get_main_ion_names(),
       impurity_names=config.plasma_composition.get_impurity_names(),
-      adaptive_dt=config.numerics.adaptive_dt,
+      numerics=config.numerics.build_static_params(),
       profile_conditions=config.profile_conditions.build_static_params(),
   )
 

torax/_src/config/numerics.py

@@ -24,7 +24,10 @@
 # TODO(b/326578331): remove density reference from DynamicNumerics entirely.
 @chex.dataclass
 class DynamicNumerics:
-  """Generic numeric parameters for the simulation."""
+  """Generic numeric parameters for the simulation.
+
+  For definitions see `Numerics`.
+  """
 
   t_initial: float
   t_final: float
@@ -40,6 +43,19 @@ class DynamicNumerics:
   calcphibdot: bool
 
 
+@chex.dataclass(frozen=True)
+class StaticNumerics:
+  """Static numerics parameters for the simulation.
+
+  For definitions see `Numerics`.
+  """
+  evolve_ion_heat: bool
+  evolve_electron_heat: bool
+  evolve_current: bool
+  evolve_density: bool
+  adaptive_dt: bool
+
+
 class Numerics(torax_pydantic.BaseModelFrozen):
   """Generic numeric parameters for the simulation.
 
@@ -119,7 +135,7 @@ def build_dynamic_params(
       self,
       t: chex.Numeric,
   ) -> DynamicNumerics:
-    """Builds a DynamicNumerics."""
+    """Builds a DynamicNumerics object for time t."""
     return DynamicNumerics(
         t_initial=self.t_initial,
         t_final=self.t_final,
@@ -134,3 +150,13 @@ def build_dynamic_params(
         adaptive_T_source_prefactor=self.adaptive_T_source_prefactor,
         adaptive_n_source_prefactor=self.adaptive_n_source_prefactor,
     )
+
+  def build_static_params(self) -> StaticNumerics:
+    """Builds a StaticNumerics object."""
+    return StaticNumerics(
+        evolve_ion_heat=self.evolve_ion_heat,
+        evolve_electron_heat=self.evolve_electron_heat,
+        evolve_current=self.evolve_current,
+        evolve_density=self.evolve_density,
+        adaptive_dt=self.adaptive_dt,
+    )

torax/_src/config/runtime_params_slice.py

@@ -113,15 +113,6 @@ class StaticRuntimeParamsSlice:
   sources: Mapping[str, sources_params.StaticRuntimeParams]
   # Torax mesh used to construct the geometry.
   torax_mesh: torax_pydantic.Grid1D
-  # Solve the ion heat equation (ion temperature evolves over time)
-  evolve_ion_heat: bool
-  # Solve the electron heat equation (electron temperature evolves over time)
-  evolve_electron_heat: bool
-  # Solve the current equation (psi evolves over time driven by the solver;
-  # q and s evolve over time as a function of psi)
-  evolve_current: bool
-  # Solve the density equation (n evolves over time)
-  evolve_density: bool
   # Ion symbols for main ion and impurity (which each could be mixtures of ions)
   # These are static to simplify source functions for fusion power and radiation
   # which are species-dependent.
@@ -130,24 +121,17 @@ class StaticRuntimeParamsSlice:
   main_ion_names: tuple[str, ...]
   impurity_names: tuple[str, ...]
   profile_conditions: profile_conditions.StaticRuntimeParams
-  # Iterative reduction of dt if nonlinear step does not converge,
-  # If nonlinear step does not converge, then the step is redone
-  # iteratively at successively lower dt until convergence is reached
-  adaptive_dt: bool
+  numerics: numerics.StaticNumerics
 
   def __hash__(self):
     return hash((
         self.solver,
         tuple(sorted(self.sources.items())),  # Hashable version of sources
         hash(self.torax_mesh),  # Grid1D has a hash method defined.
-        self.evolve_ion_heat,
-        self.evolve_electron_heat,
-        self.evolve_current,
-        self.evolve_density,
         self.main_ion_names,
         self.impurity_names,
-        self.adaptive_dt,
         self.profile_conditions,
+        self.numerics,
     ))
 
   def validate_new(self, new_params: typing_extensions.Self):

torax/_src/core_profiles/updaters.py

@@ -91,21 +91,21 @@ def get_prescribed_core_profile_values(
   """
   # If profiles are not evolved, they can still potential be time-evolving,
   # depending on the runtime params. If so, they are updated below.
-  if not static_runtime_params_slice.evolve_ion_heat:
+  if not static_runtime_params_slice.numerics.evolve_ion_heat:
     T_i = getters.get_updated_ion_temperature(
         dynamic_runtime_params_slice.profile_conditions, geo
     ).value
   else:
     T_i = core_profiles.T_i.value
-  if not static_runtime_params_slice.evolve_electron_heat:
+  if not static_runtime_params_slice.numerics.evolve_electron_heat:
     T_e_cell_variable = getters.get_updated_electron_temperature(
         dynamic_runtime_params_slice.profile_conditions, geo
     )
     T_e = T_e_cell_variable.value
   else:
     T_e_cell_variable = core_profiles.T_e
     T_e = T_e_cell_variable.value
-  if not static_runtime_params_slice.evolve_density:
+  if not static_runtime_params_slice.numerics.evolve_density:
     n_e_cell_variable = getters.get_updated_electron_density(
         static_runtime_params_slice,
         dynamic_runtime_params_slice.profile_conditions,

torax/_src/mhd/sawtooth/simple_redistribution.py

@@ -65,7 +65,7 @@ def __call__(
     """
 
     # No sawtooth redistribution if current is not being evolved.
-    if not static_runtime_params_slice.evolve_current:
+    if not static_runtime_params_slice.numerics.evolve_current:
       return core_profiles_t
 
     assert dynamic_runtime_params_slice.mhd.sawtooth is not None
@@ -88,7 +88,7 @@ def __call__(
     indices = jnp.arange(geo.rho_norm.shape[0])
     redistribution_mask = indices < idx_mixing
 
-    if static_runtime_params_slice.evolve_density:
+    if static_runtime_params_slice.numerics.evolve_density:
       n_e_redistributed = flatten_profile.flatten_density_profile(
           rho_norm_q1,
           mixing_radius,
@@ -99,7 +99,7 @@ def __call__(
       )
     else:
       n_e_redistributed = core_profiles_t.n_e
-    if static_runtime_params_slice.evolve_electron_heat:
+    if static_runtime_params_slice.numerics.evolve_electron_heat:
       te_redistributed = flatten_profile.flatten_temperature_profile(
           rho_norm_q1,
           mixing_radius,
@@ -113,8 +113,8 @@ def __call__(
     else:
       te_redistributed = core_profiles_t.T_e
     if (
-        static_runtime_params_slice.evolve_density
-        or static_runtime_params_slice.evolve_electron_heat
+        static_runtime_params_slice.numerics.evolve_density
+        or static_runtime_params_slice.numerics.evolve_electron_heat
     ):
       ions_redistributed = getters.get_updated_ions(
           static_runtime_params_slice,
@@ -136,7 +136,7 @@ def __call__(
           Z_eff=core_profiles_t.Z_eff,
           Z_eff_face=core_profiles_t.Z_eff_face,
       )
-    if static_runtime_params_slice.evolve_ion_heat:
+    if static_runtime_params_slice.numerics.evolve_ion_heat:
       ti_redistributed = flatten_profile.flatten_temperature_profile(
           rho_norm_q1,
           mixing_radius,

torax/_src/orchestration/step_function.py

@@ -255,7 +255,7 @@ def __call__(
         explicit_source_profiles,
     )
 
-    if self._static_runtime_params_slice.adaptive_dt:
+    if self._static_runtime_params_slice.numerics.adaptive_dt:
       # This is a no-op if
       # output_state.solver_numeric_outputs.solver_error_state == 0.
       (

torax/_src/solver/solver.py

@@ -88,13 +88,13 @@ def __eq__(self, other: typing_extensions.Self) -> bool:
   def evolving_names(self) -> tuple[str, ...]:
     """The names of core_profiles variables that are evolved by the solver."""
     evolving_names = []
-    if self.static_runtime_params_slice.evolve_ion_heat:
+    if self.static_runtime_params_slice.numerics.evolve_ion_heat:
       evolving_names.append('T_i')
-    if self.static_runtime_params_slice.evolve_electron_heat:
+    if self.static_runtime_params_slice.numerics.evolve_electron_heat:
       evolving_names.append('T_e')
-    if self.static_runtime_params_slice.evolve_current:
+    if self.static_runtime_params_slice.numerics.evolve_current:
       evolving_names.append('psi')
-    if self.static_runtime_params_slice.evolve_density:
+    if self.static_runtime_params_slice.numerics.evolve_density:
       evolving_names.append('n_e')
     return tuple(evolving_names)
 

torax/_src/sources/qei_source.py

@@ -119,12 +119,12 @@ def _model_based_qei(
   if (
       # if only a single heat equation is being evolved
       (
-          static_runtime_params_slice.evolve_ion_heat
-          and not static_runtime_params_slice.evolve_electron_heat
+          static_runtime_params_slice.numerics.evolve_ion_heat
+          and not static_runtime_params_slice.numerics.evolve_electron_heat
       )
       or (
-          static_runtime_params_slice.evolve_electron_heat
-          and not static_runtime_params_slice.evolve_ion_heat
+          static_runtime_params_slice.numerics.evolve_electron_heat
+          and not static_runtime_params_slice.numerics.evolve_ion_heat
       )
   ):
     explicit_i = qei_coef * core_profiles.T_e.value