Response to review

* Restore 'core_transport_models' -> 'transport_models'
* Support multiple pedestal transport models, defaulting to an empty list
* Check that rho_min/max and inner/outer patches not set on pedestal models
* Update docs and image
* Fix a test that was still using AssertRaises rather than AssertRaisesRegex, and was covering the wrong thing
* Improve readability/code reuse in summing over transport models; added comment about future vmap
* Add test to cover behaviour when no pedestal model supplied

Author: theo-brown

docs/configuration.rst

@@ -1069,17 +1069,25 @@ combined
 
 A combined (additive) model, where the total transport coefficients are
 calculated by summing contributions from a list of component models. Each
-component model is active only within its defined radial domain, set using
-``rho_min``` and ``rho_max``. These zones can be overlapping or
-non-overlapping; in regions of overlap, the total transport coefficients are
-computed by adding the contributions from component models active at those
-coordinates. Post-processing (clipping and smoothing) is performed on the
-summed value.
+component model is active only within its defined radial domain, which can
+be overlapping or non-overlapping; in regions of overlap, the total
+transport coefficients are computed by adding the contributions from
+component models active at those coordinates.
+For models defined in ``transport_models``, the active domain is set by
+``rho_min``` and ``rho_max``. For models in ``pedestal_transport_models``,
+the active domain is set by the ``rho_norm_ped_top`` parameter from the
+``pedestal`` section of the config.
+Post-processing (clipping and smoothing) is performed on the summed
+values from all component models, including in the pedestal.
 
 The runtime parameters are as follows.
 
 ``transport_models`` (list[dict])
-  A list containing config dicts for the component transport models.
+  A list containing config dicts for the component models for turbulent transport in the core.
+
+``pedestal_transport_models`` (list[dict])
+  A list containing config dicts for the component models for turbulent transport in the pedestal.
+
 
    .. warning::
     TORAX will throw a ``ValueError`` if any of the component transport
@@ -1094,29 +1102,35 @@ The runtime parameters are as follows.
 Example:
 
 .. code-block:: python
-
+    ...
     'transport': {
-      'model_name': 'combined',
-      'transport_models': [
-        {
-          'model_name': 'constant',
-            'chi_i': 1.0,
-            'rho_max': 0.3,
-        },
-        {
-          'model_name': 'constant',
-            'chi_i': 2.0,
-            'rho_min': 0.2
-            'rho_max': 0.5,
-        },
-        {
-          'model_name': 'constant',
-            'chi_i': 0.5,
-            'rho_min': 0.5
-            'rho_max': 1.0,
-        },
-      ],
-    }
+        'model_name': 'combined',
+        'transport_models': [
+            {
+               'model_name': 'constant',
+               'chi_i': 1.0,
+               'rho_max': 0.3,
+            },
+            {
+                'model_name': 'constant',
+                'chi_i': 2.0,
+                'rho_min': 0.2,
+            },
+        ],
+        'pedestal_transport_models': [
+            {
+                'model_name': 'constant',
+                'chi_i': 0.5,
+            },
+        ],
+      },
+      'pedestal': {
+          'model_name': 'set_T_ped_n_ped',
+          'set_pedestal': True,
+          'n_e_ped': 0.8,
+          'n_e_ped_is_fGW': True,
+      },
+      ...
 
 This would produce a ``chi_i`` profile that looks like the following.
 
@@ -1127,8 +1141,8 @@ This would produce a ``chi_i`` profile that looks like the following.
 Note that in the region :math:`[0, 0.2]`, only the first component is active,
 so ``chi_i = 1.0``. In :math:`(0.2, 0.3]` the first two components are both
 active, leading to a combined value of ``chi_i = 3.0``. In :math:`(0.3, 0.5]`,
-only the second model is active (``chi_i = 2.0``), and in :math:`(0.5, 1.0]`
-only the fourth model is active (``chi_i = 0.5``).
+only the second model is active (``chi_i = 2.0``), and in :math:`(0.8, 1.0]`
+only the pedestal transport model is active (``chi_i = 0.5``).
 
 
 sources

docs/images/combined_transport_example.png

@@ -14,8 +14,10 @@
 
 """Script for plotting the combined transport model in the docs."""
 from typing import Sequence
-from absl import app
+
 import matplotlib.pyplot as plt
+from absl import app
+
 import torax
 from torax._src.torax_pydantic import model_config
 
@@ -33,13 +35,18 @@ def main(argv: Sequence[str]) -> None:
           'geometry_type': 'circular',
           'n_rho': 30,  # for higher resolution plotting
       },
-      'pedestal': {},
+      'pedestal': {
+          'model_name': 'set_T_ped_n_ped',
+          'set_pedestal': True,
+          'n_e_ped': 0.8,
+          'n_e_ped_is_fGW': True,
+      },
       'neoclassical': {},
       'sources': {},
       'solver': {},
       'transport': {
           'model_name': 'combined',
-          'core_transport_models': [
+          'transport_models': [
               {
                   'model_name': 'constant',
                   'chi_i': 1.0,
@@ -49,15 +56,14 @@ def main(argv: Sequence[str]) -> None:
                   'model_name': 'constant',
                   'chi_i': 2.0,
                   'rho_min': 0.2,
-                  'rho_max': 0.5,
               },
           ],
-          'pedestal_transport_model': {
-              'model_name': 'constant',
-              'chi_i': 0.5,
-              'rho_min': 0.5,
-              'rho_max': 1.0,
-          },
+          'pedestal_transport_models': [
+              {
+                  'model_name': 'constant',
+                  'chi_i': 0.5,
+              },
+          ],
       },
   }
   torax_config = model_config.ToraxConfig.from_dict(config)

