Add pedestal_smoothing_width to transport config and smooth pedestal transport.

docs/configuration.rst

@@ -1507,6 +1507,10 @@ Top-level Transport Parameters
   is used and ``smoothing_width`` is set to ``0.0``, a warning is logged
   recommending non-zero smoothing to avoid sharp numerical artifacts.
 
+``pedestal_smoothing_width`` (float [default = 0.0])
+  Width of HWHM Gaussian smoothing kernel operating on pedestal transport model
+  outputs. If set to ``0.0``, no smoothing is applied to pedestal models.
+
 ``smoothing_zones`` (list[dict] [default = []])
   Optional list of configuration dictionaries specifying radial zones with
   distinct Gaussian smoothing widths. Each zone dictionary specifies

docs/physics_models.rst

@@ -303,7 +303,7 @@ TORAX combines one or more core and pedestal component models within a top-level
 Numerical conditioning and post-processing of transport coefficients—such as
 minimum/maximum clipping (``chi_min``, ``chi_max``, ``D_e_min/max``,
 ``V_e_min/max``) and spatial Gaussian smoothing (``smoothing_width``,
-``smoothing_zones``)—are configured on the top-level transport model rather than
+``pedestal_smoothing_width``, ``smoothing_zones``)—are configured on the top-level transport model rather than
 on individual component models. Furthermore, to combine transport models or
 prescribe transport coefficients within specific radial zones (e.g. inner core
 or outer edge regions), users can configure dictionaries of named component

torax/_src/transport_model/pydantic_model.py

@@ -470,6 +470,9 @@ class TransportModel(torax_pydantic.BaseModelFrozen):
     V_e_max: maximum electron density convection.
     smoothing_width: Width of HWHM Gaussian smoothing kernel operating on
       transport model outputs.
+    pedestal_smoothing_width: Width of HWHM Gaussian smoothing kernel operating
+      on pedestal transport model outputs. If 0.0, no smoothing is applied to
+      pedestal models.
     core_transport_models: A dict mapping user-given names to transport models,
       whose outputs will be summed to give the combined core transport
       coefficients.
@@ -488,6 +491,7 @@ class TransportModel(torax_pydantic.BaseModelFrozen):
   V_e_min: torax_pydantic.MeterPerSecond = -50.0
   V_e_max: torax_pydantic.MeterPerSecond = 50.0
   smoothing_width: pydantic.NonNegativeFloat = 0.0
