Move more config params to ProfileConditions.

After this change, the `Config` object is fully organized. There aren't extraneous params at the top level anymore.

PiperOrigin-RevId: 628058985

torax/config.py

@@ -153,6 +153,17 @@ class ProfileConditions:
   # Set ped top location.
   Ped_top: TimeDependentField = 0.91
 
+  # current profiles (broad "Ohmic" + localized "external" currents)
+  # peaking factor of "Ohmic" current: johm = j0*(1 - r^2/a^2)^nu
+  nu: float = 3.0
+  # toggles if "Ohmic" current is treated as total current upon initialization,
+  # or if non-inductive current should be included in initial jtot calculation
+  initial_j_is_total_current: bool = False
+  # toggles if the initial psi calculation is based on the "nu" current formula,
+  # or from the psi available in the numerical geometry file. This setting is
+  # ignored for the ad-hoc circular geometry, which has no numerical geometry.
+  initial_psi_from_j: bool = False
+
 
 @chex.dataclass
 class Numerics:
@@ -233,17 +244,6 @@ class Config:
   )
   numerics: Numerics = dataclasses.field(default_factory=Numerics)
 
-  # current profiles (broad "Ohmic" + localized "external" currents)
-  # peaking factor of "Ohmic" current: johm = j0*(1 - r^2/a^2)^nu
-  nu: float = 3.0
-  # toggles if "Ohmic" current is treated as total current upon initialization,
-  # or if non-inductive current should be included in initial jtot calculation
-  initial_j_is_total_current: bool = False
-  # toggles if the initial psi calculation is based on the "nu" current formula,
-  # or from the psi available in the numerical geometry file. This setting is
-  # ignored for the ad-hoc circular geometry, which has no numerical geometry.
-  initial_psi_from_j: bool = False
-
   # solver parameters
   solver: SolverConfig = dataclasses.field(default_factory=SolverConfig)
 

torax/config_slice.py

@@ -78,19 +78,6 @@ class DynamicConfigSlice:
   numerics: DynamicNumerics
   sources: Mapping[str, sources_params.DynamicRuntimeParams]
 
-  # current profiles (broad "Ohmic" + localized "external" currents)
-
-  # peaking factor of prescribed (initial) "Ohmic" current:
-  # johm = j0*(1 - r^2/a^2)^nu
-  nu: float
-  # toggles if "Ohmic" current is treated as total current upon initialization,
-  # or if non-inductive current should be included in initial jtot calculation
-  initial_j_is_total_current: bool
-  # toggles if the initial psi calculation is based on the "nu" current formula,
-  # or from the psi available in the numerical geometry file. This setting is
-  # ignored for the ad-hoc circular geometry, which has no numerical geometry.
-  initial_psi_from_j: bool
-
 
 @chex.dataclass(frozen=True)
 class DynamicSolverConfigSlice:
@@ -167,6 +154,18 @@ class DynamicProfileConditions:
   # Set ped top location.
   Ped_top: float
 
+  # current profiles (broad "Ohmic" + localized "external" currents)
+  # peaking factor of prescribed (initial) "Ohmic" current:
+  # johm = j0*(1 - r^2/a^2)^nu
+  nu: float
+  # toggles if "Ohmic" current is treated as total current upon initialization,
+  # or if non-inductive current should be included in initial jtot calculation
+  initial_j_is_total_current: bool
+  # toggles if the initial psi calculation is based on the "nu" current formula,
+  # or from the psi available in the numerical geometry file. This setting is
+  # ignored for the ad-hoc circular geometry, which has no numerical geometry.
+  initial_psi_from_j: bool
+
 
 @chex.dataclass
 class DynamicNumerics:

torax/core_profile_setters.py

@@ -211,10 +211,12 @@ def _prescribe_currents_no_bootstrap(
   )
 
   # construct prescribed current formula on grid.
-  jformula_face = (1 - geo.r_face_norm**2) ** dynamic_config_slice.nu
+  jformula_face = (
+      1 - geo.r_face_norm**2
+  ) ** dynamic_config_slice.profile_conditions.nu
   # calculate total and Ohmic current profiles
   denom = _trapz(jformula_face * geo.spr_face, geo.r_face)
