Add predictive L-H transition based on the Martin scaling.

* Add from_pedestal_model saturation model, which increases transport in the pedestal if the requested pedestal top values are exceeded, reducing the pedestal height.
* Add Martin scaling formation model, which reduces the transport in the pedestal if the LH threshold is crossed, increasing the pedestal height.
* Introduce ADAPTIVE_SOURCE and ADAPTIVE_TRANSPORT modes. If in ADAPTIVE transport mode, the formation and saturation models are used.
* Add integration test: an ITER-like case with an L-H transition.

PiperOrigin-RevId: 874140482

Author: theo-brown

torax/_src/pedestal_model/formation/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2026 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

torax/_src/pedestal_model/formation/base.py

@@ -0,0 +1,50 @@
+# Copyright 2026 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Base class for pedestal formation models."""
+
+import abc
+import dataclasses
+from torax._src import state
+from torax._src import static_dataclass
+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
+from torax._src.sources import source_profiles as source_profiles_lib
+
+
+@dataclasses.dataclass(frozen=True, eq=False)
+class FormationModel(static_dataclass.StaticDataclass, abc.ABC):
+  """Base class for pedestal formation models."""
+
+  @abc.abstractmethod
+  def __call__(
+      self,
+      runtime_params: runtime_params_lib.RuntimeParams,
+      geo: geometry.Geometry,
+      core_profiles: state.CoreProfiles,
+      source_profiles: source_profiles_lib.SourceProfiles,
+  ) -> pedestal_model_output.TransportMultipliers:
+    """Calculates the transport decrease multipliers.
+
+    Args:
+      runtime_params: Runtime parameters.
+      geo: Geometry.
+      core_profiles: Core profiles.
+      source_profiles: Source profiles.
+
+    Returns:
+      transport_decrease_multiplier: Factors to multiply transport coefficients
+        by (<= 1.0).
+    """

torax/_src/pedestal_model/formation/martin_formation_model.py

@@ -0,0 +1,106 @@
+# Copyright 2026 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Martin scaling pedestal formation model."""
+
+import dataclasses
+import jax
+import jax.numpy as jnp
+from torax._src import array_typing
+from torax._src import math_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_output
+from torax._src.pedestal_model import runtime_params as pedestal_runtime_params_lib
+from torax._src.pedestal_model.formation import base
+from torax._src.physics import scaling_laws
+from torax._src.sources import source_profiles as source_profiles_lib
+
+# pylint: disable=invalid-name
+
+
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True)
+class MartinFormationRuntimeParams(
+    pedestal_runtime_params_lib.FormationRuntimeParams
+):
+  """Runtime params for pedestal formation models."""
+
+  P_LH_prefactor: array_typing.FloatScalar = 1.0
+
+
+def _calculate_P_SOL_total(
+    core_sources: source_profiles_lib.SourceProfiles,
+    geo: geometry.Geometry,
+) -> jax.Array:
+  """Calculates the total power out of the separatrix."""
+  P_SOL_e = sum(
+      math_utils.volume_integration(source, geo)
+      for source in core_sources.T_e.values()
+  )
+  P_SOL_i = sum(
+      math_utils.volume_integration(source, geo)
+      for source in core_sources.T_i.values()
+  )
+  # TODO(b/488318267): Missing dW/dt terms.
+  return P_SOL_e + P_SOL_i
+
+
+@dataclasses.dataclass(frozen=True, eq=False)
+class MartinFormationModel(base.FormationModel):
+  """Pedestal formation based on P_SOL and P_LH, using Martin scaling."""
+
+  def __call__(
+      self,
+      runtime_params: runtime_params_lib.RuntimeParams,
+      geo: geometry.Geometry,
+      core_profiles: state.CoreProfiles,
+      core_sources: source_profiles_lib.SourceProfiles,
+  ) -> pedestal_model_output.TransportMultipliers:
+    """Calculates the transport decrease multipliers using Martin scaling."""
+    assert isinstance(
+        runtime_params.pedestal.formation, MartinFormationRuntimeParams
+    )
+
+    P_SOL_total = _calculate_P_SOL_total(core_sources, geo)
+    _, _, P_LH, _ = scaling_laws.calculate_plh_scaling_factor(
+        geo, core_profiles
+    )
+    rescaled_P_LH = P_LH * runtime_params.pedestal.formation.P_LH_prefactor
+
+    # Calculate transport_multiplier
+    # If P_SOL > P_LH, multiplier tends to 0.0
+    # If P_SOL < P_LH, multiplier tends to 1.0
+    # TODO(b/323504363): Add hysteresis logic here later
+    width = runtime_params.pedestal.formation.sigmoid_width
+    exponent = runtime_params.pedestal.formation.sigmoid_exponent
+    offset = runtime_params.pedestal.formation.sigmoid_offset
+    normalized_deviation = (
+        P_SOL_total - rescaled_P_LH
+    ) / rescaled_P_LH - offset
+    transport_multiplier = 1 - jax.nn.sigmoid(normalized_deviation / width)
+    transport_multiplier = transport_multiplier**exponent
+    transport_multiplier = jnp.clip(
+        transport_multiplier,
+        min=runtime_params.pedestal.min_transport_multiplier,
+        max=runtime_params.pedestal.max_transport_multiplier,
+    )
+
+    return pedestal_model_output.TransportMultipliers(
+        chi_e_multiplier=transport_multiplier,
+        chi_i_multiplier=transport_multiplier,
+        D_e_multiplier=transport_multiplier,
+        v_e_multiplier=transport_multiplier,
+    )

torax/_src/pedestal_model/formation/tests/martin_formation_model_test.py