+  pedestal_smoothing_width: pydantic.NonNegativeFloat = 0.0
   core_transport_models: dict[
       str, ComponentTransportModelConfig
   ] = pydantic.Field(
@@ -546,6 +550,7 @@ def build_runtime_params(
         V_e_min=self.V_e_min,
         V_e_max=self.V_e_max,
         smoothing_width=self.smoothing_width,
+        pedestal_smoothing_width=self.pedestal_smoothing_width,
         core_transport_model_params=core_transport_model_params,
         pedestal_transport_model_params=pedestal_transport_model_params,
         smoothing_zones=tuple(smoothing_zones),
@@ -554,7 +559,8 @@ def build_runtime_params(
   @pydantic.model_validator(mode='after')
   def _check_smoothing_width_minimum(self) -> Self:
     smoothing_widths = [z.smoothing_width for z in self.smoothing_zones] + [
-        self.smoothing_width
+        self.smoothing_width,
+        self.pedestal_smoothing_width,
     ]
     if any(w < 0.0 for w in smoothing_widths):
       raise ValueError(

torax/_src/transport_model/runtime_params.py

@@ -68,6 +68,7 @@ class RuntimeParams:
   V_e_min: float
   V_e_max: float
   smoothing_width: float
+  pedestal_smoothing_width: float
   core_transport_model_params: Mapping[str, ComponentRuntimeParams]
   pedestal_transport_model_params: Mapping[str, ComponentRuntimeParams]
   smoothing_zones: tuple[SmoothingZoneParams, ...]

torax/_src/transport_model/tests/transport_model_test.py

@@ -615,6 +615,54 @@ def test_build_smoothing_matrix_pedestal_boundary_isolation(self):
       expected_row[idx] = 1.0
       np.testing.assert_allclose(matrix[idx], expected_row, atol=1e-7)
 
+  def test_build_pedestal_smoothing_matrix_zero_width_is_identity(self):
+    config = default_configs.get_default_config_dict()
+    config['transport'] = {'pedestal_smoothing_width': 0.0}
+    torax_config = model_config.ToraxConfig.from_dict(config)
+    runtime_params = build_runtime_params.RuntimeParamsProvider.from_config(
+        torax_config
+    )(t=torax_config.numerics.t_initial)
+    geo = torax_config.geometry.build_provider(
+        t=torax_config.numerics.t_initial
+    )
+    transition_state = _make_transition_state(rho_norm_ped_top=0.8)
+    matrix = transport_model._build_pedestal_smoothing_matrix(
+        runtime_params.transport,
+        geo,
+        transition_state,
+    )
+    np.testing.assert_allclose(matrix, np.eye(len(geo.rho_face_norm)))
+
+  def test_build_pedestal_smoothing_matrix_isolation_and_row_sums(self):
+    config = default_configs.get_default_config_dict()
+    config['transport'] = {'pedestal_smoothing_width': 0.05}
+    torax_config = model_config.ToraxConfig.from_dict(config)
+    runtime_params = build_runtime_params.RuntimeParamsProvider.from_config(
+        torax_config
+    )(t=torax_config.numerics.t_initial)
+    geo = torax_config.geometry.build_provider(
+        t=torax_config.numerics.t_initial
+    )
+    transition_state = _make_transition_state(rho_norm_ped_top=0.8)
+    matrix = transport_model._build_pedestal_smoothing_matrix(
+        runtime_params.transport,
+        geo,
+        transition_state,
+    )
+
+    core_mask = geo.rho_face_norm < 0.8
+    ped_mask = geo.rho_face_norm >= 0.8
+    # Core rows should be identity.
+    for idx in np.where(core_mask)[0]:
+      expected_row = np.zeros(len(geo.rho_face_norm))
+      expected_row[idx] = 1.0
+      np.testing.assert_allclose(matrix[idx], expected_row, atol=1e-7)
+
+    # Pedestal rows should have zero weights in the core.
+    np.testing.assert_allclose(matrix[ped_mask][:, core_mask], 0.0)
+    # Row sums in the pedestal region should normalize to 1.
+    np.testing.assert_allclose(np.sum(matrix[ped_mask], axis=1), 1.0, atol=1e-6)
+
   def test_smoothing_zones(self):
     """Tests that smoothing_zones smoothes transport coefficients in the specified region."""
     config = default_configs.get_default_config_dict()
@@ -953,6 +1001,7 @@ def test_call_invokes_constituent_models_and_combines(self):
     combined_params.V_e_min = -100.0
     combined_params.V_e_max = 100.0
     combined_params.smoothing_width = 0.0
+    combined_params.pedestal_smoothing_width = 0.0
     combined_params.smoothing_zones = ()
 
     transition_state = _make_transition_state(rho_norm_ped_top=1.0)

torax/_src/transport_model/transport_model.py

@@ -87,25 +87,37 @@ def __call__(
         True,
     )
 
+    core_smoothing_matrix = _build_smoothing_matrix(
+        runtime_params.transport,
+        runtime_params,
+        geo,
+        pedestal_transition_state,
+    )
+    pedestal_smoothing_matrix = _build_pedestal_smoothing_matrix(
+        runtime_params.transport,
+        geo,
+        pedestal_transition_state,
+    )
+
     core, core_components = self._compute_domain_coeffs(
         self.core_transport_models,
         runtime_params.transport.core_transport_model_params,
         runtime_params,
         geo,
         core_profiles,
-        pedestal_transition_state,
         two_point_mask,
         domain_mask=core_domain_mask,
+        smoothing_matrix=core_smoothing_matrix,
     )
     pedestal, pedestal_components = self._compute_domain_coeffs(
         self.pedestal_transport_models,
         runtime_params.transport.pedestal_transport_model_params,
         runtime_params,
         geo,
         core_profiles,
-        pedestal_transition_state,
         two_point_mask,
         domain_mask=pedestal_mask,
+        smoothing_matrix=pedestal_smoothing_matrix,
     )
 
     return transport_coeffs.TurbulentTransport(
@@ -124,11 +136,9 @@ def _compute_domain_coeffs(
       runtime_params: runtime_params_lib.RuntimeParams,
       geo: geometry.Geometry,
       core_profiles: state.CoreProfiles,
-      pedestal_transition_state: (
-          pedestal_transition_state_lib.PedestalTransitionState
-      ),
       two_point_mask: array_typing.BoolVectorFace,
       domain_mask: jax.Array,
+      smoothing_matrix: jax.Array,
   ) -> tuple[
       transport_coeffs.TransportCoeffs,
       dict[str, transport_coeffs.TransportCoeffs],
@@ -152,10 +162,8 @@ def _compute_domain_coeffs(
         clipped,
     )
     smoothed = self._smooth_coeffs(
-        runtime_params,
-        geo,
         masked,
-        pedestal_transition_state,
+        smoothing_matrix,
     )
     return smoothed, model_outputs
 
@@ -269,21 +277,10 @@ def _apply_clipping(
 
   def _smooth_coeffs(
       self,
-      runtime_params: runtime_params_lib.RuntimeParams,
-      geo: geometry.Geometry,
       input_coeffs: transport_coeffs.TransportCoeffs,
-      pedestal_transition_state: (
-          pedestal_transition_state_lib.PedestalTransitionState
-      ),
+      smoothing_matrix: jax.Array,
   ) -> transport_coeffs.TransportCoeffs:
     """Gaussian smoothing of turbulent transport coefficients."""
-    smoothing_matrix = _build_smoothing_matrix(
-        runtime_params.transport,
-        runtime_params,
-        geo,
-        pedestal_transition_state,
-    )
-
     # Iterate over fields of the CoreTransport dataclass.
     # Ignore optional fields that are made all zero in post_init.
     def smooth_single_coeff(coeff):
@@ -325,12 +322,6 @@ def _build_smoothing_matrix(
     ),
 ) -> jax.Array:
   """Builds a smoothing matrix for the transport model."""
-  # To reduce the range of the convolution, weights under lower_cutoff are
-  # clipped to zero.
-  lower_cutoff = 0.01
-  # used for eps, small number to avoid divisions by zero for sigma = 0
-  consts = constants.CONSTANTS
-
   # 1. Build smoothing width profile
   # Spatially-varying sigma(rho) is constructed across radial grid points.
   # This supports multiple smoothing_zones with distinct widths.
@@ -372,6 +363,39 @@ def build_profile_fallback():
       smoothing_width_profile,
   )
 
+  return _build_smoothing_matrix_from_profile(geo, smoothing_width_profile)
+
+
+def _build_pedestal_smoothing_matrix(
+    transport_runtime_params: transport_runtime_params_lib.RuntimeParams,
+    geo: geometry.Geometry,
+    pedestal_transition_state: (
+        pedestal_transition_state_lib.PedestalTransitionState
+    ),
+) -> jax.Array:
+  """Builds a smoothing matrix localized to the pedestal region."""
+  rho_norm_ped_top = (
+      pedestal_transition_state.pedestal_model_output.rho_norm_ped_top
+  )
+  smoothing_width_profile = jnp.where(
+      geo.rho_face_norm >= rho_norm_ped_top,
+      transport_runtime_params.pedestal_smoothing_width,
+      0.0,
+  )
+  return _build_smoothing_matrix_from_profile(geo, smoothing_width_profile)
+
+
+def _build_smoothing_matrix_from_profile(
+    geo: geometry.Geometry,
+    smoothing_width_profile: jax.Array,
+) -> jax.Array:
+  """Builds a Gaussian convolution smoothing matrix from a width profile."""
+  # To reduce the range of the convolution, weights under lower_cutoff are
+  # clipped to zero.
+  lower_cutoff = 0.01
+  # used for eps, small number to avoid divisions by zero for sigma = 0
+  consts = constants.CONSTANTS
+
   # 2. Kernel matrix with variable width (sigma_i for each destination row i)
   r_diff = geo.rho_face_norm[:, jnp.newaxis] - geo.rho_face_norm
   sigma = smoothing_width_profile[:, jnp.newaxis]