diff --git a/docs/configuration.rst b/docs/configuration.rst index c4d39959e..64bbddf5a 100644 --- a/docs/configuration.rst +++ b/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 diff --git a/docs/physics_models.rst b/docs/physics_models.rst index 9ee03738b..ca2d3707f 100644 --- a/docs/physics_models.rst +++ b/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 diff --git a/torax/_src/transport_model/pydantic_model.py b/torax/_src/transport_model/pydantic_model.py index eae360006..f0aaea768 100644 --- a/torax/_src/transport_model/pydantic_model.py +++ b/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( diff --git a/torax/_src/transport_model/runtime_params.py b/torax/_src/transport_model/runtime_params.py index 5c0b2bd91..8c3df0111 100644 --- a/torax/_src/transport_model/runtime_params.py +++ b/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, ...] diff --git a/torax/_src/transport_model/tests/transport_model_test.py b/torax/_src/transport_model/tests/transport_model_test.py index 1ca6f56f9..636f28656 100644 --- a/torax/_src/transport_model/tests/transport_model_test.py +++ b/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) diff --git a/torax/_src/transport_model/transport_model.py b/torax/_src/transport_model/transport_model.py index 719ef43be..f89e1cb6c 100644 --- a/torax/_src/transport_model/transport_model.py +++ b/torax/_src/transport_model/transport_model.py @@ -87,15 +87,27 @@ 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, @@ -103,9 +115,9 @@ def __call__( 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]