@@ -0,0 +1,113 @@
+# Copyright 2026 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import dataclasses
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax.numpy as jnp
+import numpy as np
+from torax._src.orchestration import initial_state
+from torax._src.orchestration import run_simulation
+from torax._src.pedestal_model.formation import martin_formation_model
+from torax._src.test_utils import default_configs
+from torax._src.torax_pydantic import model_config
+
+# pylint: disable=invalid-name
+
+
+class MartinFormationModelTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    config = default_configs.get_default_config_dict()
+    # Switch to use the Martin formation model.
+    config['pedestal'] = {
+        'set_pedestal': True,
+        'mode': 'ADAPTIVE_TRANSPORT',
+        'formation_model': {'model_name': 'martin'},
+    }
+    # Add a source so that P_SOL is non-zero.
+    config['sources'] = {
+        'generic_heat': {
+            'gaussian_location': 0.15,
+            'gaussian_width': 0.1,
+            'P_total': 20.0e6,
+            'electron_heat_fraction': 0.8,
+        }
+    }
+    self.torax_config = model_config.ToraxConfig.from_dict(config)
+    step_fn = run_simulation.make_step_fn(self.torax_config)
+    self.initial_state, self.initial_post_processed_outputs = (
+        initial_state.get_initial_state_and_post_processed_outputs(step_fn)
+    )
+    self.runtime_params = step_fn.runtime_params_provider(t=0.0)
+
+  def test_calculate_P_SOL_total(self):
+    P_SOL_total = martin_formation_model._calculate_P_SOL_total(
+        self.initial_state.core_sources, self.initial_state.geometry
+    )
+
+    np.testing.assert_allclose(
+        P_SOL_total, self.initial_post_processed_outputs.P_SOL_total
+    )
+
+  @parameterized.named_parameters(
+      dict(
+          # If P_sol >> P_LH, we expect the suppression multiplier to be very
+          # small (significant suppression). However, it's clipped internally to
+          # be 0.1.
+          testcase_name='above_threshold',
+          power=1e6,
+          expected_multiplier=0.1,
+      ),
+      dict(
+          # If P_sol << P_LH, we expect the suppression multiplier to be 1.0
+          # (no suppression).
+          testcase_name='below_threshold',
+          # We set the aux power to be negative (equivalent to a heat sink)
+          # to make sure that P_sol < P_LH.
+          power=-1e6,
+          expected_multiplier=1.0,
+      ),
+  )
+  def test_martin_formation_model_suppression(self, power, expected_multiplier):
+    formation_model = martin_formation_model.MartinFormationModel()
+
+    aux_power_profile = power * jnp.ones_like(self.initial_state.geometry.rho)
+    high_power_profiles = dataclasses.replace(
+        self.initial_state.core_sources,
+        T_e={'aux': aux_power_profile},
+        T_i={'aux': aux_power_profile},
+    )
+
+    transport_multipliers = formation_model(
+        self.runtime_params,
+        self.initial_state.geometry,
+        self.initial_state.core_profiles,
+        high_power_profiles,
+    )
+    for k, multiplier in dataclasses.asdict(transport_multipliers).items():
+      np.testing.assert_allclose(
+          multiplier,
+          expected_multiplier,
+          atol=1e-3,
+          err_msg=(
+              f'{k}={multiplier} is not close to the expected value of'
+              f' {expected_multiplier}.'
+          ),
+      )
+
+
+if __name__ == '__main__':
+  absltest.main()

torax/_src/pedestal_model/pedestal_model.py

@@ -27,6 +27,9 @@
 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
+from torax._src.pedestal_model import runtime_params as pedestal_runtime_params_lib
+from torax._src.pedestal_model.formation import base as formation_base
+from torax._src.pedestal_model.saturation import base as saturation_base
 from torax._src.sources import source_profiles as source_profiles_lib
 
 # pylint: disable=invalid-name
@@ -37,16 +40,48 @@
 class PedestalModel(static_dataclass.StaticDataclass, abc.ABC):
   """Calculates temperature and density of the pedestal."""
 
+  formation_model: formation_base.FormationModel
+  saturation_model: saturation_base.SaturationModel
+
   def __call__(
       self,
       runtime_params: runtime_params_lib.RuntimeParams,
       geo: geometry.Geometry,
       core_profiles: state.CoreProfiles,
       source_profiles: source_profiles_lib.SourceProfiles,
   ) -> pedestal_model_output.PedestalModelOutput:
+    pedestal_output = self._call_implementation(
+        runtime_params, geo, core_profiles
+    )
+
+    # If in ADAPTIVE_TRANSPORT mode, calculate the transport multipliers based
+    # on the formation and saturation models.
+    if (
+        runtime_params.pedestal.mode
+        == pedestal_runtime_params_lib.Mode.ADAPTIVE_TRANSPORT
+    ):
+      transport_decrease_multipliers = self.formation_model(
+          runtime_params, geo, core_profiles, source_profiles
+      )
+      transport_increase_multipliers = self.saturation_model(
+          runtime_params, geo, core_profiles, pedestal_output
+      )
+
+      # Combine via exp(log) for numerical stability, as multipliers can
+      # be very small or large.
+      transport_multipliers = jax.tree.map(
+          lambda x, y: jnp.exp(jnp.log(x) + jnp.log(y)),
+          transport_decrease_multipliers,
+          transport_increase_multipliers,
+      )
+
+      pedestal_output = dataclasses.replace(
+          pedestal_output, transport_multipliers=transport_multipliers
+      )
+
     return jax.lax.cond(
         runtime_params.pedestal.set_pedestal,
-        lambda: self._call_implementation(runtime_params, geo, core_profiles),
+        lambda: pedestal_output,
         lambda: pedestal_model_output.PedestalModelOutput(
             rho_norm_ped_top=jnp.inf,
             rho_norm_ped_top_idx=geo.torax_mesh.nx,