Add pedestal transport model section to CombinedTransportModel

Author: theo-brown

docs/scripts/combined_transport_example.py

@@ -39,7 +39,7 @@ def main(argv: Sequence[str]) -> None:
       'solver': {},
       'transport': {
           'model_name': 'combined',
-          'transport_models': [
+          'core_transport_models': [
               {
                   'model_name': 'constant',
                   'chi_i': 1.0,
@@ -51,13 +51,13 @@ def main(argv: Sequence[str]) -> None:
                   'rho_min': 0.2,
                   'rho_max': 0.5,
               },
-              {
-                  'model_name': 'constant',
-                  'chi_i': 0.5,
-                  'rho_min': 0.5,
-                  'rho_max': 1.0,
-              },
           ],
+          'pedestal_transport_model': {
+              'model_name': 'constant',
+              'chi_i': 0.5,
+              'rho_min': 0.5,
+              'rho_max': 1.0,
+          },
       },
   }
   torax_config = model_config.ToraxConfig.from_dict(config)

torax/_src/transport_model/combined.py

@@ -1,4 +1,4 @@
-# Copyright 2024 DeepMind Technologies Limited
+7  # Copyright 2024 DeepMind Technologies Limited
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -18,9 +18,11 @@
 """
 
 from typing import Sequence
+import dataclasses
 
 import chex
 import jax
+import jax.numpy as jnp
 from torax._src import state
 from torax._src.config import runtime_params_slice
 from torax._src.geometry import geometry
@@ -33,19 +35,72 @@
 
 @chex.dataclass(frozen=True)
 class DynamicRuntimeParams(runtime_params_lib.DynamicRuntimeParams):
-  transport_model_params: Sequence[runtime_params_lib.DynamicRuntimeParams]
+  core_transport_model_params: Sequence[runtime_params_lib.DynamicRuntimeParams]
+  pedestal_transport_model_params: Sequence[
+      runtime_params_lib.DynamicRuntimeParams
+  ]
 
 
 class CombinedTransportModel(transport_model_lib.TransportModel):
   """Combines coefficients from a list of transport models."""
 
   def __init__(
-      self, transport_models: Sequence[transport_model_lib.TransportModel]
+      self,
+      core_transport_models: Sequence[transport_model_lib.TransportModel],
+      pedestal_transport_model: transport_model_lib.TransportModel,
   ):
     super().__init__()
-    self.transport_models = transport_models
+    self.core_transport_models = core_transport_models
+    self.pedestal_transport_model = pedestal_transport_model
     self._frozen = True
 
+  def __call__(
+      self,
+      dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
+      geo: geometry.Geometry,
+      core_profiles: state.CoreProfiles,
+      pedestal_model_output: pedestal_model_lib.PedestalModelOutput,
+  ) -> transport_model_lib.TurbulentTransport:
+    if not getattr(self, "_frozen", False):
+      raise RuntimeError(
+          f"Subclass implementation {type(self)} forgot to "
+          "freeze at the end of __init__."
+      )
+
+    transport_runtime_params = dynamic_runtime_params_slice.transport
+
+    # Calculate the transport coefficients - includes contribution from pedestal
+    # and core transport models.
+    transport_coeffs = self._call_implementation(
+        transport_runtime_params,
+        dynamic_runtime_params_slice,
+        geo,
+        core_profiles,
+        pedestal_model_output,
+    )
+
+    # In contrast to the base TransportModel, we do not apply domain restriction
+
+    # Apply min/max clipping
+    transport_coeffs = self._apply_clipping(
+        transport_runtime_params,
+        transport_coeffs,
+    )
+
+    # In contrast to the base TransportModel, we do not apply patches
+
+    # Return smoothed coefficients if smoothing is enabled
+    # TODO: what should be done about masking? In base TransportModel, the
+    # pedestal and any patches are masked out before smoothing. May require a
+    # custom implementation of _smooth_coeffs for CombinedTransportModel.
+    return self._smooth_coeffs(
+        transport_runtime_params,
+        dynamic_runtime_params_slice,
+        geo,
+        transport_coeffs,
+        pedestal_model_output,
+    )
+
   def _call_implementation(
       self,
       transport_dynamic_runtime_params: runtime_params_lib.DynamicRuntimeParams,
@@ -72,16 +127,15 @@ def _call_implementation(
     # Required for pytype
     assert isinstance(transport_dynamic_runtime_params, DynamicRuntimeParams)
 
-    component_transport_coeffs_list = []
-
+    # Core transport
+    core_transport_coeffs_list = []
     for component_model, component_params in zip(
-        self.transport_models,
-        transport_dynamic_runtime_params.transport_model_params,
+        self.core_transport_models,
+        transport_dynamic_runtime_params.core_transport_model_params,
     ):
       # Use the component model's _call_implementation, rather than __call__
       # directly. This ensures postprocessing (clipping, smoothing, patches) are
-      # performed on the output of CombinedTransportModel rather than its
-      # component models.
+      # performed on the combined output rather than the individual components.
       component_transport_coeffs = component_model._call_implementation(
           component_params,
           dynamic_runtime_params_slice,
@@ -100,20 +154,63 @@ def _call_implementation(
           pedestal_model_output,
       )
 
-      component_transport_coeffs_list.append(component_transport_coeffs)
+      core_transport_coeffs_list.append(component_transport_coeffs)
+
+    # Pedestal transport
+    pedestal_transport_coeffs = (
+        self.pedestal_transport_model._call_implementation(
+            transport_dynamic_runtime_params.pedestal_transport_model_params,
+            dynamic_runtime_params_slice,
+            geo,
+            core_profiles,
+            pedestal_model_output,
+        )
+    )
+    pedestal_transport_coeffs = self._apply_pedestal_domain_restriction(
+        transport_dynamic_runtime_params.pedestal_transport_model_params,
+        geo,
+        pedestal_transport_coeffs,
+        pedestal_model_output,
+    )
 
+    # Combine the transport coefficients from core and pedestal models.
     combined_transport_coeffs = jax.tree.map(
         lambda *leaves: sum(leaves),
-        *component_transport_coeffs_list,
+        *(core_transport_coeffs_list + [pedestal_transport_coeffs]),
     )
 
     return combined_transport_coeffs
 
   def __hash__(self):
-    return hash(tuple(self.transport_models))
+    return hash(
+        tuple(self.core_transport_models + [self.pedestal_transport_model])
+    )
 
   def __eq__(self, other):
     return (
         isinstance(other, CombinedTransportModel)
-        and self.transport_models == other.transport_models
+        and self.core_transport_models == other.core_transport_models
+        and self.pedestal_transport_model == other.pedestal_transport_model
+    )
+
+  def _apply_pedestal_domain_restriction(
+      self,
+      transport_runtime_params: runtime_params_lib.DynamicRuntimeParams,
+      geo: geometry.Geometry,
+      transport_coeffs: transport_model_lib.TurbulentTransport,
+      pedestal_model_output: pedestal_model_lib.PedestalModelOutput,
+  ) -> transport_model_lib.TurbulentTransport:
+    active_mask = geo.rho_face_norm > pedestal_model_output.rho_norm_ped_top
+
+    chi_face_ion = jnp.where(active_mask, transport_coeffs.chi_face_ion, 0.0)
+    chi_face_el = jnp.where(active_mask, transport_coeffs.chi_face_el, 0.0)
+    d_face_el = jnp.where(active_mask, transport_coeffs.d_face_el, 0.0)
+    v_face_el = jnp.where(active_mask, transport_coeffs.v_face_el, 0.0)
+
+    return dataclasses.replace(
+        transport_coeffs,
+        chi_face_ion=chi_face_ion,
+        chi_face_el=chi_face_el,
+        d_face_el=d_face_el,
+        v_face_el=v_face_el,
     )

torax/_src/transport_model/pydantic_model.py

@@ -368,50 +368,69 @@ class CombinedTransportModel(pydantic_model_base.TransportBase):
 
   Attributes:
     model_name: The transport model to use. Hardcoded to 'combined'.
-    transport_models: A sequence of transport models, whose outputs will be
-      summed to give the combined transport coefficients.
+    core_transport_models: A sequence of transport models, whose outputs will be
+      summed to give the combined core transport coefficients.
+    pedestal_transport_model: A model for pedestal transport coefficients.
   """
 
