Persist rho_norm_ped_top in PedestalTransitionState.

This change makes rho_norm_ped_top part of the PedestalTransitionState, allowing models that dynamically compute the pedestal top location (like EPEDNN) to propagate this value between timesteps. The rho_norm_ped_top from the pedestal model output at t+dt is now stored in the transition state for use in the next timestep's pre_step processing, ensuring accurate L-mode baseline capture.

PiperOrigin-RevId: 890589780

Author: jcitrin

torax/_src/fvm/calc_coeffs.py

@@ -14,6 +14,7 @@
 
 """Calculates Block1DCoeffs for a time step."""
 
+import dataclasses
 import functools
 import jax
 import jax.numpy as jnp
@@ -28,6 +29,7 @@
 from torax._src.fvm import cell_variable
 from torax._src.geometry import geometry
 from torax._src.internal_boundary_conditions import internal_boundary_conditions as internal_boundary_conditions_lib
+from torax._src.pedestal_model import pedestal_transition_state as pedestal_transition_state_lib
 from torax._src.pedestal_model import runtime_params as pedestal_runtime_params_lib
 from torax._src.sources import source_profile_builders
 from torax._src.sources import source_profiles as source_profiles_lib
@@ -43,9 +45,13 @@ def __init__(
       self,
       physics_models: physics_models_lib.PhysicsModels,
       evolving_names: tuple[str, ...],
+      pedestal_transition_state: (
+          pedestal_transition_state_lib.PedestalTransitionState | None
+      ) = None,
   ):
     self.physics_models = physics_models
     self.evolving_names = evolving_names
