Add a gas puff source with feedback on line averaged density.

PiperOrigin-RevId: 906928786

Author: theo-brown

torax/_src/sources/gas_puff_source.py

@@ -18,7 +18,9 @@
 
 import chex
 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
@@ -42,6 +44,8 @@
 class RuntimeParams(sources_runtime_params_lib.RuntimeParams):
   puff_decay_length: array_typing.FloatScalar
   S_total: array_typing.FloatScalar
+  target_line_average_n_e: array_typing.FloatScalar
+  feedback_gain: array_typing.FloatScalar
 
 
 # Default formula: exponential
@@ -66,6 +70,35 @@ def calc_puff_source(
   )
 
 
+# Gas puff with feedback on line averaged density
+def calc_puff_feedback_source(
+    runtime_params: runtime_params_lib.RuntimeParams,
+    geo: geometry.Geometry,
+    source_name: str,
+    core_profiles: state.CoreProfiles,
+    unused_calculated_source_profiles: source_profiles.SourceProfiles | None,
+    unused_conductivity: conductivity_base.Conductivity | None,
+) -> tuple[array_typing.FloatVectorCell, ...]:
+  """Calculates external source term for n from puffs with feedback."""
+  source_params = runtime_params.sources[source_name]
+  assert isinstance(source_params, RuntimeParams)
+
+  current_line_avg_n_e = math_utils.line_average(core_profiles.n_e.value, geo)
+  error = source_params.target_line_average_n_e - current_line_avg_n_e
+
+  S_total = source_params.feedback_gain * error
+  S_total = jnp.clip(S_total, 0.0, jnp.inf)
+
+  return (
+      formulas.exponential_profile(
+          decay_start=1.0,
+          width=source_params.puff_decay_length,
+          total=S_total,
+          geo=geo,
+      ),
+  )
+
+
 @dataclasses.dataclass(kw_only=True, frozen=True, eq=False)
 class GasPuffSource(source.Source):
   """Gas puff source for the n_e equation."""
