Move static_runtime_params_slice from StepFunction __call__ to constructor.

PiperOrigin-RevId: 775283537

Author: Nush395

torax/_src/mhd/sawtooth/tests/sawtooth_model_test.py

@@ -111,13 +111,11 @@ def setUp(self):
         neoclassical_models=neoclassical_models,
     )
 
-    self.geometry_provider = torax_config.geometry.build_provider
-
-    self.static_runtime_params_slice = (
+    geometry_provider = torax_config.geometry.build_provider
+    static_runtime_params_slice = (
         build_runtime_params.build_static_params_from_config(torax_config)
     )
-
-    self.dynamic_runtime_params_slice_provider = (
+    dynamic_runtime_params_slice_provider = (
         build_runtime_params.DynamicRuntimeParamsSliceProvider.from_config(
             torax_config
         )
@@ -127,24 +125,24 @@ def setUp(self):
         solver=solver,
         time_step_calculator=torax_config.time_step_calculator.time_step_calculator,
         mhd_models=mhd_models,
-        geometry_provider=self.geometry_provider,
-        dynamic_runtime_params_slice_provider=self.dynamic_runtime_params_slice_provider,
+        static_runtime_params_slice=static_runtime_params_slice,
+        geometry_provider=geometry_provider,
+        dynamic_runtime_params_slice_provider=dynamic_runtime_params_slice_provider,
     )
 
     self.initial_state, self.initial_post_processed_outputs = (
         initial_state_lib.get_initial_state_and_post_processed_outputs(
             t=torax_config.numerics.t_initial,
-            static_runtime_params_slice=self.static_runtime_params_slice,
-            dynamic_runtime_params_slice_provider=self.dynamic_runtime_params_slice_provider,
-            geometry_provider=self.geometry_provider,
+            static_runtime_params_slice=static_runtime_params_slice,
+            dynamic_runtime_params_slice_provider=dynamic_runtime_params_slice_provider,
+            geometry_provider=geometry_provider,
             step_fn=self.step_fn,
         )
     )
 
   def test_sawtooth_crash(self):
     """Tests that default values lead to crash and compares post-crash to ref."""
     output_state, _, sim_error = self.step_fn(
-        static_runtime_params_slice=self.static_runtime_params_slice,
         input_state=self.initial_state,
         previous_post_processed_outputs=self.initial_post_processed_outputs,
     )
@@ -180,7 +178,6 @@ def test_no_sawtooth_crash(self):
         ),
     )
     output_state, _, sim_error = self.step_fn(
-        static_runtime_params_slice=self.static_runtime_params_slice,
         input_state=initial_state,
         previous_post_processed_outputs=self.initial_post_processed_outputs,
     )
@@ -197,7 +194,6 @@ def test_no_subsequent_sawtooth_crashes(self):
     """Tests for no subsequent sawtooth crashes even if q in trigger condition."""
     # This crashes
     output_state0, post_processed_outputs0, _ = self.step_fn(
-        static_runtime_params_slice=self.static_runtime_params_slice,
         input_state=self.initial_state,
         previous_post_processed_outputs=self.initial_post_processed_outputs,
     )
@@ -224,7 +220,6 @@ def test_no_subsequent_sawtooth_crashes(self):
 
     with self.subTest('no_subsequent_sawtooth_crashes'):
       output_state_should_not_crash, _, sim_error = self.step_fn(
-          static_runtime_params_slice=self.static_runtime_params_slice,
           input_state=new_input_state_should_not_crash,
           previous_post_processed_outputs=post_processed_outputs0,
       )
@@ -243,7 +238,6 @@ def test_no_subsequent_sawtooth_crashes(self):
 
     with self.subTest('crashes_if_sawtooth_crash_is_false'):
       output_state_should_crash, _, sim_error = self.step_fn(
-          static_runtime_params_slice=self.static_runtime_params_slice,
           input_state=new_input_state_should_crash,
           previous_post_processed_outputs=post_processed_outputs0,
       )

torax/_src/orchestration/run_loop.py

