Implement ADD and OVERWRITE merge modes in CombinedTransportModel.

Drive-by: fixed a bug where previously sub-channels (like chi_bohm_e) were leaking outside their configured radial domain in combined_transport_model.

The combined transport model now supports two merge modes for combining coefficients from multiple sub-models: ADD and OVERWRITE. OVERWRITE mode replaces existing coefficients in the model's active domain and locks those regions, preventing subsequent ADD models from contributing. Disabled channels in an OVERWRITE model are transparent and do not affect the accumulated coefficients. The implementation uses an accumulator and lock system within a new _accumulate helper function.

PiperOrigin-RevId: 858780479

Author: jcitrin

torax/_src/output_tools/output.py

@@ -659,16 +659,17 @@ def _save_core_transport(
 
     # Save optional BohmGyroBohm attributes if present.
     core_transport = self._stacked_core_transport
-    if (
-        core_transport.chi_face_el_bohm is not None
-        or core_transport.chi_face_el_gyrobohm is not None
-        or core_transport.chi_face_ion_bohm is not None
-        or core_transport.chi_face_ion_gyrobohm is not None
-    ):
-      xr_dict[CHI_BOHM_E] = core_transport.chi_face_el_bohm
-      xr_dict[CHI_GYROBOHM_E] = core_transport.chi_face_el_gyrobohm
-      xr_dict[CHI_BOHM_I] = core_transport.chi_face_ion_bohm
-      xr_dict[CHI_GYROBOHM_I] = core_transport.chi_face_ion_gyrobohm
+    optional_transport_map = {
+        CHI_BOHM_E: core_transport.chi_face_el_bohm,
+        CHI_GYROBOHM_E: core_transport.chi_face_el_gyrobohm,
+        CHI_BOHM_I: core_transport.chi_face_ion_bohm,
+        CHI_GYROBOHM_I: core_transport.chi_face_ion_gyrobohm,
+    }
+
+    for name, data in optional_transport_map.items():
+      # Skip if None or an array of Nones from stack
+      if data is not None and data.dtype != object:
+        xr_dict[name] = data
 
     xr_dict = {
         name: self._pack_into_data_array(

torax/_src/transport_model/bohm_gyrobohm.py

@@ -50,7 +50,7 @@ class RuntimeParams(transport_runtime_params_lib.RuntimeParams):
 class BohmGyroBohmTransportModel(transport_model_lib.TransportModel):
   """Calculates various coefficients related to particle transport according to the Bohm + gyro-Bohm Model."""
 
-  def _call_implementation(
+  def call_implementation(
       self,
       transport_runtime_params: transport_runtime_params_lib.RuntimeParams,
       runtime_params: runtime_params_lib.RuntimeParams,

torax/_src/transport_model/combined.py

@@ -19,13 +19,14 @@
 
 import dataclasses
 from typing import Callable, Sequence
-
 import jax
 import jax.numpy as jnp
+from torax._src import jax_utils
 from torax._src import state
 from torax._src.config import runtime_params as runtime_params_lib
 from torax._src.geometry import geometry
 from torax._src.pedestal_model import pedestal_model as pedestal_model_lib
+from torax._src.transport_model import enums
 from torax._src.transport_model import runtime_params as transport_runtime_params_lib
 from torax._src.transport_model import transport_model as transport_model_lib
 
@@ -60,7 +61,7 @@ def __call__(
 
     # Calculate the transport coefficients - includes contribution from pedestal
     # and core transport models.
-    transport_coeffs = self._call_implementation(
+    transport_coeffs = self.call_implementation(
         transport_runtime_params,
         runtime_params,
         geo,
@@ -70,7 +71,7 @@ def __call__(
 
     # In contrast to the base TransportModel, we do not apply domain restriction
     # or output masking (enabled/disabled channels) as these are handled at the
-    # component model level in _call_implementation here.
+    # component model level in call_implementation here.
 
     # Apply min/max clipping
     transport_coeffs = self._apply_clipping(
@@ -93,7 +94,7 @@ def __call__(
         pedestal_model_output,
     )
 
-  def _call_implementation(
+  def call_implementation(
       self,
       transport_runtime_params: transport_runtime_params_lib.RuntimeParams,
       runtime_params: runtime_params_lib.RuntimeParams,
@@ -117,95 +118,142 @@ def _call_implementation(
     # Required for pytype
     assert isinstance(transport_runtime_params, RuntimeParams)
 
-    def apply_and_restrict(
-        component_model: transport_model_lib.TransportModel,
-        component_params: transport_runtime_params_lib.RuntimeParams,
-        restriction_fn: Callable[
-            [
-                transport_runtime_params_lib.RuntimeParams,
-                geometry.Geometry,
-                transport_model_lib.TurbulentTransport,
-                pedestal_model_lib.PedestalModelOutput,
-            ],
-            transport_model_lib.TurbulentTransport,
-        ],
-    ) -> transport_model_lib.TurbulentTransport:
-      # TODO(b/434175682): 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,
-          runtime_params,
-          geo,
-          core_profiles,
-          pedestal_model_output,
-      )
-      component_transport_coeffs = component_model._apply_output_mask(
-          component_params,
-          component_transport_coeffs,
-      )
-      component_transport_coeffs = restriction_fn(
-          component_params,
-          geo,
-          component_transport_coeffs,
-          pedestal_model_output,
-      )
-      return component_transport_coeffs
-
-    pedestal_coeffs = [
-        apply_and_restrict(
-            model, params, self._apply_pedestal_domain_restriction
-        )
-        for model, params in zip(
-            self.pedestal_transport_models,
-            transport_runtime_params.pedestal_transport_model_params,
-        )
-    ]
+    core_coeffs = self._combine(
+        self.transport_models,
+        transport_runtime_params.transport_model_params,
+        runtime_params,
+        geo,
+        core_profiles,
+        pedestal_model_output,
+        transport_model_lib.compute_core_domain_mask,
+    )
 
-    core_coeffs = [
-        apply_and_restrict(model, params, model._apply_domain_restriction)
-        for model, params in zip(
-            self.transport_models,
-            transport_runtime_params.transport_model_params,
-        )
-    ]
+    pedestal_coeffs = self._combine(
+        self.pedestal_transport_models,
+        transport_runtime_params.pedestal_transport_model_params,
+        runtime_params,
+        geo,
+        core_profiles,
+        pedestal_model_output,
+        _pedestal_domain_mask,
+    )
 
     # Combine the transport coefficients from core and pedestal models.
-    def _combine_maybe_none_coeffs(*leaves):
-      non_none_leaves = [leaf for leaf in leaves if leaf is not None]
-      return sum(non_none_leaves) if non_none_leaves else None
-
     combined_transport_coeffs = jax.tree.map(
-        _combine_maybe_none_coeffs,
-        *pedestal_coeffs,
-        *core_coeffs,
-        # Needed to handle the case where some coefficients are None and others
-        # are not.
-        is_leaf=lambda x: x is None,
+        _add_optional, core_coeffs, pedestal_coeffs
     )
 
     return combined_transport_coeffs
 
-  def _apply_pedestal_domain_restriction(
+  def _combine(
       self,
-      unused_transport_runtime_params: transport_runtime_params_lib.RuntimeParams,
+      models: tuple[transport_model_lib.TransportModel, ...],
+      params_list: Sequence[transport_runtime_params_lib.RuntimeParams],
+      runtime_params: runtime_params_lib.RuntimeParams,
       geo: geometry.Geometry,
-      transport_coeffs: transport_model_lib.TurbulentTransport,
+      core_profiles: state.CoreProfiles,
       pedestal_model_output: pedestal_model_lib.PedestalModelOutput,
+      domain_mask_fn: Callable[
+          [
+              transport_runtime_params_lib.RuntimeParams,
+              geometry.Geometry,
+              pedestal_model_lib.PedestalModelOutput,
+          ],
+          jax.Array,
+      ],
   ) -> transport_model_lib.TurbulentTransport:
-    del unused_transport_runtime_params
-    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)
+    """Calculates and combines transport coefficients from a list of models."""
 
-    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,
+    # Initialize accumulators with zeros. Will be iteratively updated based on
+    # model outputs and merge modes.
+    zero_profile = jnp.zeros_like(
+        geo.rho_face_norm, dtype=jax_utils.get_dtype()
     )
+    accumulators = {}
+    locks = {}
+
+    for channel, config in transport_model_lib.CHANNEL_CONFIG_STRUCT.items():
+      accumulators[channel] = zero_profile
+      locks[channel] = jnp.zeros_like(geo.rho_face_norm, dtype=bool)
+      for sub in config['sub_channels']:
+        accumulators[sub] = None
+
+    # TODO(b/344023668) explore batching or fori_loop for performance.
+    for model, params in zip(models, params_list, strict=True):
+      # 1. Calculate raw coefficients
+      coeffs = model.call_implementation(
+          params, runtime_params, geo, core_profiles, pedestal_model_output
+      )
+
+      # 2. Zero out disabled channels. Unused subchannels returned as None.
+      coeffs = model.zero_out_disabled_channels(params, coeffs)
+
+      # 3. Calculate active domain mask. Values outside this are set to 0.
+      domain_mask = domain_mask_fn(params, geo, pedestal_model_output)
+
+      coeffs_dict = dataclasses.asdict(coeffs)
+      for k in coeffs_dict:
+        # Apply domain restriction to values.
+        if coeffs_dict[k] is not None:
+          coeffs_dict[k] = jnp.where(domain_mask, coeffs_dict[k], 0.0)
+
+      for channel, config in transport_model_lib.CHANNEL_CONFIG_STRUCT.items():
+        disable_flag_name = config['disable_flag']
+        is_disabled = getattr(params, disable_flag_name)
+
+        # A channel is active for this model if it's in the domain AND enabled.
+        # Note that this is a boolean array over the face grid.
+        channel_active = jnp.logical_and(
+            domain_mask, jnp.logical_not(is_disabled)
+        )
+
+        val = coeffs_dict[channel]
+        if params.merge_mode == enums.MergeMode.OVERWRITE:
+          # Wiping: Replace accumulator values where active.
+          accumulators[channel] = jnp.where(
+              channel_active, val, accumulators[channel]
+          )
+          # Update lock.
+          locks[channel] = jnp.logical_or(locks[channel], channel_active)
+        else:  # ADD
+          # Add where not locked.
+          factor = jnp.where(locks[channel], 0.0, 1.0)
+          accumulators[channel] = accumulators[channel] + val * factor
+
+        # Handle sub-channels.
+        for sub in config['sub_channels']:
+          sub_val = coeffs_dict[sub]
+          if sub_val is not None:
+            if accumulators[sub] is None:
+              accumulators[sub] = zero_profile
+
+            if params.merge_mode == enums.MergeMode.OVERWRITE:
+              accumulators[sub] = jnp.where(
+                  channel_active, sub_val, accumulators[sub]
+              )
+            else:  # ADD
+              # Add where not locked (using main channel lock).
+              factor = jnp.where(locks[channel], 0.0, 1.0)
+              accumulators[sub] = accumulators[sub] + sub_val * factor
+
+    return transport_model_lib.TurbulentTransport(**accumulators)
+
+
+def _add_optional(
+    core_value: jax.Array | None, pedestal_value: jax.Array | None
+) -> jax.Array | None:
+  """Adds two values, treating None as zero. Returns None if both are None."""
+  if core_value is None:
+    return pedestal_value
+  if pedestal_value is None:
+    return core_value
+  return core_value + pedestal_value
+
+
+def _pedestal_domain_mask(
+    unused_params: transport_runtime_params_lib.RuntimeParams,
+    geo: geometry.Geometry,
+    pedestal_output: pedestal_model_lib.PedestalModelOutput,
+) -> jax.Array:
+  """Calculates the active domain mask for pedestal transport models."""
+  return jnp.asarray(geo.rho_face_norm > pedestal_output.rho_norm_ped_top)

torax/_src/transport_model/constant.py

@@ -50,7 +50,7 @@ class RuntimeParams(transport_runtime_params_lib.RuntimeParams):
 class ConstantTransportModel(transport_model_lib.TransportModel):
   """Calculates various coefficients related to particle transport."""
 
-  def _call_implementation(
+  def call_implementation(
       self,
       transport_runtime_params: transport_runtime_params_lib.RuntimeParams,
       runtime_params: runtime_params_lib.RuntimeParams,

torax/_src/transport_model/critical_gradient.py

@@ -44,7 +44,7 @@ class RuntimeParams(transport_runtime_params_lib.RuntimeParams):
 class CriticalGradientTransportModel(transport_model.TransportModel):
   """Calculates various coefficients related to particle transport."""
 
-  def _call_implementation(
+  def call_implementation(
       self,
       transport_runtime_params: transport_runtime_params_lib.RuntimeParams,
       runtime_params: runtime_params_lib.RuntimeParams,

torax/_src/transport_model/qlknn_transport_model.py

@@ -217,7 +217,7 @@ class QLKNNTransportModel(
   path: str
   name: str
 
-  def _call_implementation(
+  def call_implementation(
       self,
       transport_runtime_params: RuntimeParams,
       runtime_params: runtime_params_lib.RuntimeParams,

torax/_src/transport_model/qualikiz_transport_model.py

@@ -81,7 +81,7 @@ def __init__(self):
     )
     self._runpath = os.path.join(self._qlkrun_parentdir.name, self._qlkrun_name)
 
-  def _call_implementation(
+  def call_implementation(
       self,
       transport_runtime_params: transport_runtime_params_lib.RuntimeParams,
       runtime_params: runtime_params_lib.RuntimeParams,

torax/_src/transport_model/tests/bohm_gyrobohm_test.py

@@ -134,14 +134,14 @@ def test_coeff_multiplier_feature(self):
         chi_i_gyrobohm_multiplier=5.0,
     )
 
-    output_A = self.model._call_implementation(
+    output_A = self.model.call_implementation(
         dyn_params_A.transport,
         dyn_params_A,
         self.geo,
         self.core_profiles,
         self.pedestal_outputs,
     )
-    output_B = self.model._call_implementation(
+    output_B = self.model.call_implementation(
         dyn_params_B.transport,
         dyn_params_B,
         self.geo,
@@ -179,14 +179,14 @@ def test_raw_bohm_and_gyrobohm_fields(self):
         chi_i_gyrobohm_multiplier=5.0,
     )
 
-    output_A = self.model._call_implementation(
+    output_A = self.model.call_implementation(
         dyn_params_A.transport,
         dyn_params_A,
         self.geo,
         self.core_profiles,
         self.pedestal_outputs,
     )
-    output_B = self.model._call_implementation(
+    output_B = self.model.call_implementation(
         dyn_params_B.transport,
         dyn_params_B,
         self.geo,