-  transport_models: Sequence[CombinedCompatibleTransportModel] = pydantic.Field(
+  core_transport_models: Sequence[
+      CombinedCompatibleTransportModel
+  ] = pydantic.Field(
       default_factory=list
   )  # pytype: disable=invalid-annotation
+  pedestal_transport_model: CombinedCompatibleTransportModel
   model_name: Literal['combined'] = 'combined'
 
   def build_transport_model(self) -> combined.CombinedTransportModel:
-    model_list = [
-        model.build_transport_model() for model in self.transport_models
+    core_transport_model_list = [
+        model.build_transport_model() for model in self.core_transport_models
     ]
-    return combined.CombinedTransportModel(transport_models=model_list)
+    pedestal_transport_model = (
+        self.pedestal_transport_model.build_transport_model()
+    )
+
+    return combined.CombinedTransportModel(
+        core_transport_models=core_transport_model_list,
+        pedestal_transport_model=pedestal_transport_model,
+    )
 
   def build_dynamic_params(
       self, t: chex.Numeric
   ) -> combined.DynamicRuntimeParams:
     base_kwargs = dataclasses.asdict(super().build_dynamic_params(t))
-    model_params_list = [
-        model.build_dynamic_params(t) for model in self.transport_models
+    core_transport_model_params = [
+        model.build_dynamic_params(t) for model in self.core_transport_models
     ]
+    pedestal_transport_model_params = (
+        self.pedestal_transport_model.build_dynamic_params(t)
+    )
+
     return combined.DynamicRuntimeParams(
-        transport_model_params=model_params_list,
+        core_transport_model_params=core_transport_model_params,
+        pedestal_transport_model_params=pedestal_transport_model_params,
         **base_kwargs,
     )
 
   @pydantic.model_validator(mode='after')
   def _check_fields(self) -> typing_extensions.Self:
     super()._check_fields()
-    if not self.transport_models:
+    if not self.core_transport_models:
       raise ValueError(
-          'transport_models cannot be empty for CombinedTransportModel. '
+          'core_transport_models cannot be empty for CombinedTransportModel. '
           'Please provide at least one transport model configuration.'
       )
-    if any([
-        np.any(model.apply_inner_patch.value)
-        or np.any(model.apply_outer_patch.value)
-        for model in self.transport_models
-    ]):
+    if (
+        any([
+            np.any(model.apply_inner_patch.value)
+            or np.any(model.apply_outer_patch.value)
+            for model in self.core_transport_models
+        ])
+        or np.any(self.apply_inner_patch.value)
+        or np.any(self.apply_outer_patch.value)
+    ):
       raise ValueError(
-          'apply_inner_patch and apply_outer_patch and should be set in the'
-          ' config for CombinedTransportModel only, rather than its component'
-          ' models.'
+          'apply_inner_patch and apply_outer_patch not supported for'
+          ' CombinedTransportModel or its component models.'
       )
     if np.any(self.rho_min.value != 0.0) or np.any(self.rho_max.value != 1.0):
       raise ValueError(
@@ -422,4 +441,3 @@ def _check_fields(self) -> typing_extensions.Self:
 
 
 TransportConfig = CombinedTransportModel | CombinedCompatibleTransportModel  # pytype: disable=invalid-annotation
-

torax/_src/transport_model/tests/combined_test.py

@@ -31,7 +31,7 @@ def test_call_implementation(self):
     config = default_configs.get_default_config_dict()
     config['transport'] = {
         'model_name': 'combined',
-        'transport_models': [
+        'core_transport_models': [
             {'model_name': 'constant', 'rho_max': 0.2, 'chi_i': 1.0},
             {
                 'model_name': 'constant',
@@ -41,6 +41,7 @@ def test_call_implementation(self):
             },
             {'model_name': 'constant', 'rho_min': 0.5, 'chi_i': 3.0},
         ],
+        'pedestal_transport_model': {'model_name': 'constant', 'chi_i': 0.1},
     }
     torax_config = model_config.ToraxConfig.from_dict(config)
     model = torax_config.transport.build_transport_model()
@@ -80,12 +81,12 @@ def test_call_implementation(self):
         mock_pedestal_outputs,
     )
     # Target:
-    # - 0 for rho = [rho_ped_top, rho_max]
+    # - 0.1 for rho = [rho_ped_top, rho_max]
     # - 3 for rho = (0.8, rho_ped_top), to check pedestal overrides it
     # - 5 for rho = (0.5, 0.8], to check case where models overlap
     # - 2 for rho = (0.2, 0.5], to check case rho_min_1 == rho_max_2
     # - 1 for rho = [0, 0.2], to check case where rho_min = 0
-    target = jnp.where(geo.rho_face_norm <= 0.91, 3.0, 0.0)
+    target = jnp.where(geo.rho_face_norm <= 0.91, 3.0, 0.1)
     target = jnp.where(geo.rho_face_norm <= 0.8, 5.0, target)
     target = jnp.where(geo.rho_face_norm <= 0.5, 2.0, target)
     target = jnp.where(geo.rho_face_norm <= 0.2, 1.0, target)
@@ -95,10 +96,27 @@ def test_error_if_patches_set_on_children(self):
     config = default_configs.get_default_config_dict()
     config['transport'] = {
         'model_name': 'combined',
-        'transport_models': [
+        'core_transport_models': [
             {'model_name': 'constant', 'apply_inner_patch': True},
             {'model_name': 'constant'},
         ],
+        'pedestal_transport_model': {'model_name': 'constant'},
+    }
+    with self.assertRaisesRegex(
+        ValueError, '(?=.*patch)(?=.*CombinedTransportModel)'
+    ):
+      model_config.ToraxConfig.from_dict(config)
+
+  def test_error_if_patches_set_on_self(self):
+    config = default_configs.get_default_config_dict()
+    config['transport'] = {
+        'model_name': 'combined',
+        'core_transport_models': [
+            {'model_name': 'constant'},
+            {'model_name': 'constant'},
+        ],
+        'pedestal_transport_model': {'model_name': 'constant'},
+        'apply_inner_patch': True,
     }
     with self.assertRaisesRegex(
         ValueError, '(?=.*patch)(?=.*CombinedTransportModel)'
@@ -109,10 +127,11 @@ def test_error_if_rho_min_or_rho_max_set(self):
     config = default_configs.get_default_config_dict()
     config['transport'] = {
         'model_name': 'combined',
-        'transport_models': [
+        'core_transport_models': [
             {'model_name': 'constant'},
             {'model_name': 'constant'},
         ],
+        'pedestal_transport_model': {'model_name': 'constant'},
         'rho_min': 0.1,
     }
     with self.assertRaisesRegex(
@@ -122,10 +141,11 @@ def test_error_if_rho_min_or_rho_max_set(self):
 
       config['transport'] = {
           'model_name': 'combined',
-          'transport_models': [
+          'core_transport_models': [
               {'model_name': 'constant'},
               {'model_name': 'constant'},
           ],
+          'pedestal_transport_model': {'model_name': 'constant'},
           'rho_max': 0.9,
       }
     with self.assertRaises(ValueError):