@@ -152,7 +152,6 @@ def run_loop(
         _log_timestep(current_state)
 
       current_state, post_processed_outputs, sim_error = step_fn(
-          static_runtime_params_slice,
           current_state,
           post_processing_history[-1],
       )

torax/_src/orchestration/run_simulation.py

@@ -102,6 +102,7 @@ def prepare_simulation(
       solver=solver,
       time_step_calculator=torax_config.time_step_calculator.time_step_calculator,
       mhd_models=mhd_models,
+      static_runtime_params_slice=static_runtime_params_slice,
       geometry_provider=geometry_provider,
       dynamic_runtime_params_slice_provider=dynamic_runtime_params_slice_provider,
   )

torax/_src/orchestration/step_function.py

@@ -76,6 +76,7 @@ def __init__(
       solver: solver_lib.Solver,
       time_step_calculator: ts.TimeStepCalculator,
       mhd_models: mhd_base.MHDModels,
+      static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
       dynamic_runtime_params_slice_provider: build_runtime_params.DynamicRuntimeParamsSliceProvider,
       geometry_provider: geometry_provider_lib.GeometryProvider,
   ):
@@ -84,7 +85,9 @@ def __init__(
     Args:
       solver: Evolves the core profiles.
       time_step_calculator: Calculates the dt for each time step.
-      mhd_models: Collection of MHD models applied, e.g. sawtooth
+      mhd_models: Collection of MHD models applied, e.g. sawtooth.
+      static_runtime_params_slice: Static parameters that, if they change,
+        should trigger a recompilation of the SimulationStepFn.
       dynamic_runtime_params_slice_provider: Object that returns a set of
         runtime parameters which may change from time step to time step or
         simulation run to run. If these runtime parameters change, it does NOT
@@ -105,6 +108,7 @@ def __init__(
     self._dynamic_runtime_params_slice_provider = (
         dynamic_runtime_params_slice_provider
     )
+    self._static_runtime_params_slice = static_runtime_params_slice
 
   @property
   def solver(self) -> solver_lib.Solver:
@@ -120,7 +124,6 @@ def time_step_calculator(self) -> ts.TimeStepCalculator:
 
   def __call__(
       self,
-      static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
       input_state: sim_state.ToraxSimState,
       previous_post_processed_outputs: post_processing.PostProcessedOutputs,
   ) -> tuple[
@@ -136,8 +139,6 @@ def __call__(
     sawtooth redistribution, at a t+dt set by the sawtooth model.
 
     Args:
-      static_runtime_params_slice: Static parameters that, if they change,
-        should trigger a recompilation of the SimulationStepFn.
       input_state: State at the start of the time step, including the core
         profiles which are being evolved.
       previous_post_processed_outputs: Post-processed outputs from the previous
@@ -172,7 +173,7 @@ def __call__(
     # set to 0.
     explicit_source_profiles = source_profile_builders.build_source_profiles(
         dynamic_runtime_params_slice=dynamic_runtime_params_slice_t,
-        static_runtime_params_slice=static_runtime_params_slice,
+        static_runtime_params_slice=self._static_runtime_params_slice,
         geo=geo_t,
         core_profiles=input_state.core_profiles,
         source_models=self.solver.source_models,
@@ -209,7 +210,7 @@ def __call__(
       # previous_post_processed_outputs.
       output_state, post_processed_outputs = _sawtooth_step(
           sawtooth_solver=self.mhd_models.sawtooth,
-          static_runtime_params_slice=static_runtime_params_slice,
+          static_runtime_params_slice=self._static_runtime_params_slice,
           dynamic_runtime_params_slice_t=dynamic_runtime_params_slice_t,
           dynamic_runtime_params_slice_t_plus_crash_dt=dynamic_runtime_params_slice_t_plus_crash_dt,
           geo_t=geo_t,
@@ -246,7 +247,6 @@ def __call__(
 
     x_new, solver_numeric_outputs = self.step(
         dt,
-        static_runtime_params_slice,
         dynamic_runtime_params_slice_t,
         dynamic_runtime_params_slice_t_plus_dt,
         geo_t,
@@ -255,7 +255,7 @@ def __call__(
         explicit_source_profiles,
     )
 
-    if static_runtime_params_slice.adaptive_dt:
+    if self._static_runtime_params_slice.adaptive_dt:
       # This is a no-op if
       # output_state.solver_numeric_outputs.solver_error_state == 0.
       (
@@ -268,7 +268,6 @@ def __call__(
           x_new,
           dt,
           solver_numeric_outputs,
-          static_runtime_params_slice,
           dynamic_runtime_params_slice_t,
           dynamic_runtime_params_slice_t_plus_dt,
           geo_t,
@@ -282,7 +281,7 @@ def __call__(
     # the dt may be adaptive, we need to recompute it here.
     core_profiles_t_plus_dt = updaters.provide_core_profiles_t_plus_dt(
         dt=dt,
-        static_runtime_params_slice=static_runtime_params_slice,
+        static_runtime_params_slice=self._static_runtime_params_slice,
         dynamic_runtime_params_slice_t=dynamic_runtime_params_slice_t,
         dynamic_runtime_params_slice_t_plus_dt=dynamic_runtime_params_slice_t_plus_dt,
         geo_t_plus_dt=geo_t_plus_dt,
@@ -294,7 +293,7 @@ def __call__(
         dt=dt,
         x_new=x_new,
         solver_numeric_outputs=solver_numeric_outputs,
-        static_runtime_params_slice=self.solver.static_runtime_params_slice,
+        static_runtime_params_slice=self._static_runtime_params_slice,
         dynamic_runtime_params_slice_t_plus_dt=dynamic_runtime_params_slice_t_plus_dt,
         geometry_t_plus_dt=geo_t_plus_dt,
         explicit_source_profiles=explicit_source_profiles,
@@ -377,7 +376,6 @@ def init_time_step_calculator(
   def step(
       self,
       dt: jax.Array,
-      static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
       dynamic_runtime_params_slice_t: runtime_params_slice.DynamicRuntimeParamsSlice,
       dynamic_runtime_params_slice_t_plus_dt: runtime_params_slice.DynamicRuntimeParamsSlice,
       geo_t: geometry.Geometry,
@@ -395,8 +393,6 @@ def step(
 
     Args:
       dt: Time step duration.
-      static_runtime_params_slice: Static parameters that, if they change,
-        should trigger a recompilation of the SimulationStepFn.
       dynamic_runtime_params_slice_t: Runtime parameters at time t.
       dynamic_runtime_params_slice_t_plus_dt: Runtime parameters at time t + dt.
       geo_t: The geometry of the torus during this time step of the simulation.
@@ -421,7 +417,7 @@ def step(
     # PDE system.
     core_profiles_t_plus_dt = updaters.provide_core_profiles_t_plus_dt(
         dt=dt,
-        static_runtime_params_slice=static_runtime_params_slice,
+        static_runtime_params_slice=self._static_runtime_params_slice,
         dynamic_runtime_params_slice_t=dynamic_runtime_params_slice_t,
         dynamic_runtime_params_slice_t_plus_dt=dynamic_runtime_params_slice_t_plus_dt,
         geo_t_plus_dt=geo_t_plus_dt,
@@ -447,7 +443,6 @@ def _adaptive_step(
       x_old: tuple[cell_variable.CellVariable, ...],
       dt: jax.Array,
       solver_numeric_outputs: state.SolverNumericOutputs,
-      static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
       dynamic_runtime_params_slice_t: runtime_params_slice.DynamicRuntimeParamsSlice,
       dynamic_runtime_params_slice_t_plus_dt: runtime_params_slice.DynamicRuntimeParamsSlice,
       geo_t: geometry.Geometry,
@@ -473,8 +468,6 @@ def _adaptive_step(
       dt: Time step duration for the initial step.
       solver_numeric_outputs: Solver-specific numeric outputs from the initial
         step.
-      static_runtime_params_slice: Static parameters that, if they change,
-        should trigger a recompilation of the SimulationStepFn.
       dynamic_runtime_params_slice_t: Runtime parameters at time t.
       dynamic_runtime_params_slice_t_plus_dt: Runtime parameters at time t +
         dt.
@@ -558,7 +551,7 @@ def body_fun(
 
       core_profiles_t_plus_dt = updaters.provide_core_profiles_t_plus_dt(
           dt=dt,
-          static_runtime_params_slice=static_runtime_params_slice,
+          static_runtime_params_slice=self._static_runtime_params_slice,
           dynamic_runtime_params_slice_t=dynamic_runtime_params_slice_t,
           dynamic_runtime_params_slice_t_plus_dt=dynamic_runtime_params_slice_t_plus_dt,
           geo_t_plus_dt=geo_t_plus_dt,

torax/tests/sim_time_dependence_test.py

@@ -124,9 +124,9 @@ def _fake_run_loop(
           restart_case,
           geometry_provider,
           dynamic_runtime_params_slice_provider,
+          static_runtime_params_slice,
       )
       output_state, post_processed_outputs, error = step_fn(
-          static_runtime_params_slice,
           initial_state,
           initial_post_processed_outputs,
       )