Bugfix: apply domain restrictions to optional transport sub-channels.

Drive-by: avoid explicitly writing out sub-channels elsewhere and use CHANNEL_CONFIG_STRUCT for iterations. Will help with subsequent additions to optional transport coefficients.

bohmgyrobohm test updated due to the bohm and gyrobohm transport coefficient outputs being different in the pedestal region due to the now correct domain restriction. No impact on actual results (since these coefficients were not actually used directly)

PiperOrigin-RevId: 860323344

Author: jcitrin

torax/_src/transport_model/tests/transport_model_test.py

@@ -596,6 +596,53 @@ def test_preserves_none_disabled(self):
     )
     self.assertIsNone(new_coeffs_disabled.chi_face_ion_bohm)
 
+  def test_sub_channel_domain_restriction(self):
+    """Tests that sub-channels are masked by domain restriction."""
+    config = default_configs.get_default_config_dict()
+    config['transport'] = {
+        'model_name': 'fixed',
+        'rho_max': 0.8,
+        'smoothing_width': 0.0,
+        'chi_min': 0.0,
+        'D_e_min': 0.0,
+        'V_e_min': 0.0,
+    }
+    torax_config = model_config.ToraxConfig.from_dict(config)
+    runtime_params = build_runtime_params.RuntimeParamsProvider.from_config(
+        torax_config
+    )(t=0.0)
+    geo = torax_config.geometry.build_provider(t=0.0)
+    # We need a pedestal model even if unused by the fixed transport
+    pedestal_model = torax_config.pedestal.build_pedestal_model()
+    # Mock core profiles (not used by FixedTransportModel but needed for API)
+    core_profiles = initialization.initial_core_profiles(
+        runtime_params,
+        geo,
+        torax_config.sources.build_models(),
+        torax_config.neoclassical.build_models(),
+    )
+    pedestal_outputs = pedestal_model(runtime_params, geo, core_profiles)
+
+    transport_model = torax_config.transport.build_transport_model()
+    coeffs = transport_model(
+        runtime_params, geo, core_profiles, pedestal_outputs
+    )
+
+    # Find index where rho > 0.8
+    cutoff_idx = np.searchsorted(geo.rho_face_norm, 0.8, side='right')
+
+    # Verify main channel is zeroed
+    np.testing.assert_allclose(coeffs.chi_face_ion[cutoff_idx:], 0.0)
+
+    # Verify sub-channels are also zeroed
+    # FixedTransportModel sets chi_face_ion_bohm = chi_face_ion * 0.3
+    # If not masked, it would be non-zero because FixedTransportModel computes
+    # it everywhere
+    self.assertIsNotNone(coeffs.chi_face_ion_bohm)
+    np.testing.assert_allclose(coeffs.chi_face_ion_bohm[cutoff_idx:], 0.0)
+    self.assertIsNotNone(coeffs.chi_face_ion_gyrobohm)
+    np.testing.assert_allclose(coeffs.chi_face_ion_gyrobohm[cutoff_idx:], 0.0)
+
 
 if __name__ == '__main__':
   absltest.main()

torax/_src/transport_model/transport_model.py

@@ -193,18 +193,21 @@ def _apply_domain_restriction(
         transport_runtime_params, geo, pedestal_model_output
     )
 
-    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)
+    coeffs_dict = dataclasses.asdict(transport_coeffs)
+    to_replace = {}
 
-    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,
-    )
+    for channel_name, config in CHANNEL_CONFIG_STRUCT.items():
+      # Mask main channel
+      val = coeffs_dict[channel_name]
+      to_replace[channel_name] = jnp.where(active_mask, val, 0.0)
+
+      # Mask sub-channels
+      for sub_channel in config['sub_channels']:
+        sub_val = coeffs_dict[sub_channel]
+        if sub_val is not None:
+          to_replace[sub_channel] = jnp.where(active_mask, sub_val, 0.0)
+
+    return dataclasses.replace(transport_coeffs, **to_replace)
 
   def _apply_clipping(
       self,