Add simulation time t to RuntimeParams.

This is a prerequisite for pedestal transition ramp scaling, which needs access
to the simulation time in calc_coeffs for ADAPTIVE_SOURCE transition ramps.

PiperOrigin-RevId: 890681046

Author: jcitrin

torax/_src/config/build_runtime_params.py

@@ -28,6 +28,7 @@
 import equinox as eqx
 import jax
 from jax import numpy as jnp
+from torax._src import jax_utils
 from torax._src import state as state_lib
 from torax._src.config import numerics as numerics_lib
 from torax._src.config import runtime_params as runtime_params_lib
@@ -119,6 +120,7 @@ def __call__(
   ) -> runtime_params_lib.RuntimeParams:
     """Returns a runtime_params.RuntimeParams to use during time t."""
     return runtime_params_lib.RuntimeParams(
+        t=jnp.asarray(t, dtype=jax_utils.get_dtype()),
         transport=self.transport_model.build_runtime_params(t),
         solver=self.solver.build_runtime_params,
         sources={
@@ -352,8 +354,8 @@ def get_consistent_runtime_params_and_geometry(
       ``density_fraction`` boundary condition mode to derive ``n_e_right_bc``
       from the evolved electron density. If ``None`` (e.g. at initialization),
       the prescribed ``n_e`` from the runtime params is used.
-    is_initialization: Indicates if the simulation is currently initializing.
-      If this is `False`, then `core_profiles` cannot be `None`.
+    is_initialization: Indicates if the simulation is currently initializing. If
+      this is `False`, then `core_profiles` cannot be `None`.
 
   Returns:
     A tuple of (RuntimeParams, Geometry).

torax/_src/config/runtime_params.py

@@ -40,6 +40,7 @@
 import dataclasses
 
 import jax
+from torax._src import array_typing
 from torax._src.config import numerics
 from torax._src.core_profiles import profile_conditions
 from torax._src.core_profiles.plasma_composition import plasma_composition
@@ -75,6 +76,7 @@ class RuntimeParams:
   be changed without recompilation.
   """
 
+  t: array_typing.FloatScalar
   edge: edge_runtime_params.RuntimeParams | None
   mhd: mhd_runtime_params.RuntimeParams
   neoclassical: neoclassical_params.RuntimeParams