-  if dynamic_config_slice.initial_j_is_total_current:
+  if dynamic_config_slice.profile_conditions.initial_j_is_total_current:
     Ctot = Ip * 1e6 / denom
     jtot_face = jformula_face * Ctot
     johm_face = jtot_face - jext_face
@@ -318,10 +320,12 @@ def _prescribe_currents_with_bootstrap(
   )
 
   # construct prescribed current formula on grid.
-  jformula_face = (1 - geo.r_face_norm**2) ** dynamic_config_slice.nu
+  jformula_face = (
+      1 - geo.r_face_norm**2
+  ) ** dynamic_config_slice.profile_conditions.nu
   denom = _trapz(jformula_face * geo.spr_face, geo.r_face)
   # calculate total and Ohmic current profiles
-  if dynamic_config_slice.initial_j_is_total_current:
+  if dynamic_config_slice.profile_conditions.initial_j_is_total_current:
     Ctot = Ip * 1e6 / denom
     jtot_face = jformula_face * Ctot
     johm_face = jtot_face - jext_face - bootstrap_profile.j_bootstrap_face
@@ -548,7 +552,7 @@ def initial_core_profiles(
   # set up initial psi profile based on current profile
   if (
       isinstance(geo, geometry.CircularGeometry)
-      or dynamic_config_slice.initial_psi_from_j
+      or dynamic_config_slice.profile_conditions.initial_psi_from_j
   ):
     # set up initial current profile without bootstrap current, to get
     # q-profile approximation (needed for bootstrap)
@@ -593,7 +597,7 @@ def initial_core_profiles(
 
   elif (
       isinstance(geo, geometry.CHEASEGeometry)
-      and not dynamic_config_slice.initial_psi_from_j
+      and not dynamic_config_slice.profile_conditions.initial_psi_from_j
   ):
     # psi is already provided from the CHEASE equilibrium, so no need to first
     # calculate currents. However, non-inductive currents are still calculated
@@ -868,9 +872,11 @@ def _get_jtot_hires(
   )
 
   # calculate high resolution jtot and Ohmic current profile
-  jformula_hires = (1 - geo.r_hires_norm**2) ** dynamic_config_slice.nu
+  jformula_hires = (
+      1 - geo.r_hires_norm**2
+  ) ** dynamic_config_slice.profile_conditions.nu
   denom = _trapz(jformula_hires * geo.spr_hires, geo.r_hires)
-  if dynamic_config_slice.initial_j_is_total_current:
+  if dynamic_config_slice.profile_conditions.initial_j_is_total_current:
     Ctot_hires = dynamic_config_slice.profile_conditions.Ip * 1e6 / denom
     jtot_hires = jformula_hires * Ctot_hires
   else:

torax/sources/tests/formulas.py

@@ -52,6 +52,7 @@ def test_custom_exponential_source_can_replace_puff_source(self):
         profile_conditions=config_lib.ProfileConditions(
             set_pedestal=True,
             nbar=0.85,
+            nu=0,
         ),
         numerics=config_lib.Numerics(
             ion_heat_eq=True,
@@ -61,7 +62,6 @@ def test_custom_exponential_source_can_replace_puff_source(self):
             resistivity_mult=100,
             t_final=2,
         ),
-        nu=0,
         solver=config_lib.SolverConfig(
             predictor_corrector=False,
         ),

torax/tests/sim_custom_sources.py

@@ -159,6 +159,7 @@ def custom_source_formula(
         profile_conditions=config_lib.ProfileConditions(
             set_pedestal=True,
             nbar=0.85,
+            nu=0,
         ),
         numerics=config_lib.Numerics(
             ion_heat_eq=True,
@@ -168,7 +169,6 @@ def custom_source_formula(
             resistivity_mult=100,
             t_final=2,
         ),
-        nu=0,
         solver=config_lib.SolverConfig(
             predictor_corrector=False,
         ),

torax/tests/state.py

@@ -186,25 +186,33 @@ def test_initial_psi_from_j(
   ):
     """Tests expected behaviour of initial psi and current options."""
     config1 = config_lib.Config(
-        initial_j_is_total_current=True,
-        initial_psi_from_j=True,
-        nu=2,
+        profile_conditions=config_lib.ProfileConditions(
+            initial_j_is_total_current=True,
+            initial_psi_from_j=True,
+            nu=2,
+        ),
     )
     config2 = config_lib.Config(
-        initial_j_is_total_current=False,
-        initial_psi_from_j=True,
-        nu=2,
+        profile_conditions=config_lib.ProfileConditions(
+            initial_j_is_total_current=False,
+            initial_psi_from_j=True,
+            nu=2,
+        ),
     )
     config3 = config_lib.Config(
-        initial_j_is_total_current=False,
-        initial_psi_from_j=True,
-        nu=2,
+        profile_conditions=config_lib.ProfileConditions(
+            initial_j_is_total_current=False,
+            initial_psi_from_j=True,
+            nu=2,
+        ),
     )
     # Needed to generate psi for bootstrap calculation
     config3_helper = config_lib.Config(
-        initial_j_is_total_current=True,
-        initial_psi_from_j=True,
-        nu=2,
+        profile_conditions=config_lib.ProfileConditions(
+            initial_j_is_total_current=True,
+            initial_psi_from_j=True,
+            nu=2,
+        ),
     )
     geo = geo_builder(config1)
     source_models = source_models_lib.SourceModels()
@@ -325,13 +333,17 @@ def test_initial_psi_from_geo_noop_circular(self):
     """Tests expected behaviour of initial psi and current options."""
     source_models = source_models_lib.SourceModels()
     config1 = config_lib.Config(
-        initial_psi_from_j=False,
+        profile_conditions=config_lib.ProfileConditions(
+            initial_psi_from_j=False,
+        ),
     )
     dcs1 = config_slice.build_dynamic_config_slice(
         config1, sources=source_models.runtime_params
     )
     config2 = config_lib.Config(
-        initial_psi_from_j=True,
+        profile_conditions=config_lib.ProfileConditions(
+            initial_psi_from_j=True,
+        ),
     )
     dcs2 = config_slice.build_dynamic_config_slice(
         config2, sources=source_models.runtime_params

torax/tests/test_data/compilation_benchmark.py

@@ -40,6 +40,7 @@ def get_config() -> config_lib.Config:
           nbar=0.85,  # initial density (Greenwald fraction units)
           ne_bound_right=0.2,
           neped=1.0,
+          nu=0,
       ),
       numerics=config_lib.Numerics(
           ion_heat_eq=True,
@@ -49,7 +50,6 @@ def get_config() -> config_lib.Config:
           resistivity_mult=100,  # to shorten current diffusion time
           t_final=0.0007944 * 2,
       ),
-      nu=0,
       solver=config_lib.SolverConfig(
           use_pereverzev=False,
       ),

torax/tests/test_data/test_absolute_jext.py

@@ -36,15 +36,15 @@ def get_config() -> config_lib.Config:
       profile_conditions=config_lib.ProfileConditions(
           Ti_bound_left=8,
           Te_bound_left=8,
+          # set flat Ohmic current to provide larger range of current evolution
+          # for test
+          nu=0,
       ),
       numerics=config_lib.Numerics(
           current_eq=True,
           resistivity_mult=100,  # to shorten current diffusion time
           t_final=2,
       ),
-      # set flat Ohmic current to provide larger range of current evolution for
-      # test
-      nu=0,
       solver=config_lib.SolverConfig(
           predictor_corrector=False,
           use_pereverzev=True,

torax/tests/test_data/test_all_transport_crank_nicolson.py

@@ -37,6 +37,9 @@ def get_config() -> config_lib.Config:
           nbar=0.85,  # initial density (using Greenwald fraction default)
           ne_bound_right=0.2,
           neped=1.0,
+          # set flat Ohmic current to provide larger range of current evolution
+          # for test
+          nu=0,
       ),
       numerics=config_lib.Numerics(
           ion_heat_eq=True,
@@ -47,9 +50,6 @@ def get_config() -> config_lib.Config:
           t_final=2,
           largeValue_n=1.0e5,
       ),
-      # set flat Ohmic current to provide larger range of current evolution for
-      # test
-      nu=0,
       solver=config_lib.SolverConfig(
           predictor_corrector=False,
           use_pereverzev=True,

torax/tests/test_data/test_all_transport_fusion_qlknn.py

@@ -36,6 +36,9 @@ def get_config() -> config_lib.Config:
           nbar=0.85,  # initial density (Greenwald fraction units)
           ne_bound_right=0.2,
           neped=1.0,
+          # set flat Ohmic current to provide larger range of current evolution
+          # for test
+          nu=0,
       ),
       numerics=config_lib.Numerics(
           ion_heat_eq=True,
@@ -45,9 +48,6 @@ def get_config() -> config_lib.Config:
           resistivity_mult=100,  # to shorten current diffusion time
           t_final=2,
       ),
-      # set flat Ohmic current to provide larger range of current evolution for
-      # test
-      nu=0,
       solver=config_lib.SolverConfig(
           predictor_corrector=False,
           use_pereverzev=True,

torax/tests/test_data/test_bootstrap.py

@@ -32,6 +32,9 @@ def get_config() -> config_lib.Config:
       profile_conditions=config_lib.ProfileConditions(
           set_pedestal=False,
           nbar=0.85,  # initial density (in Greenwald fraction units)
+          # set flat Ohmic current to provide larger range of current evolution
+          # for test
+          nu=0,
       ),
       numerics=config_lib.Numerics(
           ion_heat_eq=True,
@@ -41,9 +44,6 @@ def get_config() -> config_lib.Config:
           resistivity_mult=100,  # to shorten current diffusion time
           t_final=1,
       ),
-      # set flat Ohmic current to provide larger range of current evolution for
-      # test
-      nu=0,
       solver=config_lib.SolverConfig(
           predictor_corrector=False,
       ),

torax/tests/test_data/test_exact_finaltime.py

@@ -29,16 +29,16 @@ def get_config() -> config_lib.Config:
       profile_conditions=config_lib.ProfileConditions(
           Ti_bound_left=8,
           Te_bound_left=8,
+          # set flat Ohmic current to provide larger range of current evolution
+          # for test
+          nu=0,
       ),
       numerics=config_lib.Numerics(
           current_eq=True,
           resistivity_mult=100,  # to shorten current diffusion time
           t_final=2,
           exact_t_final=True,
       ),
-      # set flat Ohmic current to provide larger range of current evolution for
-      # test
-      nu=0,
       solver=config_lib.SolverConfig(
           predictor_corrector=False,
           use_pereverzev=True,

torax/tests/test_data/test_fusion_power.py

@@ -36,6 +36,9 @@ def get_config() -> config_lib.Config:
           nbar=0.85,  # initial density (in Greenwald fraction units)
           ne_bound_right=0.2,
           neped=1.0,
+          # set flat Ohmic current to provide larger range of current evolution
+          # for test
+          nu=0,
       ),
       numerics=config_lib.Numerics(
           ion_heat_eq=True,
@@ -45,9 +48,6 @@ def get_config() -> config_lib.Config:
           resistivity_mult=100,  # to shorten current diffusion time
           t_final=1,
       ),
-      # set flat Ohmic current to provide larger range of current evolution for
-      # test
-      nu=0,
       solver=config_lib.SolverConfig(
           predictor_corrector=False,
           use_pereverzev=True,

torax/tests/test_data/test_ne_qlknn_deff_veff.py

@@ -36,6 +36,9 @@ def get_config() -> config_lib.Config:
           nbar=0.85,  # initial density (in Greenwald fraction units)
           ne_bound_right=0.2,
           neped=1.0,
+          # set flat Ohmic current to provide larger range of current evolution
+          # for test
+          nu=0,
       ),
       numerics=config_lib.Numerics(
           ion_heat_eq=True,
@@ -45,9 +48,6 @@ def get_config() -> config_lib.Config:
           resistivity_mult=100,  # to shorten current diffusion time
           t_final=2,
       ),
-      # set flat Ohmic current to provide larger range of current evolution for
-      # test
-      nu=0,
       solver=config_lib.SolverConfig(
           predictor_corrector=False,
           use_pereverzev=True,