+    self.pedestal_transition_state = pedestal_transition_state
 
   def __hash__(self) -> int:
     return hash((
@@ -84,8 +90,8 @@ def __call__(
         state x.
       geo: The geometry of the system at this time step.
       core_profiles: The core profiles of the system at this time step.
-      prev_core_profiles: The core profiles of the system at the previous
-        time step.
+      prev_core_profiles: The core profiles of the system at the previous time
+        step.
       dt: The time step size.
       x: The state with cell-grid values of the evolving variables.
       explicit_source_profiles: Precomputed explicit source profiles. These
@@ -133,6 +139,7 @@ def __call__(
         evolving_names=self.evolving_names,
         use_pereverzev=use_pereverzev,
         explicit_call=explicit_call,
+        pedestal_transition_state=self.pedestal_transition_state,
     )
 
 
@@ -145,6 +152,9 @@ def calc_coeffs(
     evolving_names: tuple[str, ...],
     use_pereverzev: bool = False,
     explicit_call: bool = False,
+    pedestal_transition_state: (
+        pedestal_transition_state_lib.PedestalTransitionState | None
+    ) = None,
 ) -> block_1d_coeffs.Block1DCoeffs:
   """Calculates Block1DCoeffs for the time step described by `core_profiles`.
 
@@ -170,6 +180,9 @@ def calc_coeffs(
       explicit component of the PDE. Then calculates a reduced Block1DCoeffs if
       theta_implicit=1. This saves computation for the default fully implicit
       implementation.
+    pedestal_transition_state: State for tracking pedestal L-H and H-L
+      transitions. Only used when the pedestal mode is ADAPTIVE_SOURCE with
+      use_formation_model_with_adaptive_source=True. None otherwise.
 
   Returns:
     coeffs: Block1DCoeffs containing the coefficients at this time step.
@@ -192,6 +205,7 @@ def calc_coeffs(
         physics_models=physics_models,
         evolving_names=evolving_names,
         use_pereverzev=use_pereverzev,
+        pedestal_transition_state=pedestal_transition_state,
     )
 
 
@@ -210,6 +224,9 @@ def _calc_coeffs_full(
     physics_models: physics_models_lib.PhysicsModels,
     evolving_names: tuple[str, ...],
     use_pereverzev: bool = False,
+    pedestal_transition_state: (
+        pedestal_transition_state_lib.PedestalTransitionState | None
+    ) = None,
 ) -> block_1d_coeffs.Block1DCoeffs:
   """See `calc_coeffs` for details."""
 
@@ -415,26 +432,75 @@ def _calc_coeffs_full(
       runtime_params.pedestal.mode
       == pedestal_runtime_params_lib.Mode.ADAPTIVE_SOURCE
   ):
+    # Get the pedestal-top target values from the pedestal model.
+    pedestal_top_values = (
+        pedestal_model_output.to_internal_boundary_conditions(geo)
+    )
+
+    # Apply ramp scaling if use_formation_model_with_adaptive_source is
+    # enabled.
+    if runtime_params.pedestal.use_formation_model_with_adaptive_source:
+      assert pedestal_transition_state is not None, (
+          'pedestal_transition_state must not be None when'
+          ' use_formation_model_with_adaptive_source is True.'
+      )
+      # Scale the pedestal-top values from the pedestal model by the ramp
+      # fraction. Will be a no-op in H-mode following the transition_time_width.
+      internal_boundary_conditions = _apply_transition_ramp_scaling(
+          pedestal_top_values=pedestal_top_values,
+          pedestal_transition_state=pedestal_transition_state,
+          runtime_params=runtime_params,
+      )
+      # If in L-mode and the H->L ramp has completed (fraction >= 1.0), skip
+      # the adaptive source entirely to revert to standard L-mode modeling.
+      # ramp_fraction will be 1.0 if simulation initialized in L-mode and has
+      # remained in L-mode, since initial transition_start_time is -inf.
+      ramp_fraction = _compute_ramp_fraction(
+          pedestal_transition_state=pedestal_transition_state,
+          transition_time_width=runtime_params.pedestal.transition_time_width,
+          t=runtime_params.t,
+      )
+      # Skip adaptive source if in L-mode and the H->L ramp has completed.
+      skip_adaptive_source = ~pedestal_transition_state.in_H_mode & (
+          ramp_fraction >= 1.0
+      )
+    else:
+      internal_boundary_conditions = pedestal_top_values
+      skip_adaptive_source = jnp.bool_(False)
+
+    def _apply_source():
+      return internal_boundary_conditions_lib.apply_adaptive_source(
+          source_T_i=source_i,
+          source_T_e=source_e,
+          source_n_e=source_n_e,
+          source_mat_ii=source_mat_ii,
+          source_mat_ee=source_mat_ee,
+          source_mat_nn=source_mat_nn,
+          runtime_params=runtime_params,
+          internal_boundary_conditions=internal_boundary_conditions,
+      )
+
+    def _skip_source():
+      return (
+          source_i,
+          source_e,
+          source_n_e,
+          source_mat_ii,
+          source_mat_ee,
+          source_mat_nn,
+      )
+
     (
         source_i,
         source_e,
         source_n_e,
         source_mat_ii,
         source_mat_ee,
         source_mat_nn,
-    ) = internal_boundary_conditions_lib.apply_adaptive_source(
-        source_T_i=source_i,
-        source_T_e=source_e,
-        source_n_e=source_n_e,
-        source_mat_ii=source_mat_ii,
-        source_mat_ee=source_mat_ee,
-        source_mat_nn=source_mat_nn,
-        runtime_params=runtime_params,
-        # Pedestal contributes an internal boundary condition to the source
-        # terms at the pedestal top.
-        internal_boundary_conditions=pedestal_model_output.to_internal_boundary_conditions(
-            geo
-        ),
+    ) = jax.lax.cond(
+        skip_adaptive_source,
+        _skip_source,
+        _apply_source,
     )
 
   # --- Build arguments to solver  --- #
@@ -539,3 +605,93 @@ def _calc_coeffs_reduced(
       transient_in_cell=transient_in_cell,
   )
   return coeffs
+
+
+def _compute_ramp_fraction(
+    pedestal_transition_state: pedestal_transition_state_lib.PedestalTransitionState,
+    transition_time_width: array_typing.FloatScalar,
+    t: array_typing.FloatScalar,
+) -> array_typing.FloatScalar:
+  """Computes the ramp fraction for a pedestal transition.
+
+  Returns a value in [0, 1] representing the progress of the current
+  transition. 0 means the transition just started, 1 means it is complete.
+
+  Args:
+    pedestal_transition_state: Current transition state.
+    transition_time_width: Duration of the transition ramp.
+    t: Current simulation time (i.e. t + dt when called from the solver).
+
+  Returns:
+    Ramp fraction clipped to [0, 1].
+  """
+  elapsed = t - pedestal_transition_state.transition_start_time
+  fraction = elapsed / transition_time_width
+  return jnp.clip(fraction, 0.0, 1.0)
+
+
+def _apply_transition_ramp_scaling(
+    pedestal_top_values: internal_boundary_conditions_lib.InternalBoundaryConditions,
+    pedestal_transition_state: pedestal_transition_state_lib.PedestalTransitionState,
+    runtime_params: runtime_params_lib.RuntimeParams,
+) -> internal_boundary_conditions_lib.InternalBoundaryConditions:
+  """Applies ramp scaling to internal boundary conditions during transitions.
+
+  During an L-H transition, linearly ramps from L-mode values to the H-mode
+  targets. During an H-L transition, ramps from the H-mode targets back to
+  the L-mode values.
+
+  The L-mode values are stored in the pedestal_transition_state (captured
+  at the start of an L->H transition). The H-mode targets are the full
+  pedestal model output.
+
+  Args:
+    pedestal_top_values: Pedestal-top target internal boundary conditions from
+      the pedestal model.
+    pedestal_transition_state: Current transition state containing L-mode
+      baseline values.
+    runtime_params: Runtime parameters (provides time t and pedestal config).
+
+  Returns:
+    Scaled internal boundary conditions.
+  """
+  ramp_fraction = _compute_ramp_fraction(
+      pedestal_transition_state=pedestal_transition_state,
+      transition_time_width=runtime_params.pedestal.transition_time_width,
+      t=runtime_params.t,
+  )
+
+  # Extract the nonzero pedestal-top values from the IBC. The IBC arrays are
+  # cell-grid sized with a single nonzero element at the pedestal top. We use
+  # jnp.max to extract the nonzero value.
+  h_mode_T_i_ped = jnp.max(pedestal_top_values.T_i)
+  h_mode_T_e_ped = jnp.max(pedestal_top_values.T_e)
+  h_mode_n_e_ped = jnp.max(pedestal_top_values.n_e)
+
+  l_mode_T_i_ped = pedestal_transition_state.T_i_ped_L_mode
+  l_mode_T_e_ped = pedestal_transition_state.T_e_ped_L_mode
+  l_mode_n_e_ped = pedestal_transition_state.n_e_ped_L_mode
+
+  # In H-mode: ramp from L-mode to H-mode (L + fraction * (H - L))
+  # In L-mode (H->L ramp): ramp from H-mode to L-mode (H + fraction * (L - H))
+  def _lerp(l_val, h_val, frac, in_h_mode):
+    return jnp.where(
+        in_h_mode,
+        l_val + frac * (h_val - l_val),  # L->H ramp
+        h_val + frac * (l_val - h_val),  # H->L ramp
+    )
+
+  in_h_mode = pedestal_transition_state.in_H_mode
+  scaled_T_i = _lerp(l_mode_T_i_ped, h_mode_T_i_ped, ramp_fraction, in_h_mode)
+  scaled_T_e = _lerp(l_mode_T_e_ped, h_mode_T_e_ped, ramp_fraction, in_h_mode)
+  scaled_n_e = _lerp(l_mode_n_e_ped, h_mode_n_e_ped, ramp_fraction, in_h_mode)
+
+  # Reconstruct IBC with scaled values at the same pedestal-top location.
+  # The nonzero mask from the original pedestal_top_values gives us the
+  # location.
+  return dataclasses.replace(
+      pedestal_top_values,
+      T_i=jnp.where(pedestal_top_values.T_i != 0.0, scaled_T_i, 0.0),
+      T_e=jnp.where(pedestal_top_values.T_e != 0.0, scaled_T_e, 0.0),
+      n_e=jnp.where(pedestal_top_values.n_e != 0.0, scaled_n_e, 0.0),
+  )

torax/_src/fvm/tests/calc_coeffs_test.py

@@ -16,9 +16,12 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
+import jax.numpy as jnp
 from torax._src.config import build_runtime_params
 from torax._src.core_profiles import initialization
 from torax._src.fvm import calc_coeffs
+from torax._src.internal_boundary_conditions import internal_boundary_conditions
+from torax._src.pedestal_model import pedestal_transition_state
 from torax._src.sources import source_profile_builders
 from torax._src.test_utils import default_sources
 from torax._src.torax_pydantic import model_config
@@ -159,5 +162,92 @@ def create_coeffs_callback(
       )
 
 
+class TransitionCalculationsTest(parameterized.TestCase):
+
+  def test_pedestal_transition_state_initial_state(self):
+    state = pedestal_transition_state.PedestalTransitionState.initial_state()
+    self.assertTrue(jnp.isneginf(state.transition_start_time))
+    self.assertEqual(state.T_i_ped_L_mode, 0.0)
+    self.assertFalse(state.in_H_mode)
+
+  def test_compute_ramp_fraction_very_small_width(self):
+    state = pedestal_transition_state.PedestalTransitionState(
+        transition_start_time=jnp.array(1.0),
+        T_i_ped_L_mode=jnp.array(0.0),
+        T_e_ped_L_mode=jnp.array(0.0),
+        n_e_ped_L_mode=jnp.array(0.0),
+        in_H_mode=jnp.array(True),
+        rho_norm_ped_top=jnp.array(0.9),
+    )
+    # Very small transition_time_width clips to 1.0 when elapsed > 0.
+    self.assertEqual(
+        calc_coeffs._compute_ramp_fraction(state, 1e-10, 1.5), 1.0
+    )
+
+  def test_compute_ramp_fraction_ramp(self):
+    state = pedestal_transition_state.PedestalTransitionState(
+        transition_start_time=jnp.array(1.0),
+        T_i_ped_L_mode=jnp.array(0.0),
+        T_e_ped_L_mode=jnp.array(0.0),
+        n_e_ped_L_mode=jnp.array(0.0),
+        in_H_mode=jnp.array(True),
+        rho_norm_ped_top=jnp.array(0.9),
+    )
+    # transition_time_width = 1.0. Start at 1.0.
+    # Clip at both ends
+    self.assertEqual(
+        calc_coeffs._compute_ramp_fraction(state, 1.0, 0.5), 0.0
+    )  # t < start
+    self.assertEqual(
+        calc_coeffs._compute_ramp_fraction(state, 1.0, 1.0), 0.0
+    )  # t = start
+    self.assertEqual(
+        calc_coeffs._compute_ramp_fraction(state, 1.0, 1.5), 0.5
+    )  # t = start + 0.5
+    self.assertEqual(
+        calc_coeffs._compute_ramp_fraction(state, 1.0, 2.0), 1.0
+    )  # t = start + 1.0
+    self.assertEqual(
+        calc_coeffs._compute_ramp_fraction(state, 1.0, 2.5), 1.0
+    )  # t = start + 1.5
+
+  def test_apply_transition_ramp_scaling_l_to_h(self):
+    l_mode_baseline = 1.0
+    h_mode_target = 3.0
+
+    state = pedestal_transition_state.PedestalTransitionState(
+        transition_start_time=jnp.array(1.0),
+        T_i_ped_L_mode=jnp.array(l_mode_baseline),
+        T_e_ped_L_mode=jnp.array(l_mode_baseline),
+        n_e_ped_L_mode=jnp.array(l_mode_baseline),
+        in_H_mode=jnp.array(True),  # L -> H
+        rho_norm_ped_top=jnp.array(0.9),
+    )
+
+    pedestal_top_values = (
+        internal_boundary_conditions.InternalBoundaryConditions(
+            T_i=jnp.array([0.0, h_mode_target, 0.0]),
+            T_e=jnp.array([0.0, h_mode_target, 0.0]),
+            n_e=jnp.array([0.0, h_mode_target, 0.0]),
+        )
+    )
+
+    class MockPedestalRuntimeParams:
+      transition_time_width = 1.0
+
+    class MockRuntimeParams:
+      pedestal = MockPedestalRuntimeParams()
+      t = 1.5  # halfway
+
+    scaled_ibc = calc_coeffs._apply_transition_ramp_scaling(  # pytype: disable=wrong-arg-types
+        pedestal_top_values=pedestal_top_values,
+        pedestal_transition_state=state,
+        runtime_params=MockRuntimeParams(),
+    )
+
+    # Expected: 1.0 + 0.5 * (3.0 - 1.0) = 2.0
+    self.assertTrue(jnp.allclose(scaled_ibc.T_i, jnp.array([0.0, 2.0, 0.0])))
+
+
 if __name__ == '__main__':
   absltest.main()