Improve numerics of LH transition test scenario.

Main changes:
- Switch from chi to fixed timestep calculator
- Add smoothing to chi profiles
- Remove resistivity multiplier, which was causing numerical instability
- Switch to COMBINED transport model for clarity

PiperOrigin-RevId: 878555971

Author: theo-brown

torax/_src/pedestal_model/formation/martin_formation_model.py

@@ -16,7 +16,6 @@
 
 import dataclasses
 import jax
-import jax.numpy as jnp
 from torax._src import array_typing
 from torax._src import math_utils
 from torax._src import state
@@ -86,19 +85,13 @@ def __call__(
     # If P_SOL > P_LH, multiplier tends to 0.0
     # If P_SOL < P_LH, multiplier tends to 1.0
     # TODO(b/488393318): Add hysteresis to the LH-HL transition.
-    width = runtime_params.pedestal.formation.sigmoid_width
-    exponent = runtime_params.pedestal.formation.sigmoid_exponent
-    offset = runtime_params.pedestal.formation.sigmoid_offset
-    normalized_deviation = (
-        P_SOL_total - rescaled_P_LH
-    ) / rescaled_P_LH - offset
-    transport_multiplier = 1 - jax.nn.sigmoid(normalized_deviation / width)
-    transport_multiplier = transport_multiplier**exponent
-    transport_multiplier = jnp.clip(
-        transport_multiplier,
-        min=runtime_params.pedestal.min_transport_multiplier,
-        max=runtime_params.pedestal.max_transport_multiplier,
-    )
+    sharpness = runtime_params.pedestal.formation.sharpness
+    offset = runtime_params.pedestal.formation.offset
+    base_multiplier = runtime_params.pedestal.formation.base_multiplier
+    normalized_deviation = (P_SOL_total - rescaled_P_LH) / rescaled_P_LH
+    shifted_deviation = normalized_deviation - offset
+    alpha = jax.nn.sigmoid(shifted_deviation * sharpness)
+    transport_multiplier = (1.0 - alpha) * 1.0 + alpha * base_multiplier
 
     return pedestal_model_output.TransportMultipliers(
         chi_e_multiplier=transport_multiplier,

torax/_src/pedestal_model/pydantic_model.py

@@ -38,36 +38,35 @@ class MartinFormation(torax_pydantic.BaseModelFrozen):
   """Configuration for Martin formation model.
 
   This formation model triggers a reduction in pedestal transport when P_SOL >
-  P_LH, where P_LH is from the Martin scaling. The reduction is a smooth sigmoid
-  function of the ratio P_SOL / (P_LH * P_LH_prefactor).
+  P_LH, where P_LH is from the Martin scaling. The reduction is a multiplicative
+  factors between 1.0 and base_multiplier.
 
   The formula is
-    rescaled_P_LH = P_LH * P_LH_prefactor
-      normalized_deviation = (P_SOL - rescaled_P_LH) / rescaled_P_LH - offset
-      transport_multiplier = 1 - sigmoid(normalized_deviation / width)
-      transport_multiplier = transport_multiplier**exponent
-
-  The transport multiplier is later clipped to the range
-  [min_transport_multiplier, max_transport_multiplier].
+    transport_multiplier = (1.0 - alpha) * 1.0 + alpha * base_multiplier
+  where alpha is a smooth sigmoid function of
+    (P_SOL - P_LH * P_LH_prefactor) / (P_LH * P_LH_prefactor)
+  with given sharpness and offset.
 
   Attributes:
-    sigmoid_width: Dimensionless width of the sigmoid function for smoothing the
-      formation window. Increase for a smoother L-H transition, but doing so may
-      lead to starting the L-H transition at a power below P_LH.
-    sigmoid_offset: Dimensionless offset of sigmoid function from P_LH / P_SOL =
-      1. Increase to start the L-H transition at a higher P_SOL / P_LH ratio.
-    sigmoid_exponent: The exponent of the transport multiplier. Increase for a
-      faster reduction in transport once the L-H transition starts.
+    sharpness: The steepness of the smooth formation window. Decrease for a
+      smoother formation, which may be more numerically stable but may lead to
+      starting formation at a temperature or density below the target  values.
+    offset: Dimensionless offset of the formation window. Increase to start
+      formation at a higher P_SOL.
+    base_multiplier: The base value of the transport multiplier. Increase for
+      stronger decreases in transport once formation starts.
     P_LH_prefactor: Dimensionless multiplier for P_LH. Increase to scale up
       P_LH, and therefore start the L-H transition at a higher P_SOL.
   """
 
   model_name: Annotated[Literal["martin"], torax_pydantic.JAX_STATIC] = "martin"
-  sigmoid_width: pydantic.PositiveFloat = 1e-3
-  sigmoid_offset: Annotated[
+  sharpness: pydantic.PositiveFloat = 1e3
+  offset: Annotated[
       array_typing.FloatScalar, pydantic.Field(ge=-10.0, le=10.0)
   ] = 0.0
-  sigmoid_exponent: pydantic.PositiveFloat = 3.0
+  base_multiplier: Annotated[
+      array_typing.FloatScalar, pydantic.Field(gt=0.0, le=1.0)
+  ] = 0.1
   P_LH_prefactor: pydantic.PositiveFloat = 1.0
 
   def build_formation_model(
@@ -80,9 +79,9 @@ def build_runtime_params(
   ) -> martin_formation_model.MartinFormationRuntimeParams:
     del t
     return martin_formation_model.MartinFormationRuntimeParams(
-        sigmoid_width=self.sigmoid_width,
-        sigmoid_offset=self.sigmoid_offset,
-        sigmoid_exponent=self.sigmoid_exponent,
+        sharpness=self.sharpness,
+        offset=self.offset,
+        base_multiplier=self.base_multiplier,
         P_LH_prefactor=self.P_LH_prefactor,
     )
 
@@ -92,36 +91,34 @@ class ProfileValueSaturation(torax_pydantic.BaseModelFrozen):
 
   This saturation model triggers an increase in pedestal transport when the
   pedestal temperature and density are above the values requested by the
-  pedestal model. The increase is a smooth sigmoid function of the ratio of the
+  pedestal model. The increase is a smooth linear function of the ratio of the
   current value to the value requested by the pedestal model.
 
   The formula is
-    normalized_deviation = (current - target) / target - offset
-    transport_multiplier = 1 / (1 - sigmoid(normalized_deviation / width))
-    transport_multiplier = transport_multiplier**exponent
-
-  The transport multiplier is then clipped to the range
-  [min_transport_multiplier, max_transport_multiplier].
+    transport_multiplier = 1 + base_multiplier * softplus(actual, target)
+  where softplus(actual, target) is a smooth ReLU function of the deviation
+  from the target value, with given sharpness and offset.
 
   Attributes:
-    sigmoid_width: Dimensionless width of the sigmoid function for smoothing the
-      saturation window. Increase for a smoother saturation, but doing so may
-      lead to starting saturation at a temperature or density below the target
-      values.
-    sigmoid_offset: Dimensionless offset of the saturation window. Increase to
-      start saturation at a higher temperature or density.
-    sigmoid_exponent: The exponent of the transport multiplier. Increase for a
-      faster increase in transport once saturation starts.
+    sharpness: The steepness of the smooth saturation function. Decrease for a a
+      smoother saturation, which may be more numerically stable but may lead to
+      starting saturation at a temperature or density below the target values.
+    offset: Dimensionless offset of the saturation window. Increase to start
+      saturation at a higher temperature or density.
+    base_multiplier: The base value of the transport multiplier. Increase for
+      stronger increases in transport once saturation starts.
   """
 
   model_name: Annotated[Literal["profile_value"], torax_pydantic.JAX_STATIC] = (
       "profile_value"
   )
-  sigmoid_width: pydantic.PositiveFloat = 0.1
-  sigmoid_offset: Annotated[
+  sharpness: pydantic.PositiveFloat = 1e3
+  offset: Annotated[
       array_typing.FloatScalar, pydantic.Field(ge=-10.0, le=10.0)
   ] = 0.0
-  sigmoid_exponent: pydantic.PositiveFloat = 1.0
+  base_multiplier: Annotated[
+      array_typing.FloatScalar, pydantic.Field(gt=1.0)
+  ] = 10.0
 
   def build_saturation_model(
       self,
@@ -133,9 +130,9 @@ def build_runtime_params(
   ) -> runtime_params.SaturationRuntimeParams:
     del t
     return runtime_params.SaturationRuntimeParams(
-        sigmoid_width=self.sigmoid_width,
-        sigmoid_offset=self.sigmoid_offset,
-        sigmoid_exponent=self.sigmoid_exponent,
+        sharpness=self.sharpness,
+        offset=self.offset,
+        base_multiplier=self.base_multiplier,
     )
 
 
@@ -170,10 +167,10 @@ class BasePedestal(torax_pydantic.BaseModelFrozen, abc.ABC):
       ProfileValueSaturation()
   )
   max_transport_multiplier: torax_pydantic.TimeVaryingScalar = (
-      torax_pydantic.ValidatedDefault(10.0)
+      torax_pydantic.ValidatedDefault(1e3)
   )
   min_transport_multiplier: torax_pydantic.TimeVaryingScalar = (
-      torax_pydantic.ValidatedDefault(0.1)
+      torax_pydantic.ValidatedDefault(1e-3)
   )
 
   @pydantic.model_validator(mode="before")

torax/_src/pedestal_model/runtime_params.py

@@ -37,19 +37,19 @@ class Mode(enum.Enum):
 class FormationRuntimeParams:
   """Runtime params for pedestal formation models."""
 
-  sigmoid_width: array_typing.FloatScalar = 0.1
-  sigmoid_offset: array_typing.FloatScalar = 0.0
-  sigmoid_exponent: array_typing.FloatScalar = 1.0
+  sharpness: array_typing.FloatScalar
+  offset: array_typing.FloatScalar
+  base_multiplier: array_typing.FloatScalar
 
 
 @jax.tree_util.register_dataclass
 @dataclasses.dataclass(frozen=True)
 class SaturationRuntimeParams:
   """Runtime params for pedestal saturation models."""
 
-  sigmoid_width: array_typing.FloatScalar = 0.1
-  sigmoid_offset: array_typing.FloatScalar = 0.0
-  sigmoid_exponent: array_typing.FloatScalar = 1.0
+  sharpness: array_typing.FloatScalar
+  offset: array_typing.FloatScalar
+  base_multiplier: array_typing.FloatScalar
 
 
 @jax.tree_util.register_dataclass

torax/_src/pedestal_model/saturation/profile_value_saturation_model.py

@@ -91,18 +91,11 @@ def _calculate_multiplier(
     Returns:
       The transport increase multiplier.
     """
-    width = pedestal_runtime_params.saturation.sigmoid_width
-    exponent = pedestal_runtime_params.saturation.sigmoid_exponent
-    offset = pedestal_runtime_params.saturation.sigmoid_offset
+    sharpness = pedestal_runtime_params.saturation.sharpness
+    offset = pedestal_runtime_params.saturation.offset
+    base_multiplier = pedestal_runtime_params.saturation.base_multiplier
     normalized_deviation = (current - target) / target - offset
-    transport_multiplier = 1 / (
-        1 - jax.nn.sigmoid(normalized_deviation / width)
+    transport_multiplier = 1 + base_multiplier * jax.nn.softplus(
+        normalized_deviation * sharpness
     )
-    transport_multiplier = transport_multiplier**exponent
-    transport_multiplier = jnp.clip(
-        transport_multiplier,
-        min=pedestal_runtime_params.min_transport_multiplier,
-        max=pedestal_runtime_params.max_transport_multiplier,
-    )
-
     return transport_multiplier

torax/_src/pedestal_model/saturation/tests/profile_value_saturation_model_test.py

@@ -11,11 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import dataclasses
 from absl.testing import absltest
 from absl.testing import parameterized
 import numpy as np
-from torax._src import constants
 from torax._src.config import build_runtime_params
 from torax._src.core_profiles import initialization
 from torax._src.pedestal_model import pedestal_model_output
@@ -49,43 +47,35 @@ def setUp(self):
   @parameterized.named_parameters(
       dict(
           testcase_name='active',
-          # P_current >> P_target -> saturation is active.
-          P_current_over_P_target=1e3,
+          # T_current >> T_target -> saturation is active.
+          T_target_over_T_current=1e-3,
       ),
       dict(
           testcase_name='inactive',
-          # P_current << P_target -> no saturation.
-          P_current_over_P_target=1e-3,
+          # T_current << T_target -> no saturation.
+          T_target_over_T_current=1e3,
       ),
   )
   def test_saturation_multiplier(
       self,
-      P_current_over_P_target,
+      T_target_over_T_current,
   ):
     saturation_model = (
         profile_value_saturation_model.ProfileValueSaturationModel()
     )
 
     # For this test, we put the pedestal top at the last grid point.
     ped_top_idx = -1
-
-    # Get the current pressure at the pedestal top.
-    current_pressure = self.core_profiles.pressure_thermal_total.value[
-        ped_top_idx
-    ]
+    current_T_e_ped = self.core_profiles.T_e.face_value()[ped_top_idx]
 
     # Construct a pedestal output that is asking for a pedestal with
-    # target_pressure.
-    target_pressure = current_pressure / P_current_over_P_target
+    # target temperature.
     pedestal_output = pedestal_model_output.PedestalModelOutput(
         rho_norm_ped_top=self.geo.rho_face[ped_top_idx],
         rho_norm_ped_top_idx=ped_top_idx,
-        # Set T_i_ped, T_e_ped, and n_e_ped such that target_pressure is
-        # achieved. This case has n_impurity = 0 and Z_i = 1, so
-        # P = (T_i + T_e)*keV_to_J*n_e.
-        T_i_ped=(1.0 / constants.CONSTANTS.keV_to_J) / 2,
-        T_e_ped=(1.0 / constants.CONSTANTS.keV_to_J) / 2,
-        n_e_ped=target_pressure,
+        T_i_ped=1.0,
+        T_e_ped=current_T_e_ped * T_target_over_T_current,
+        n_e_ped=1.0,
     )
 
     transport_multipliers = saturation_model(
@@ -95,16 +85,14 @@ def test_saturation_multiplier(
         pedestal_output,
     )
 
-    if P_current_over_P_target > 1.0:
-      # If the current pressure is above the target pressure, we expect the
-      # multiplier to be greater than 1.0.
-      for multiplier in dataclasses.asdict(transport_multipliers).values():
-        self.assertGreater(multiplier, 1.0)
+    if T_target_over_T_current > 1.0:
+      # If the target temperature is above the current temperature, we expect
+      # the multiplier to be equal to 1.0 - the pedestal is not saturated.
+      np.testing.assert_allclose(transport_multipliers.chi_e_multiplier, 1.0)
     else:
-      # If the current pressure is below the target pressure, we expect the
-      # multiplier to be 1.0.
-      for multiplier in dataclasses.asdict(transport_multipliers).values():
-        np.testing.assert_allclose(multiplier, 1.0, rtol=1e-3)
+      # If the target temperature is below the current temperature, we expect
+      # the multiplier to be greater than 1.0 - the pedestal is saturated.
+      self.assertGreater(transport_multipliers.chi_e_multiplier, 1.0)
 
 
 if __name__ == '__main__':

torax/_src/transport_model/combined.py

@@ -84,15 +84,8 @@ def __call__(
     # should be handled by instantiating constant component models instead.
     # However, the rho_inner and rho_outer arguments are currently required
     # in the case where the inner/outer region are to be excluded from
-    # smoothing. Smoothing is applied to
-    # rho_inner < rho_norm < min(rho_ped_top, rho_outer) unless
-    # smooth_everywhere is True.
-    return self._smooth_coeffs(
-        runtime_params,
-        geo,
-        transport_coeffs,
-        pedestal_model_output,
-    )
+    # smoothing.
+    return transport_coeffs
 
   def call_implementation(
       self,

torax/_src/transport_model/tests/transport_model_test.py

@@ -309,6 +309,13 @@ def test_smoothing_everywhere(self):
         core_profiles,
         pedestal_model_outputs,
     )
+
+    # Apply the smoothing
+    transport_coeffs = transport_model.smooth_coeffs(
+        runtime_params, geo, transport_coeffs, pedestal_model_outputs
+    )
+
+    # Set up original transport coefficients for comparison
     inner_patch_idx = np.searchsorted(
         geo.rho_face_norm, runtime_params.transport.rho_inner
     )

torax/_src/transport_model/transport_coefficients_builder.py

@@ -85,7 +85,6 @@ def calculate_all_transport_coeffs(
 
   # Modify the turbulent + Pereverzev transport coefficients if the pedestal
   # model is in ADAPTIVE_TRANSPORT mode.
-  # TODO(b/488980968): Identify speed issue with ADAPTIVE_TRANSPORT mode.
   if (
       runtime_params.pedestal.mode
       == pedestal_runtime_params_lib.Mode.ADAPTIVE_TRANSPORT
@@ -111,4 +110,35 @@ def calculate_all_transport_coeffs(
         **dataclasses.asdict(pereverzev_transport_coeffs),
     )
 
+  # TODO(b/485147781) Clean up this post-processing.
+  # Apply smoothing to the final turbulent transport coefficients
+  turbulent_transport_coeffs = transport_model_lib.TurbulentTransport(
+      chi_face_ion=core_transport.chi_face_ion,
+      chi_face_el=core_transport.chi_face_el,
+      d_face_el=core_transport.d_face_el,
+      v_face_el=core_transport.v_face_el,
+      chi_face_ion_bohm=core_transport.chi_face_ion_bohm,
+      chi_face_ion_gyrobohm=core_transport.chi_face_ion_gyrobohm,
+      chi_face_ion_itg=core_transport.chi_face_ion_itg,
+      chi_face_ion_tem=core_transport.chi_face_ion_tem,
+      chi_face_el_bohm=core_transport.chi_face_el_bohm,
+      chi_face_el_gyrobohm=core_transport.chi_face_el_gyrobohm,
+      chi_face_el_itg=core_transport.chi_face_el_itg,
+      chi_face_el_tem=core_transport.chi_face_el_tem,
+      chi_face_el_etg=core_transport.chi_face_el_etg,
+      d_face_el_itg=core_transport.d_face_el_itg,
+      d_face_el_tem=core_transport.d_face_el_tem,
+      v_face_el_itg=core_transport.v_face_el_itg,
+      v_face_el_tem=core_transport.v_face_el_tem,
+  )
+  turbulent_transport_coeffs = transport_model.smooth_coeffs(
+      runtime_params=runtime_params,
+      geo=geo,
+      pedestal_model_output=pedestal_model_output,
+      transport_coeffs=turbulent_transport_coeffs,
+  )
+  core_transport = dataclasses.replace(
+      core_transport,
+      **dataclasses.asdict(turbulent_transport_coeffs),
+  )
   return core_transport