Change build_static_params to take runtime_params separatly

PiperOrigin-RevId: 743963715

Author: tamaranorman

torax/config/build_runtime_params.py

@@ -23,6 +23,9 @@
 DynamicRuntimeParamsSlice and a corresponding geometry with consistent Ip.
 """
 import chex
+from torax.config import numerics as numerics_lib
+from torax.config import plasma_composition as plasma_composition_lib
+from torax.config import profile_conditions as profile_conditions_lib
 from torax.config import runtime_params as general_runtime_params_lib
 from torax.config import runtime_params_slice
 from torax.geometry import geometry
@@ -38,16 +41,23 @@
 
 def build_static_runtime_params_slice(
     *,
-    runtime_params: general_runtime_params_lib.GeneralRuntimeParams,
+    profile_conditions: profile_conditions_lib.ProfileConditions,
+    numerics: numerics_lib.Numerics,
+    plasma_composition: plasma_composition_lib.PlasmaComposition,
     sources: sources_pydantic_model.Sources,
     torax_mesh: torax_pydantic.Grid1D,
     stepper: stepper_pydantic_model.Stepper | None = None,
 ) -> runtime_params_slice.StaticRuntimeParamsSlice:
   """Builds a StaticRuntimeParamsSlice.
 
   Args:
-    runtime_params: General runtime params from which static params are taken,
-      which are the choices on equations being solved, and adaptive dt.
+    profile_conditions: Profile conditions from which the profile conditions
+      static variables are taken, which are the boundary conditions for the
+      plasma.
+    numerics: Numerics from which the numerics static variables are taken, which
+      are the equations being solved, adaptive dt, and the fixed dt.
+    plasma_composition: Plasma composition from which the plasma composition
+      static variables are taken, which are the main and impurity ion names.
     sources: data from which the source related static variables are taken,
       which are the explicit/implicit toggle and calculation mode for each
       source.
@@ -70,14 +80,14 @@ def build_static_runtime_params_slice(
       },
       torax_mesh=torax_mesh,
       stepper=stepper.build_static_params(),
-      ion_heat_eq=runtime_params.numerics.ion_heat_eq,
-      el_heat_eq=runtime_params.numerics.el_heat_eq,
-      current_eq=runtime_params.numerics.current_eq,
-      dens_eq=runtime_params.numerics.dens_eq,
-      main_ion_names=runtime_params.plasma_composition.get_main_ion_names(),
-      impurity_names=runtime_params.plasma_composition.get_impurity_names(),
-      adaptive_dt=runtime_params.numerics.adaptive_dt,
-      use_vloop_lcfs_boundary_condition=runtime_params.profile_conditions.use_vloop_lcfs_boundary_condition,
+      ion_heat_eq=numerics.ion_heat_eq,
+      el_heat_eq=numerics.el_heat_eq,
+      current_eq=numerics.current_eq,
+      dens_eq=numerics.dens_eq,
+      main_ion_names=plasma_composition.get_main_ion_names(),
+      impurity_names=plasma_composition.get_impurity_names(),
+      adaptive_dt=numerics.adaptive_dt,
+      use_vloop_lcfs_boundary_condition=profile_conditions.use_vloop_lcfs_boundary_condition,
   )
 
 

torax/config/tests/runtime_params_slice_test.py

@@ -13,22 +13,34 @@
 # limitations under the License.
 
 
+import copy
+
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
 from torax.config import build_runtime_params
-from torax.config import runtime_params as general_runtime_params
 from torax.config import runtime_params_slice as runtime_params_slice_lib
-from torax.geometry import pydantic_model as geometry_pydantic_model
 from torax.tests.test_lib import default_sources
-from torax.torax_pydantic import torax_pydantic
+from torax.torax_pydantic import model_config
 
 
 class RuntimeParamsSliceTest(parameterized.TestCase):
 
   def setUp(self):
     super().setUp()