docs/scripts/combined_transport_example.py

@@ -17,8 +17,7 @@
 A class for combining transport models.
 """
 import dataclasses
-from typing import Sequence
-import dataclasses
+from typing import Callable, Sequence
 
 import jax
 import jax.numpy as jnp
@@ -35,7 +34,7 @@
 @jax.tree_util.register_dataclass
 @dataclasses.dataclass(frozen=True)
 class DynamicRuntimeParams(runtime_params_lib.DynamicRuntimeParams):
-  core_transport_model_params: Sequence[runtime_params_lib.DynamicRuntimeParams]
+  transport_model_params: Sequence[runtime_params_lib.DynamicRuntimeParams]
   pedestal_transport_model_params: Sequence[
       runtime_params_lib.DynamicRuntimeParams
   ]
@@ -46,12 +45,12 @@ class CombinedTransportModel(transport_model_lib.TransportModel):
 
   def __init__(
       self,
-      core_transport_models: Sequence[transport_model_lib.TransportModel],
-      pedestal_transport_model: transport_model_lib.TransportModel,
+      transport_models: Sequence[transport_model_lib.TransportModel],
+      pedestal_transport_models: Sequence[transport_model_lib.TransportModel],
   ):
     super().__init__()
-    self.core_transport_models = core_transport_models
-    self.pedestal_transport_model = pedestal_transport_model
+    self.transport_models = transport_models
+    self.pedestal_transport_models = pedestal_transport_models
     self._frozen = True
 
   def __call__(
@@ -128,72 +127,57 @@ def _call_implementation(
     # Required for pytype
     assert isinstance(transport_dynamic_runtime_params, DynamicRuntimeParams)
 
-    # Core transport
-    core_transport_coeffs_list = []
-    for component_model, component_params in zip(
-        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 combined output rather than the individual components.
+    def apply_and_restrict(
+        component_model: transport_model_lib.TransportModel,
+        component_params: runtime_params_lib.DynamicRuntimeParams,
+        restriction_fn: Callable,
+    ) -> transport_model_lib.TurbulentTransport:
+      # TODO: Consider only computing transport coefficients for the active
+      # domain, rather than masking them out later. This could be significantly
+      # more efficient especially for pedestal models, as these are only active
+      # in a small region of the domain.
       component_transport_coeffs = component_model._call_implementation(
           component_params,
           dynamic_runtime_params_slice,
           geo,
           core_profiles,
           pedestal_model_output,
       )
-
-      # Apply domain restriction
-      # This is a property of each component_model, so needs to be applied
-      # at the component model level rather than the global level
-      component_transport_coeffs = component_model._apply_domain_restriction(
+      component_transport_coeffs = restriction_fn(
           component_params,
           geo,
           component_transport_coeffs,
           pedestal_model_output,
       )
+      return component_transport_coeffs
 
-      core_transport_coeffs_list.append(component_transport_coeffs)
-
-    # Pedestal transport
-    pedestal_transport_coeffs = (
-        self.pedestal_transport_model._call_implementation(
+    pedestal_coeffs = [
+        apply_and_restrict(
+            model, params, self._apply_pedestal_domain_restriction
+        )
+        for model, params in zip(
+            self.pedestal_transport_models,
             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,
-    )
+    ]
+
+    core_coeffs = [
+        apply_and_restrict(model, params, model._apply_domain_restriction)
+        for model, params in zip(
+            self.transport_models,
+            transport_dynamic_runtime_params.transport_model_params,
+        )
+    ]
 
     # Combine the transport coefficients from core and pedestal models.
     combined_transport_coeffs = jax.tree.map(
         lambda *leaves: sum(leaves),
-        *(core_transport_coeffs_list + [pedestal_transport_coeffs]),
+        *pedestal_coeffs,
+        *core_coeffs,
     )
 
     return combined_transport_coeffs
 
-  def __hash__(self):
-    return hash(
-        tuple(self.core_transport_models + [self.pedestal_transport_model])
-    )
-
-  def __eq__(self, other):
-    return (
-        isinstance(other, CombinedTransportModel)
-        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,
@@ -215,3 +199,13 @@ def _apply_pedestal_domain_restriction(
         d_face_el=d_face_el,
         v_face_el=v_face_el,
     )
+
+  def __hash__(self):
+    return hash(tuple(self.transport_models + self.pedestal_transport_models))
+
+  def __eq__(self, other):
+    return (
+        isinstance(other, CombinedTransportModel)
+        and self.transport_models == other.transport_models
+        and self.pedestal_transport_models == other.pedestal_transport_models
+    )