Flatten __call__ and simplify _combine in CombinedTransportModel.

- Merge call_implementation() body into __call__(), removing the separate method. __call__ now directly combines core and pedestal coefficients, applies clipping, and applies smoothing.
- Simplify _combine loop: call model() directly instead of separate model.call_implementation() + model.zero_out_disabled_channels(), since ComponentTransportModel.__call__ now wraps both.
- Update combined_test.py to call model() instead of model.call_implementation().

PiperOrigin-RevId: 964556808

Author: tamaranorman

torax/_src/mhd/tests/mhd_pydantic_model_test.py

@@ -28,14 +28,14 @@
 from torax._src.sources import source_models as source_models_lib
 from torax._src.test_utils import default_configs
 from torax._src.torax_pydantic import model_config
-from torax._src.transport_model import transport_model as transport_model_lib
+from torax._src.transport_model import component
 
 
 class MHDPydanticModelTest(parameterized.TestCase):
 
   def setUp(self):
     super().setUp()
-    self.transport_model = mock.Mock(spec=transport_model_lib.TransportModel)
+    self.transport_model = mock.Mock(spec=component.ComponentTransportModel)
     self.source_models = mock.Mock(spec=source_models_lib.SourceModels)
     self.pedestal_model = mock.Mock(spec=pedestal_model_lib.PedestalModel)
     self.neoclassical_models = mock.Mock(
@@ -48,10 +48,8 @@ def test_no_mhd_config_makes_empty_runtime_params(self):
     )
 
     self.assertIsInstance(torax_config.mhd, mhd_pydantic_model.MHD)