-    self._geo = geometry_pydantic_model.CircularConfig().build_geometry()
+    self._python_config = {
+        'runtime_params': {'numerics': {}},
+        'geometry': {
+            'geometry_type': 'circular',
+        },
+        'pedestal': {},
+        'sources': default_sources.get_default_source_config(),
+        'stepper': {},
+        'time_step_calculator': {},
+        'transport': {},
+    }
+    self._torax_config = model_config.ToraxConfig.from_dict(self._python_config)
+    self._torax_mesh = self._torax_config.geometry.build_provider.torax_mesh
 
   def test_dynamic_slice_can_be_input_to_jitted_function(self):
     """Tests that the slice can be input to a jitted function."""
@@ -39,53 +51,56 @@ def foo(
       _ = runtime_params_slice  # do nothing.
 
     foo_jitted = jax.jit(foo)
-    runtime_params = general_runtime_params.GeneralRuntimeParams()
-    torax_pydantic.set_grid(runtime_params, self._geo.torax_mesh)
+    runtime_params = self._torax_config.runtime_params
     dynamic_slice = build_runtime_params.DynamicRuntimeParamsSliceProvider(
         runtime_params,
-        torax_mesh=self._geo.torax_mesh,
+        torax_mesh=self._torax_mesh,
     )(
-        t=runtime_params.numerics.t_initial,
+        t=self._torax_config.numerics.t_initial,
     )
     # Make sure you can call the function with dynamic_slice as an arg.
     foo_jitted(dynamic_slice)
 
   def test_static_runtime_params_slice_hash_same_for_same_params(self):
     """Tests that the hash is the same for the same static params."""
-    runtime_params = general_runtime_params.GeneralRuntimeParams()
-    torax_pydantic.set_grid(runtime_params, self._geo.torax_mesh)
-    sources = default_sources.get_default_sources()
     static_slice1 = build_runtime_params.build_static_runtime_params_slice(
-        runtime_params=runtime_params,
-        sources=sources,
-        torax_mesh=self._geo.torax_mesh,
+        profile_conditions=self._torax_config.profile_conditions,
+        numerics=self._torax_config.numerics,
+        plasma_composition=self._torax_config.plasma_composition,
+        sources=self._torax_config.sources,
+        torax_mesh=self._torax_mesh,
     )
     static_slice2 = build_runtime_params.build_static_runtime_params_slice(
-        runtime_params=runtime_params,
-        sources=sources,
-        torax_mesh=self._geo.torax_mesh,
+        profile_conditions=self._torax_config.profile_conditions,
+        numerics=self._torax_config.numerics,
+        plasma_composition=self._torax_config.plasma_composition,
+        sources=self._torax_config.sources,
+        torax_mesh=self._torax_mesh,
     )
     self.assertEqual(hash(static_slice1), hash(static_slice2))
 
   def test_static_runtime_params_slice_hash_different_for_different_params(
       self,
   ):
     """Test that the hash changes when the static params change."""
-    runtime_params = general_runtime_params.GeneralRuntimeParams()
-    sources = default_sources.get_default_sources()
     static_slice1 = build_runtime_params.build_static_runtime_params_slice(
-        runtime_params=runtime_params,
-        sources=sources,
-        torax_mesh=self._geo.torax_mesh,
+        profile_conditions=self._torax_config.profile_conditions,
+        numerics=self._torax_config.numerics,
+        plasma_composition=self._torax_config.plasma_composition,
+        sources=self._torax_config.sources,
+        torax_mesh=self._torax_mesh,
     )
-    runtime_params_mod = runtime_params.model_copy()
-    runtime_params_mod._update_fields(
-        {'numerics.ion_heat_eq': not runtime_params.numerics.ion_heat_eq}
+    new_config = copy.deepcopy(self._python_config)
+    new_config['runtime_params']['numerics']['ion_heat_eq'] = (
+        not self._torax_config.runtime_params.numerics.ion_heat_eq
     )
+    new_torax_config = model_config.ToraxConfig.from_dict(new_config)
     static_slice2 = build_runtime_params.build_static_runtime_params_slice(
-        runtime_params=runtime_params_mod,
-        sources=sources,
-        torax_mesh=self._geo.torax_mesh,
+        profile_conditions=new_torax_config.profile_conditions,
+        numerics=new_torax_config.numerics,
+        plasma_composition=new_torax_config.plasma_composition,
+        sources=new_torax_config.sources,
+        torax_mesh=self._torax_mesh,
     )
     self.assertNotEqual(hash(static_slice1), hash(static_slice2))
 

torax/core_profiles/tests/getters_test.py

@@ -266,7 +266,9 @@ def test_get_ion_density_and_charge_states(self):
         torax_mesh=self.geo.torax_mesh,
     )
     static_slice = build_runtime_params.build_static_runtime_params_slice(
-        runtime_params=runtime_params,
+        profile_conditions=runtime_params.profile_conditions,
+        numerics=runtime_params.numerics,
+        plasma_composition=runtime_params.plasma_composition,
         sources=sources,
         torax_mesh=self.geo.torax_mesh,
     )

torax/core_profiles/tests/initialization_test.py

@@ -106,7 +106,9 @@ def test_initial_psi(
         t=1.0,
     )
     static_slice = build_runtime_params.build_static_runtime_params_slice(
-        runtime_params=runtime_params,
+        profile_conditions=runtime_params.profile_conditions,
+        numerics=runtime_params.numerics,
+        plasma_composition=runtime_params.plasma_composition,
         sources=sources,
         torax_mesh=self.geo.torax_mesh,
     )

torax/core_profiles/tests/updaters_test.py

@@ -134,7 +134,9 @@ def test_compute_boundary_conditions_ne(
     )
     sources = source_pydantic_model.Sources()
     static_slice = build_runtime_params.build_static_runtime_params_slice(
-        runtime_params=runtime_params,
+        profile_conditions=runtime_params.profile_conditions,
+        numerics=runtime_params.numerics,
+        plasma_composition=runtime_params.plasma_composition,
         sources=sources,
         torax_mesh=self.geo.torax_mesh,
     )
@@ -199,7 +201,9 @@ def test_compute_boundary_conditions_Te(
     )
     sources = source_pydantic_model.Sources.from_dict({})
     static_slice = build_runtime_params.build_static_runtime_params_slice(
-        runtime_params=runtime_params,
+        profile_conditions=runtime_params.profile_conditions,
+        numerics=runtime_params.numerics,
+        plasma_composition=runtime_params.plasma_composition,
         sources=sources,
         torax_mesh=self.geo.torax_mesh,
     )
@@ -250,7 +254,9 @@ def test_compute_boundary_conditions_Ti(
     )
     sources = source_pydantic_model.Sources.from_dict({})
     static_slice = build_runtime_params.build_static_runtime_params_slice(
-        runtime_params=runtime_params,
+        profile_conditions=runtime_params.profile_conditions,
+        numerics=runtime_params.numerics,
+        plasma_composition=runtime_params.plasma_composition,
         sources=sources,
         torax_mesh=self.geo.torax_mesh,
     )

torax/fvm/tests/calc_coeffs_test.py

@@ -15,20 +15,13 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 from torax.config import build_runtime_params
-from torax.config import numerics as numerics_lib
-from torax.config import profile_conditions as profile_conditions_lib
-from torax.config import runtime_params as general_runtime_params
 from torax.core_profiles import initialization
 from torax.fvm import calc_coeffs
-from torax.geometry import pydantic_model as geometry_pydantic_model
-from torax.pedestal_model import pydantic_model as pedestal_pydantic_model
-from torax.sources import pydantic_model as sources_pydantic_model
 from torax.sources import runtime_params as source_runtime_params
 from torax.sources import source_models as source_models_lib
 from torax.sources import source_profile_builders
-from torax.stepper import pydantic_model as stepper_pydantic_model
 from torax.tests.test_lib import default_sources
-from torax.transport_model import pydantic_model as transport_pydantic_model
+from torax.torax_pydantic import model_config
 
 
 class CoreProfileSettersTest(parameterized.TestCase):
@@ -42,66 +35,55 @@ class CoreProfileSettersTest(parameterized.TestCase):
   def test_calc_coeffs_smoke_test(
       self, num_cells, theta_imp, set_pedestal
   ):
-    runtime_params = general_runtime_params.GeneralRuntimeParams(
-        profile_conditions=profile_conditions_lib.ProfileConditions(
-            set_pedestal=set_pedestal,
-        ),
-        numerics=numerics_lib.Numerics(
-            el_heat_eq=False,
-        ),
-    )
-    stepper_params = stepper_pydantic_model.Stepper.from_dict(
-        dict(
-            predictor_corrector=False,
-            theta_imp=theta_imp,
-        )
-    )
-    geo = geometry_pydantic_model.CircularConfig(
-        n_rho=num_cells
-    ).build_geometry()
-
-    transport = transport_pydantic_model.Transport.from_dict(
-        {'transport_model': 'constant', 'chimin': 0, 'chii_const': 1}
-    )
-    pedestal = pedestal_pydantic_model.Pedestal()
-    pedestal_model = pedestal.build_pedestal_model()
-    transport_model = transport.build_transport_model()
-    sources = default_sources.get_default_sources()
-    sources_dict = sources.to_dict()
-    sources_dict['qei_source']['Qei_mult'] = 0.0
-    sources_dict['generic_ion_el_heat_source']['Ptot'] = (
-        0.0
-    )
-    sources_dict['fusion_heat_source']['mode'] = (
+    sources_config = default_sources.get_default_source_config()
+    sources_config['qei_source']['Qei_mult'] = 0.0
+    sources_config['generic_ion_el_heat_source']['Ptot'] = 0.0
+    sources_config['fusion_heat_source']['mode'] = (
         source_runtime_params.Mode.ZERO
     )
-    sources_dict['ohmic_heat_source']['mode'] = (
+    sources_config['ohmic_heat_source']['mode'] = (
         source_runtime_params.Mode.ZERO
     )
-    sources = sources_pydantic_model.Sources.from_dict(sources_dict)
+    torax_config = model_config.ToraxConfig.from_dict(
+        dict(
+            runtime_params=dict(
+                profile_conditions=dict(set_pedestal=set_pedestal),
+                numerics=dict(el_heat_eq=False),
+            ),
+            geometry=dict(geometry_type='circular', n_rho=num_cells),
+            pedestal=dict(),
+            sources=sources_config,
+            stepper=dict(predictor_corrector=False, theta_imp=theta_imp),
+            transport=dict(transport_model='constant', chimin=0, chii_const=1),
+            time_step_calculator=dict(),
+        )
+    )
     source_models = source_models_lib.SourceModels(
-        sources=sources.source_model_config
+        sources=torax_config.sources.source_model_config
     )
     dynamic_runtime_params_slice = (
         build_runtime_params.DynamicRuntimeParamsSliceProvider(
-            runtime_params,
-            transport=transport,
-            sources=sources,
-            stepper=stepper_params,
-            pedestal=pedestal,
-            torax_mesh=geo.torax_mesh,
+            torax_config.runtime_params,
+            transport=torax_config.transport,
+            sources=torax_config.sources,
+            stepper=torax_config.stepper,
+            pedestal=torax_config.pedestal,
+            torax_mesh=torax_config.geometry.build_provider.torax_mesh,
         )(
-            t=runtime_params.numerics.t_initial,
+            t=torax_config.numerics.t_initial,
         )
     )
     static_runtime_params_slice = (
         build_runtime_params.build_static_runtime_params_slice(
-            runtime_params=runtime_params,
-            sources=sources,
-            torax_mesh=geo.torax_mesh,
-            stepper=stepper_params,
+            profile_conditions=torax_config.profile_conditions,
+            numerics=torax_config.numerics,
+            plasma_composition=torax_config.plasma_composition,
+            sources=torax_config.sources,
+            torax_mesh=torax_config.geometry.build_provider.torax_mesh,
+            stepper=torax_config.stepper,
         )
     )
+    geo = torax_config.geometry.build_provider(torax_config.numerics.t_initial)
     core_profiles = initialization.initial_core_profiles(
         static_runtime_params_slice,
         dynamic_runtime_params_slice,
@@ -117,6 +99,8 @@ def test_calc_coeffs_smoke_test(
         core_profiles=core_profiles,
         explicit=True,
     )
+    pedestal_model = torax_config.pedestal.build_pedestal_model()
+    transport_model = torax_config.transport.build_transport_model()
     calc_coeffs.calc_coeffs(
         static_runtime_params_slice=static_runtime_params_slice,
         dynamic_runtime_params_slice=dynamic_runtime_params_slice,