@@ -86,35 +119,50 @@ class GasPuffSourceConfig(base.SourceModelBase):
     S_total: total gas puff particles/s
   """
 
-  model_name: Annotated[Literal['exponential'], torax_pydantic.JAX_STATIC] = (
-      'exponential'
-  )
+  model_name: Annotated[
+      Literal['exponential', 'feedback'], torax_pydantic.JAX_STATIC
+  ] = 'exponential'
   puff_decay_length: torax_pydantic.TimeVaryingScalar = (
       torax_pydantic.ValidatedDefault(0.05)
   )
   S_total: torax_pydantic.TimeVaryingScalar = torax_pydantic.ValidatedDefault(
       1e22
   )
+  target_line_average_n_e: torax_pydantic.TimeVaryingScalar = (
+      torax_pydantic.ValidatedDefault(0.0)
+  )
+  feedback_gain: torax_pydantic.TimeVaryingScalar = (
+      torax_pydantic.ValidatedDefault(0.0)
+  )
   mode: Annotated[
       sources_runtime_params_lib.Mode, torax_pydantic.JAX_STATIC
   ] = sources_runtime_params_lib.Mode.MODEL_BASED
 
   @property
   def model_func(self) -> source.SourceProfileFunction:
+    if self.model_name == 'feedback':
+      return calc_puff_feedback_source
     return calc_puff_source
 
   def build_runtime_params(
       self,
       t: chex.Numeric,
   ) -> RuntimeParams:
+    if self.model_name == 'feedback':
+      is_explicit = True
+    else:
+      is_explicit = self.is_explicit
+
     return RuntimeParams(
         prescribed_values=tuple(
             [v.get_value(t) for v in self.prescribed_values]
         ),
         mode=self.mode,
-        is_explicit=self.is_explicit,
+        is_explicit=is_explicit,
         puff_decay_length=self.puff_decay_length.get_value(t),
         S_total=self.S_total.get_value(t),
+        target_line_average_n_e=self.target_line_average_n_e.get_value(t),
+        feedback_gain=self.feedback_gain.get_value(t),
     )
 
   def build_source(self) -> GasPuffSource:

torax/_src/sources/tests/gas_puff_source_test.py

@@ -12,8 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from absl.testing import absltest
+import chex
+from torax._src import math_utils
+from torax._src.config import build_runtime_params
+from torax._src.core_profiles import initialization
 from torax._src.sources import gas_puff_source
 from torax._src.sources.tests import test_lib
+from torax._src.test_utils import default_configs
+from torax._src.torax_pydantic import model_config
 
 
 class GasPuffSourceTest(test_lib.SingleProfileSourceTestCase):
@@ -25,6 +31,66 @@ def setUp(self):
         source_name=gas_puff_source.GasPuffSource.SOURCE_NAME,
     )
 
+  def test_feedback_mode(self):
+    """Tests calc_puff_feedback_source with real objects."""
+    config = default_configs.get_default_config_dict()
+    config['sources'] = {
+        'gas_puff': {
+            'model_name': 'feedback',
+            'feedback_gain': 10.0,
+        }
+    }
+    torax_config = model_config.ToraxConfig.from_dict(config)
+    geo = torax_config.geometry.build_provider(torax_config.numerics.t_initial)
+
+    runtime_params_provider = (
+        build_runtime_params.RuntimeParamsProvider.from_config(torax_config)
+    )
+    runtime_params = runtime_params_provider(t=torax_config.numerics.t_initial)
+
+    source_models = torax_config.sources.build_models()
+    neoclassical_models = torax_config.neoclassical.build_models()
+    core_profiles = initialization.initial_core_profiles(
+        runtime_params=runtime_params,
+        geo=geo,
+        source_models=source_models,
+        neoclassical_models=neoclassical_models,
+    )
+
+    initial_line_avg = math_utils.line_average(core_profiles.n_e.value, geo)
+
+    # Rebuild with specific requested value
+    config['sources']['gas_puff']['model_name'] = 'feedback'
+    config['sources']['gas_puff']['target_line_average_n_e'] = float(
+        initial_line_avg + 1e19
+    )
+    config['sources']['gas_puff']['feedback_gain'] = 10.0
+    torax_config = model_config.ToraxConfig.from_dict(config)
+    runtime_params = build_runtime_params.RuntimeParamsProvider.from_config(
+        torax_config
+    )(t=torax_config.numerics.t_initial)
+
+    gas_puff_params = runtime_params.sources['gas_puff']
+    chex.assert_trees_all_close(
+        gas_puff_params.target_line_average_n_e, float(initial_line_avg + 1e19)
+    )
+    chex.assert_trees_all_close(gas_puff_params.feedback_gain, 10.0)
+
+    profile = gas_puff_source.calc_puff_feedback_source(
+        runtime_params=runtime_params,
+        geo=geo,
+        source_name='gas_puff',
+        core_profiles=core_profiles,
+        unused_calculated_source_profiles=None,
+        unused_conductivity=None,
+    )[0]
+
+    total_particles = math_utils.volume_integration(profile, geo)
+
+    # Expected error = (initial_line_avg + 1e19) - initial_line_avg = 1e19
+    # Expected S_total = 10.0 * 1e19 = 1e20
+    chex.assert_trees_all_close(total_particles, 1e20, rtol=1e-5)
+
 
 if __name__ == '__main__':
   absltest.main()

torax/tests/test_data/test_iterhybrid_gas_puff_feedback.py

@@ -0,0 +1,36 @@
+# 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.
+
+"""ITER hybrid scenario with gas puff feedback control."""
+
+import copy
+from torax.tests.test_data import test_iterhybrid_predictor_corrector
+
+CONFIG = copy.deepcopy(test_iterhybrid_predictor_corrector.CONFIG)
+
+CONFIG['numerics']['t_final'] = 20.0
+
+# Configure gas puff to use feedback
+CONFIG['sources']['gas_puff'] = {
+    'model_name': 'feedback',
+    'puff_decay_length': 0.3,
+    'feedback_gain': 1e5,
+    'target_line_average_n_e': (
+        {
+            0.0: 8e19,
+            10.0: 9e19,
+        },
+        'STEP',
+    ),
+}