Switch while_loop_bounded to use implementation='while_loop'.

PiperOrigin-RevId: 950859512

Author: sbodenstein

torax/_src/orchestration/jit_run_loop.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """JITted run_loop for iterating over the simulation step function."""
+
 import chex
 import jax
 import jax.numpy as jnp
@@ -62,7 +63,7 @@ def _step_fn(inputs):
       _step_fn,
       (initial_state, initial_post_processed_outputs),
       max_steps,
-      implementation='scan',
+      implementation='while_loop',
   )
 
   # Prepend initial state to give (max_steps + 1, ...) output.
@@ -119,6 +120,8 @@ def run_loop(
     runtime_params_overrides: (
         build_runtime_params.RuntimeParamsProvider | None
     ) = None,
+    log_timestep_info: bool = False,
+    progress_bar: bool = True,
     max_steps: int | None = None,
 ) -> tuple[
     list[sim_state.SimState],
@@ -135,6 +138,9 @@ def run_loop(
       is). The state_history that run_simulation() outputs comes from these
       ToraxSimState objects.
     runtime_params_overrides: Optional runtime params overrides to use.
+    log_timestep_info: If True, logs basic timestep info, like time, dt, on
+      every step.
+    progress_bar: If True, displays a progress bar.
     max_steps: Optional maximum number of steps to take. If not provided, then
       the maximum number of steps will be determined by the numerics.t_final and
       numerics.min_dt.
@@ -154,6 +160,16 @@ def run_loop(
         the last valid timestep.
       - The sim error state.
   """
+
+  if progress_bar:
+    raise NotImplementedError(
+        'Progress bar is not supported with the jitted run loop.'
+    )
+  if log_timestep_info:
+    raise NotImplementedError(
+        'Log timestep info is not supported with the jitted run loop.'
+    )
+
   numerics = step_fn.runtime_params_provider.numerics
   if max_steps is None:
     max_steps = int(

torax/_src/orchestration/run_simulation.py

@@ -105,6 +105,7 @@ def run_simulation(
     log_timestep_info: bool = False,
     progress_bar: bool = True,
     max_steps: int | None = None,
+    _use_jitted_run_loop: bool = False,  # pylint: disable=invalid-name
 ) -> tuple[xr.DataTree, output.StateHistory]:
   """Runs a TORAX simulation using the config and returns the outputs.
 
@@ -114,6 +115,8 @@ def run_simulation(
     progress_bar: Whether to show a progress bar.
     max_steps: The maximum number of steps to take, if not provided, then the
       simulation will run until the maximum time is reached.
+    _use_jitted_run_loop: If True, then a jitted run loop will be used. A
+      temporary private argument used for testing.
 
   Returns:
     A tuple of the simulation outputs in the form of a DataTree and the state
@@ -128,14 +131,27 @@ def run_simulation(
       step_fn,
   ) = prepare_simulation(torax_config)
 
-  state_history, post_processed_outputs_history, sim_error = run_loop.run_loop(
-      initial_state=initial_state,
-      initial_post_processed_outputs=post_processed_outputs,
-      step_fn=step_fn,
-      log_timestep_info=log_timestep_info,
-      progress_bar=progress_bar,
-      max_steps=max_steps,
-  )
+  if _use_jitted_run_loop:
+
+    state_history, post_processed_outputs_history, sim_error = (
+        jit_run_loop.run_loop(
+            step_fn,
+            max_steps=max_steps,
+            log_timestep_info=log_timestep_info,
+            progress_bar=progress_bar,
+        )
+    )
+  else:
+    state_history, post_processed_outputs_history, sim_error = (
+        run_loop.run_loop(
+            initial_state=initial_state,
+            initial_post_processed_outputs=post_processed_outputs,
+            step_fn=step_fn,
+            log_timestep_info=log_timestep_info,
+            progress_bar=progress_bar,
+            max_steps=max_steps,
+        )
+    )
 
   state_history = output.StateHistory(
       state_history=state_history,
@@ -148,45 +164,3 @@ def run_simulation(
       state_history.simulation_output_to_xr(),
       state_history,
   )
-
-
-def run_simulation_jitted(
-    torax_config: model_config.ToraxConfig,
-    max_steps: int | None = None,
-) -> tuple[xr.DataTree, output.StateHistory]:
-  """Runs a TORAX simulation using the config and returns the outputs.
-
-  NOTE: This function doesn't guarantee that the simulation will complete. If
-  the simulation does not complete successfully, then the state history will
-  contain the error state `SimError.DID_NOT_REACH_T_FINAL` and a truncated
-  simulation history.
-
-  Args:
-    torax_config: The TORAX config to use for the simulation.
-    max_steps: The maximum number of steps to take, if not provided, then the
-      maximum number of steps will be determined by the numerics.t_final and
-      numerics.min_dt.
-
-  Returns:
-    A tuple of the simulation outputs in the form of a DataTree and the state
-    history which is intended for helpful use with debugging as it contains
-    the `CoreProfiles`, `CoreTransport`, `CoreSources`, `Geometry`, and
-    `PostProcessedOutputs` dataclasses for each step of the simulation.
-  """
-  step_fn = make_step_fn(torax_config)
-  states_history, post_processed_outputs_history, sim_error = (
-      jit_run_loop.run_loop(
-          step_fn,
-          max_steps=max_steps,
-      )
-  )
-  state_history = output.StateHistory(
-      state_history=states_history,
-      post_processed_outputs_history=post_processed_outputs_history,
-      sim_error=sim_error,
-      torax_config=torax_config,
-  )
-  return (
-      state_history.simulation_output_to_xr(),
-      state_history,
-  )

torax/_src/simulation_app.py

@@ -50,10 +50,12 @@ def run(_):
     logging.set_verbosity(logging.INFO)
     app.run(run)
 """
+
 from collections.abc import Sequence
 import datetime
 import enum
 import os
+
 import shutil
 import sys
 from typing import Callable, Final
@@ -154,6 +156,7 @@ def main(
     log_sim_output: bool = False,
     plot_sim_progress: bool = False,
     log_sim_progress_bar: bool = True,
+    _use_jitted_run_loop: bool = False,
 ) -> str:
   """Runs a simulation obtained via `get_config`.
 
@@ -175,6 +178,8 @@ def main(
     plot_sim_progress: If True, then a plotting spectator will be attached to
       the sim.
     log_sim_progress_bar: If True, then a progress bar will be logged.
+    _use_jitted_run_loop: If True, then a jitted run loop will be used. A
+      temporary private argument used for testing.
 
   Returns:
     The output state file path.
@@ -187,6 +192,7 @@ def main(
       torax_config,
       log_sim_progress,
       progress_bar=log_sim_progress_bar,
+      _use_jitted_run_loop=_use_jitted_run_loop,
   )
 
   # Check if simulation encountered an error

torax/experimental/__init__.py

@@ -21,7 +21,6 @@
 from torax._src.orchestration.initial_state import get_initial_state_and_post_processed_outputs
 from torax._src.orchestration.jit_run_loop import run_loop_jit
 from torax._src.orchestration.run_simulation import make_step_fn
-from torax._src.orchestration.run_simulation import run_simulation_jitted
 from torax._src.orchestration.sim_state import SimState
 from torax._src.orchestration.step_function import SimulationStepFn
 from torax._src.plotting.plotruns_lib import create_plotly_figure
@@ -44,5 +43,4 @@
     'get_initial_state_and_post_processed_outputs',
     'SimState',
     'ExtendedLengyelConfig',
-    'run_simulation_jitted',
 ]

torax/tests/jit_sim_test.py

@@ -275,7 +275,9 @@ def test_run_simulation_jitted(
     torax_config = config_loader.build_torax_config_from_file(config_path)
     reference_file = output.load_state_file(data_path)
 
-    xr_data_tree, _ = run_simulation.run_simulation_jitted(torax_config)
+    xr_data_tree, _ = run_simulation.run_simulation(
+        torax_config, progress_bar=False, _use_jitted_run_loop=True
+    )
 
     # Allow for small numerical differences due to the change in the order of
     # operations in the jitted versus non-jitted case.
@@ -290,8 +292,8 @@ def test_did_not_reach_t_final_error_when_max_steps_too_low(self):
     torax_config = config_loader.build_torax_config_from_file(config_path)
 
     # Use max_steps=1 which is far too few to reach t_final=1.
-    _, state_history = run_simulation.run_simulation_jitted(
-        torax_config, max_steps=1
+    _, state_history = run_simulation.run_simulation(
+        torax_config, max_steps=1, progress_bar=False, _use_jitted_run_loop=True
     )
 
     self.assertEqual(