Create NoSaturationModel and NoFormationModel

Useful for testing and debugging

PiperOrigin-RevId: 877364049

Author: theo-brown

torax/_src/pedestal_model/formation/no_formation_model.py

@@ -0,0 +1,42 @@
+# 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.
+
+"""No pedestal formation model."""
+
+import dataclasses
+import jax
+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.formation import base
+from torax._src.sources import source_profiles as source_profiles_lib
+
+# pylint: disable=invalid-name
+
+
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True, eq=False)
+class NoFormationModel(base.FormationModel):
+  """No pedestal formation model."""
+
+  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:
+    """Returns no transport multipliers."""
+    return pedestal_model_output.TransportMultipliers.default()

torax/_src/pedestal_model/pydantic_model.py

@@ -26,12 +26,37 @@
 from torax._src.pedestal_model import set_pped_tpedratio_nped
 from torax._src.pedestal_model import set_tped_nped
 from torax._src.pedestal_model.formation import martin_formation_model
+from torax._src.pedestal_model.formation import no_formation_model
+from torax._src.pedestal_model.saturation import no_saturation_model
 from torax._src.pedestal_model.saturation import profile_value_saturation_model
 from torax._src.torax_pydantic import torax_pydantic
 
 # pylint: disable=invalid-name
 
 
+class NoFormation(torax_pydantic.BaseModelFrozen):
+  """Configuration for No formation model."""
+
+  model_name: Annotated[Literal["no_formation"], torax_pydantic.JAX_STATIC] = (
+      "no_formation"
+  )
+
+  def build_formation_model(
+      self,
+  ) -> no_formation_model.NoFormationModel:
+    return no_formation_model.NoFormationModel()
+
+  def build_runtime_params(
+      self, t: chex.Numeric
+  ) -> runtime_params.FormationRuntimeParams:
+    del t
+    return runtime_params.FormationRuntimeParams(
+        sigmoid_width=0.0,
+        sigmoid_offset=0.0,
+        sigmoid_exponent=1.0,
+    )
+
+
 # TODO(b/323504363): Generalise to pedestal formation models based on power
 # thresholds (e.g. Metal Wall scaling), not just Martin scaling.
 class MartinFormation(torax_pydantic.BaseModelFrozen):
@@ -87,6 +112,29 @@ def build_runtime_params(
     )
 
 
+class NoSaturation(torax_pydantic.BaseModelFrozen):
+  """Configuration for No saturation model."""
+
+  model_name: Annotated[Literal["no_saturation"], torax_pydantic.JAX_STATIC] = (
+      "no_saturation"
+  )
+
+  def build_saturation_model(
+      self,
+  ) -> no_saturation_model.NoSaturationModel:
+    return no_saturation_model.NoSaturationModel()
+
+  def build_runtime_params(
+      self, t: chex.Numeric
+  ) -> runtime_params.SaturationRuntimeParams:
+    del t
+    return runtime_params.SaturationRuntimeParams(
+        sigmoid_width=0.0,
+        sigmoid_offset=0.0,
+        sigmoid_exponent=1.0,
+    )
+
+
 class ProfileValueSaturation(torax_pydantic.BaseModelFrozen):
   """Configuration for ProfileValueSaturation model.
 
@@ -140,8 +188,8 @@ def build_runtime_params(
 
 
 # For new formation and saturation models, add to these TypeAliases via Union.
-FormationConfig: TypeAlias = MartinFormation
-SaturationConfig: TypeAlias = ProfileValueSaturation
+FormationConfig: TypeAlias = MartinFormation | NoFormation
+SaturationConfig: TypeAlias = ProfileValueSaturation | NoSaturation
 
 
 class BasePedestal(torax_pydantic.BaseModelFrozen, abc.ABC):

torax/_src/pedestal_model/saturation/no_saturation_model.py

@@ -0,0 +1,40 @@
+# 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.
+
+"""No pedestal saturation model."""
+
+import dataclasses
+from torax._src import array_typing
+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.saturation import base
+
+# pylint: disable=invalid-name
+
+
+@dataclasses.dataclass(frozen=True, eq=False)
+class NoSaturationModel(base.SaturationModel):
+  """No pedestal saturation model."""
+
+  def __call__(
+      self,
+      runtime_params: runtime_params_lib.RuntimeParams,
+      geo: geometry.Geometry,
+      core_profiles: state.CoreProfiles,
+      pedestal_output: pedestal_model_output.PedestalModelOutput,
+  ) -> array_typing.FloatScalar:
+    """Returns no transport multipliers."""
+    return pedestal_model_output.TransportMultipliers.default()