-    provider = (
-        build_runtime_params.RuntimeParamsProvider.from_config(
-            torax_config
-        )
+    provider = build_runtime_params.RuntimeParamsProvider.from_config(
+        torax_config
     )
     runtime_params = provider(t=0.0)
     self.assertIsInstance(runtime_params.mhd, mhd_runtime_params.RuntimeParams)
@@ -67,10 +65,8 @@ def test_empty_mhd_config(self):
     assert isinstance(torax_config.mhd, mhd_pydantic_model.MHD)
     mhd_models = torax_config.mhd.build_mhd_models()
     self.assertIs(mhd_models.sawtooth_models, None)
-    provider = (
-        build_runtime_params.RuntimeParamsProvider.from_config(
-            torax_config
-        )
+    provider = build_runtime_params.RuntimeParamsProvider.from_config(
+        torax_config
     )
     runtime_params = provider(t=0.0)
     self.assertIsInstance(runtime_params.mhd, mhd_runtime_params.RuntimeParams)
@@ -96,10 +92,8 @@ def test_mhd_config_with_sawtooth(self):
         torax_config.mhd.sawtooth, sawtooth_pydantic_model.SawtoothConfig
     )
 
-    provider = (
-        build_runtime_params.RuntimeParamsProvider.from_config(
-            torax_config
-        )
+    provider = build_runtime_params.RuntimeParamsProvider.from_config(
+        torax_config
     )
     runtime_params = provider(t=0.0)
     sawtooth_params = runtime_params.mhd.sawtooth

torax/_src/transport_model/bohm_gyrobohm.py

@@ -24,8 +24,8 @@
 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_output as pedestal_model_output_lib
+from torax._src.transport_model import component
 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
 
 # pylint: disable=invalid-name
 
@@ -49,7 +49,7 @@ class RuntimeParams(transport_runtime_params_lib.RuntimeParams):
 
 
 @dataclasses.dataclass(kw_only=True, frozen=True, eq=False)
-class BohmGyroBohmTransportModel(transport_model_lib.TransportModel):
+class BohmGyroBohmTransportModel(component.ComponentTransportModel):
   """Calculates various coefficients related to particle transport according to the Bohm + gyro-Bohm Model."""
 
   def call_implementation(
@@ -59,7 +59,7 @@ def call_implementation(
       geo: geometry.Geometry,
       core_profiles: state.CoreProfiles,
       pedestal_model_output: pedestal_model_output_lib.PedestalModelOutput,
-  ) -> transport_model_lib.TurbulentTransport:
+  ) -> component.TurbulentTransport:
     r"""Calculates transport coefficients using the BohmGyroBohm model.
 
     We use the implementation from Tholerus et al, Section 3.3.
@@ -171,7 +171,7 @@ def call_implementation(
     # Electron convectivity set proportional to the electron diffusivity
     v_face_el = transport_runtime_params.V_face_coeff * d_face_el
 
-    return transport_model_lib.TurbulentTransport(
+    return component.TurbulentTransport(
         chi_face_ion=chi_i,  # pyrefly: ignore[bad-argument-type]
         chi_face_el=chi_e,  # pyrefly: ignore[bad-argument-type]
         d_face_el=d_face_el,

torax/_src/transport_model/combined.py

@@ -50,47 +50,12 @@ def __call__(
       core_profiles: state.CoreProfiles,
       pedestal_model_output: pedestal_model_output_lib.PedestalModelOutput,
   ) -> transport_model_lib.TurbulentTransport:
-
-    transport_runtime_params = runtime_params.transport
-
-    # Calculate the transport coefficients - includes contribution from pedestal
-    # and core transport models.
-    transport_coeffs = self.call_implementation(
-        transport_runtime_params,
-        runtime_params,
-        geo,
-        core_profiles,
-        pedestal_model_output,
-    )
-    # Apply min/max clipping
-    transport_coeffs = self._apply_clipping(
-        transport_runtime_params,
-        transport_coeffs,
-    )
-
-    transport_coeffs = self._smooth_coeffs(
-        runtime_params,
-        geo,
-        transport_coeffs,
-        pedestal_model_output,
-    )
-
-    return transport_coeffs
-
-  def call_implementation(
-      self,
-      transport_runtime_params: transport_runtime_params_lib.CombinedRuntimeParams,
-      runtime_params: runtime_params_lib.RuntimeParams,
-      geo: geometry.Geometry,
-      core_profiles: state.CoreProfiles,
-      pedestal_model_output: pedestal_model_output_lib.PedestalModelOutput,
-  ) -> transport_model_lib.TurbulentTransport:
     r"""Calculates transport coefficients using the Combined model.
 
+    Combines coefficients from core and pedestal transport models, applies
+    min/max clipping, and smooths the result.
+
     Args:
-      transport_runtime_params: Input runtime parameters for this transport
-        model (expected to be an instance of CombinedRuntimeParams at runtime).
-        Can change without triggering a JAX recompilation.
       runtime_params: Runtime parameters for the simulation at the current time.
       geo: Geometry of the torus at the current time.
       core_profiles: Core plasma profiles.
@@ -99,16 +64,20 @@ def call_implementation(
     Returns:
       coeffs: The transport coefficients
     """
+    transport_runtime_params = runtime_params.transport
+
+    # Calculate transport coefficients from core models.
     core_coeffs = self._combine(
         self.transport_models,
-        transport_runtime_params.transport_model_params,
+        transport_runtime_params.core_transport_model_params,
         runtime_params,
         geo,
         core_profiles,
         pedestal_model_output,
         transport_model_lib.compute_core_domain_mask,
     )
 
+    # Calculate transport coefficients from pedestal models.
     pedestal_coeffs = self._combine(
         self.pedestal_transport_models,
         transport_runtime_params.pedestal_transport_model_params,
@@ -120,11 +89,25 @@ def call_implementation(
     )
 
     # Combine the transport coefficients from core and pedestal models.
-    combined_transport_coeffs = jax.tree.map(
+    transport_coeffs = jax.tree.map(
         _add_optional, core_coeffs, pedestal_coeffs
     )
 
-    return combined_transport_coeffs
+    # Apply min/max clipping.
+    transport_coeffs = self._apply_clipping(
+        transport_runtime_params,
+        transport_coeffs,
+    )
+
+    # Apply smoothing.
+    transport_coeffs = self._smooth_coeffs(
+        runtime_params,
+        geo,
+        transport_coeffs,
+        pedestal_model_output,
+    )
+
+    return transport_coeffs
 
   def _combine(
       self,
@@ -162,15 +145,12 @@ def _combine(
 
     # 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(
+      # 1. Calculate raw coefficients and zero out disabled channels.
+      coeffs = model(
           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.
+      # 2. Calculate active domain mask. Values outside this are set to 0.
       domain_mask = domain_mask_fn(
           params, runtime_params, geo, pedestal_model_output
       )
@@ -357,8 +337,7 @@ def build_profile_fallback():
 
     def apply_pedestal_mask(profile):
       return jnp.where(
-          geo.rho_face_norm
-          < pedestal_model_output.rho_norm_ped_top,
+          geo.rho_face_norm < pedestal_model_output.rho_norm_ped_top,
           profile,
           0.0,
       )