Sim test and documentation for ADAPTIVE_SOURCE pedestal mode with formation model.

PiperOrigin-RevId: 890699552

Author: jcitrin

docs/configuration.rst

@@ -764,6 +764,67 @@ top. These models will only be used if the ``set_pedestal`` flag is set to True.
   mock up a pedestal, this feature can also be used for L-mode modeling with a
   desired internal boundary condition below :math:`\hat{\rho}=1`.
 
+``mode`` (str [default = 'ADAPTIVE_SOURCE'])
+  Defines how the pedestal is generated. Options:
+
+  * ``'ADAPTIVE_SOURCE'``: Sets the pedestal by adding a source/sink term at the
+    pedestal top, forcing the pedestal top values to be as prescribed by the
+    pedestal model. This is the default mode whenever ``set_pedestal`` is True.
+  * ``'ADAPTIVE_TRANSPORT'``: Sets the pedestal by modifying the transport
+    coefficients in the pedestal region, allowing the pedestal to
+    self-consistently evolve. Transport coefficients are scaled to allow the
+    temperature and density to evolve towards the prescribed pedestal values.
+
+``use_formation_model_with_adaptive_source`` (bool [default = False])
+  Only applicable when ``mode`` is ``'ADAPTIVE_SOURCE'``. When True, enables
+  state-dependent L-H and H-L transitions based on comparison of the power
+  crossing the separatrix (:math:`P_{SOL}`) with the L-H power threshold
+  (:math:`P_{LH}`), as determined by the formation model. Pedestal values are
+  ramped over the ``transition_time_width`` during transitions. When False,
+  ``ADAPTIVE_SOURCE`` mode always applies the prescribed pedestal values
+  (legacy behavior). Raises an error if set to True when ``mode`` is not
+  ``'ADAPTIVE_SOURCE'``.
+
+``transition_time_width`` (**time-varying-scalar** [default = 0.5])
+  Duration of the L-H or H-L transition ramp in seconds. During a transition,
+  pedestal values are linearly interpolated between L-mode baseline values
+  and H-mode target values over this time window. Must be strictly positive.
+  Only used when ``use_formation_model_with_adaptive_source`` is True.
+
+``formation_model`` (dict)
+  Configuration for the pedestal formation model, which determines when L-H
+  and H-L transitions occur. The ``model_name`` key selects the model:
+
+  * ``'martin_scaling'``: Uses the Martin scaling law to determine the L-H
+    power threshold. Additional parameters:
+
+    * ``sharpness`` (float [default = 10.0]): Controls the sharpness of the
+      transition sigmoid function.
+
+  * ``'delabie_scaling'``: Uses the Delabie scaling law for the L-H power
+    threshold. Additional parameters:
+
+    * ``sharpness`` (float [default = 10.0]): Controls the sharpness of the
+      transition sigmoid function.
+
+``saturation_model`` (dict)
+  Configuration for the pedestal saturation model, which determines how the
+  pedestal values scale relative to H-mode target values. The ``model_name``
+  key selects the model:
+
+  * ``'profile_value'``: Uses the current profile value at the pedestal top
+    for saturation. Additional parameters:
+
+    * ``steepness`` (float [default = 100.0]): Controls the steepness of the
+      saturation function.
+    * ``offset`` (float [default = 0.0]): Offset for the saturation function.
+    * ``base_multiplier`` (float [default = 0.0]): Base multiplier for the
+      saturation function.
+
+``pedestal_top_smoothing_width`` (**time-varying-scalar** [default = 0.02])
+  Width of the smoothing region at the pedestal top boundary, in units of
+  :math:`\hat{\rho}`.
+
 The following ``model_name`` options are currently supported:
 
 no_pedestal

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
@@ -72,6 +74,9 @@ def __call__(
       # Checks if reduced calc_coeffs for explicit terms when theta_implicit=1
       # should be called
       explicit_call: bool = False,
+      pedestal_transition_state: (
+          pedestal_transition_state_lib.PedestalTransitionState | None
+      ) = None,
   ) -> block_1d_coeffs.Block1DCoeffs:
     """Returns coefficients given a state x.
 
@@ -84,8 +89,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
@@ -104,6 +109,9 @@ def __call__(
       explicit_call: If True, then if theta_implicit=1, only a reduced
         Block1DCoeffs is calculated since most explicit coefficients will not be
         used.
+      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: The diffusion, convection, etc. coefficients for this state.
@@ -133,6 +141,7 @@ def __call__(
         evolving_names=self.evolving_names,
         use_pereverzev=use_pereverzev,
         explicit_call=explicit_call,
+        pedestal_transition_state=pedestal_transition_state,
     )
 
 
@@ -145,6 +154,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 +182,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 +207,7 @@ def calc_coeffs(
         physics_models=physics_models,
         evolving_names=evolving_names,
         use_pereverzev=use_pereverzev,
+        pedestal_transition_state=pedestal_transition_state,
     )
 
 
@@ -210,6 +226,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."""
 
@@ -268,6 +287,7 @@ def _calc_coeffs_full(
           core_profiles,
           merged_source_profiles,
           use_pereverzev,
+          pedestal_transition_state=pedestal_transition_state,
       )
   )
 
@@ -415,26 +435,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 +608,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),
+  )