Support prescribed sources for multiple affected profiles

Fixes #806

PiperOrigin-RevId: 744684193

Author: tamaranorman

docs/configuration.rst

@@ -918,9 +918,9 @@ The configurable runtime parameters of each source are as follows:
     This is documented in the individual source sections.
 
 * ``'PRESCRIBED'``
-    Source values are arbitrarily prescribed by the user. The value is set by ``prescribed_values``, and can contain the same
-    data structures as :ref:`Time-varying arrays`. Currently, this is only supported for sources that have a 1D output
-    along the cell grid or face grid.
+    Source values are arbitrarily prescribed by the user. The value is set by
+    ``prescribed_values``, and  should be a tuple of values. Each value can
+    contain the same data structures as :ref:`Time-varying arrays`.
 
 For example, to set 'fusion_power' to zero, e.g. for testing or sensitivity purposes, set:
 
@@ -938,7 +938,7 @@ preamble to the CONFIG dict within config module, set:
     'sources': {
         'generic_current_source': {
             'mode': 'PRESCRIBED',
-            'prescribed_values': (times, rhon, current_profiles),
+            'prescribed_values': ((times, rhon, current_profiles),),
         },
 
 where the example ``times`` is a 1D numpy array of times, ``rhon`` is a 1D numpy array of normalized toroidal flux

torax/sources/base.py

@@ -40,15 +40,15 @@ class SourceModelBase(torax_pydantic.BaseModelFrozen, abc.ABC):
       implicit or explicit. For example, file-based sources are always explicit.
       If an incorrect combination of source type and is_explicit is passed in,
       an error will be thrown when running the simulation.
-    prescribed_values: Prescribed values for the source. Used only when the
-      source is fully prescribed (i.e. source.mode == Mode.PRESCRIBED). The
-      default here is a vector of all zeros along for all rho and time, and the
-      output vector is along the cell grid.
+    prescribed_values: Tuple of prescribed values for the source, one for each
+      affected core profile. Used only when thesource is fully prescribed (i.e.
+      source.mode == Mode.PRESCRIBED). The default here is a vector of all zeros
+      along for all rho and time, and the output vector is along the cell grid.
   """
   mode: runtime_params.Mode = runtime_params.Mode.ZERO
   is_explicit: bool = False
-  prescribed_values: torax_pydantic.TimeVaryingArray = (
-      torax_pydantic.ValidatedDefault({0: {0: 0, 1: 0}})
+  prescribed_values: tuple[torax_pydantic.TimeVaryingArray, ...] = (
+      torax_pydantic.ValidatedDefault(({0: {0: 0, 1: 0}},))
   )
 
   def build_static_params(self) -> runtime_params.StaticRuntimeParams:

torax/sources/bootstrap_current_source.py

@@ -151,7 +151,9 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> DynamicRuntimeParams:
     return DynamicRuntimeParams(
-        prescribed_values=self.prescribed_values.get_value(t),
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
         bootstrap_mult=self.bootstrap_mult,
     )
 

torax/sources/bremsstrahlung_heat_sink.py

@@ -159,7 +159,9 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> 'DynamicRuntimeParams':
     return DynamicRuntimeParams(
-        prescribed_values=self.prescribed_values.get_value(t),
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
         use_relativistic_correction=self.use_relativistic_correction,
     )
 

torax/sources/cyclotron_radiation_heat_sink.py

@@ -415,7 +415,9 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> 'DynamicRuntimeParams':
     return DynamicRuntimeParams(
-        prescribed_values=self.prescribed_values.get_value(t),
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
         wall_reflection_coeff=self.wall_reflection_coeff,
     )
 

torax/sources/electron_cyclotron_source.py

@@ -176,7 +176,9 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> DynamicRuntimeParams:
     return DynamicRuntimeParams(
-        prescribed_values=self.prescribed_values.get_value(t),
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
         cd_efficiency=self.cd_efficiency.get_value(t),
         manual_ec_power_density=self.manual_ec_power_density.get_value(t),
         gaussian_ec_power_density_width=self.gaussian_ec_power_density_width.get_value(

torax/sources/fusion_heat_source.py

@@ -198,7 +198,9 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> runtime_params_lib.DynamicRuntimeParams:
     return runtime_params_lib.DynamicRuntimeParams(
-        prescribed_values=self.prescribed_values.get_value(t),
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
     )
 
   def build_source(self) -> FusionHeatSource:

torax/sources/gas_puff_source.py

@@ -110,7 +110,9 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> DynamicGasPuffRuntimeParams:
     return DynamicGasPuffRuntimeParams(
-        prescribed_values=self.prescribed_values.get_value(t),
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
         puff_decay_length=self.puff_decay_length.get_value(t),
         S_puff_tot=self.S_puff_tot.get_value(t),
     )

torax/sources/generic_current_source.py

@@ -151,7 +151,9 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> DynamicRuntimeParams:
     return DynamicRuntimeParams(
-        prescribed_values=self.prescribed_values.get_value(t),
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
         Iext=self.Iext.get_value(t),
         fext=self.fext.get_value(t),
         wext=self.wext.get_value(t),

torax/sources/generic_ion_el_heat_source.py

@@ -158,7 +158,9 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> DynamicRuntimeParams:
     return DynamicRuntimeParams(
-        prescribed_values=self.prescribed_values.get_value(t),
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
         w=self.w.get_value(t),
         rsource=self.rsource.get_value(t),
         Ptot=self.Ptot.get_value(t),

torax/sources/generic_particle_source.py

@@ -121,7 +121,9 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> DynamicParticleRuntimeParams:
     return DynamicParticleRuntimeParams(
-        prescribed_values=self.prescribed_values.get_value(t),
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
         particle_width=self.particle_width.get_value(t),
         deposition_location=self.deposition_location.get_value(t),
         S_tot=self.S_tot.get_value(t),

torax/sources/impurity_radiation_heat_sink/impurity_radiation_constant_fraction.py

@@ -112,7 +112,9 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> 'DynamicRuntimeParams':
     return DynamicRuntimeParams(
-        prescribed_values=self.prescribed_values.get_value(t),
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
         fraction_of_total_power_density=self.fraction_of_total_power_density.get_value(
             t
         ),

torax/sources/impurity_radiation_heat_sink/impurity_radiation_mavrin_fit.py

@@ -242,7 +242,9 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> 'DynamicRuntimeParams':
     return DynamicRuntimeParams(
-        prescribed_values=self.prescribed_values.get_value(t),
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
         radiation_multiplier=self.radiation_multiplier,
     )
 

torax/sources/ion_cyclotron_source.py

@@ -489,7 +489,9 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> DynamicRuntimeParams:
     return DynamicRuntimeParams(
-        prescribed_values=self.prescribed_values.get_value(t),
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
         wall_inner=self.wall_inner,
         wall_outer=self.wall_outer,
         frequency=self.frequency.get_value(t),

torax/sources/ohmic_heat_source.py

@@ -103,7 +103,9 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> runtime_params_lib.DynamicRuntimeParams:
     return runtime_params_lib.DynamicRuntimeParams(
-        prescribed_values=self.prescribed_values.get_value(t),
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
     )
 
   def build_source(self) -> OhmicHeatSource:

torax/sources/pellet_source.py

@@ -116,7 +116,9 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> DynamicPelletRuntimeParams:
     return DynamicPelletRuntimeParams(
-        prescribed_values=self.prescribed_values.get_value(t),
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
         pellet_width=self.pellet_width.get_value(t),
         pellet_deposition_location=self.pellet_deposition_location.get_value(t),
         S_pellet_tot=self.S_pellet_tot.get_value(t),

torax/sources/qei_source.py

@@ -166,7 +166,9 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> DynamicRuntimeParams:
     return DynamicRuntimeParams(
-        prescribed_values=self.prescribed_values.get_value(t),
+        prescribed_values=tuple(
+            [v.get_value(t) for v in self.prescribed_values]
+        ),
         Qei_mult=self.Qei_mult,
     )
 

torax/sources/runtime_params.py

@@ -48,7 +48,7 @@ class DynamicRuntimeParams:
   stateless, so these params are their inputs to determine their output
   profiles.
   """
-  prescribed_values: array_typing.ArrayFloat
+  prescribed_values: tuple[array_typing.ArrayFloat, ...]
 
 
 @chex.dataclass(frozen=True)

torax/sources/source.py

@@ -163,9 +163,16 @@ def get_value(
             calculated_source_profiles,
         )
       case runtime_params_lib.Mode.PRESCRIBED.value:
-        # TODO(b/395854896) add support for sources that affect multiple core
-        # profiles.
-        return (dynamic_source_runtime_params.prescribed_values,)
+        if len(self.affected_core_profiles) != len(
+            dynamic_source_runtime_params.prescribed_values
+        ):
+          raise ValueError(
+              'When using PRESCRIBED mode, the number of prescribed values must'
+              ' match the number of affected core profiles. Was: '
+              f'{len(dynamic_source_runtime_params.prescribed_values)} '
+              f' Expected: {len(self.affected_core_profiles)}.'
+          )
+        return dynamic_source_runtime_params.prescribed_values
       case runtime_params_lib.Mode.ZERO.value:
         zeros = jnp.zeros(geo.rho_norm.shape)
         return (zeros,) * len(self.affected_core_profiles)

torax/sources/tests/pydantic_model_test.py

@@ -15,6 +15,7 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
+import numpy as np
 from torax.sources import base
 from torax.sources import bootstrap_current_source
 from torax.sources import fusion_heat_source
@@ -26,6 +27,7 @@
 from torax.sources import source_models as source_models_lib
 from torax.sources.impurity_radiation_heat_sink import impurity_radiation_constant_fraction
 from torax.sources.impurity_radiation_heat_sink import impurity_radiation_mavrin_fit
+from torax.torax_pydantic import torax_pydantic
 
 
 class PydanticModelTest(parameterized.TestCase):
@@ -150,6 +152,86 @@ def test_empty_source_config_only_has_defaults_turned_off(self):
     )
     self.assertLen(sources.source_model_config, 3)
 
+  def test_adding_a_source_with_prescribed_values(self):
+    """Tests that a source can be added with overriding defaults."""
+    sources = pydantic_model.Sources.from_dict({
+        'generic_current_source': {
+            'mode': 'PRESCRIBED',
+            'prescribed_values': ((
+                np.array([0.0, 1.0, 2.0, 3.0]),
+                np.array([0., 0.5, 1.0]),
+                np.full([4, 3], 42)
+            ),),
+        },
+        'electron_cyclotron_source': {
+            'mode': 'PRESCRIBED',
+            'prescribed_values': (
+                3.,
+                4.,
+            ),
+        }
+    })
+    mesh = torax_pydantic.Grid1D(nx=4, dx=0.25)
+    torax_pydantic.set_grid(sources, mesh)
+    source = sources.source_model_config['generic_current_source']
+    self.assertLen(source.prescribed_values, 1)
+    self.assertIsInstance(
+        source.prescribed_values[0], torax_pydantic.TimeVaryingArray)
+    source = sources.source_model_config['electron_cyclotron_source']
+    self.assertLen(source.prescribed_values, 2)
+    self.assertIsInstance(
+        source.prescribed_values[0], torax_pydantic.TimeVaryingArray)
+    self.assertIsInstance(
+        source.prescribed_values[1], torax_pydantic.TimeVaryingArray)
+    value = source.prescribed_values[0].get_value(0.0)
+    np.testing.assert_equal(value, 3.)
+    value = source.prescribed_values[1].get_value(0.0)
+    np.testing.assert_equal(value, 4.)
+
+  def test_bremsstrahlung_and_mavrin_validator_with_bremsstrahlung_zero(self):
+    valid_config = {
+        'bremsstrahlung_heat_sink': {'mode': 'ZERO'},
+        'impurity_radiation_heat_sink': {
+            'mode': 'PRESCRIBED',
+            'model_function_name': 'impurity_radiation_mavrin_fit',
+        },
+    }
+    pydantic_model.Sources.from_dict(valid_config)
+
+  def test_bremsstrahlung_and_mavrin_validator_with_mavrin_zero(self):
+    valid_config = {
+        'bremsstrahlung_heat_sink': {'mode': 'PRESCRIBED'},
+        'impurity_radiation_heat_sink': {
+            'mode': 'ZERO',
+            'model_function_name': 'impurity_radiation_mavrin_fit',
+        },
+    }
+    pydantic_model.Sources.from_dict(valid_config)
+
+  def test_bremsstrahlung_and_mavrin_validator_with_constant_fraction(self):
+    valid_config = {
+        'bremsstrahlung_heat_sink': {'mode': 'PRESCRIBED'},
+        'impurity_radiation_heat_sink': {
+            'mode': 'PRESCRIBED',
+            'model_function_name': 'radially_constant_fraction_of_Pin',
+        },
+    }
+    pydantic_model.Sources.from_dict(valid_config)
+
+  def test_bremsstrahlung_and_mavrin_validator_with_invalid_config(self):
+    invalid_config = {
+        'bremsstrahlung_heat_sink': {'mode': 'PRESCRIBED'},
+        'impurity_radiation_heat_sink': {
+            'mode': 'PRESCRIBED',
+            'model_function_name': 'impurity_radiation_mavrin_fit',
+        },
+    }
+    with self.assertRaisesRegex(
+        ValueError,
+        'Both bremsstrahlung_heat_sink and impurity_radiation_heat_sink',
+    ):
+      pydantic_model.Sources.from_dict(invalid_config)
+
 
 if __name__ == '__main__':
   absltest.main()

torax/sources/tests/source_profile_builders_test.py

@@ -130,7 +130,7 @@ def affected_core_profiles(
         runtime_params_slice.DynamicRuntimeParamsSlice,
         sources={
             'foo': source_runtime_params.DynamicRuntimeParams(
-                prescribed_values=jnp.ones(self.geo.rho.shape)
+                prescribed_values=(jnp.ones(self.geo.rho.shape),)
             )
         },
     )
@@ -191,7 +191,8 @@ def affected_core_profiles(
         runtime_params_slice.DynamicRuntimeParamsSlice,
         sources={
             'foo': source_runtime_params.DynamicRuntimeParams(
-                prescribed_values=jnp.ones(self.geo.rho.shape)
+                prescribed_values=(jnp.ones(self.geo.rho.shape),
+                                   jnp.ones(self.geo.rho.shape))
             )
         },
     )
@@ -279,7 +280,7 @@ def affected_core_profiles(
         runtime_params_slice.DynamicRuntimeParamsSlice,
         sources={
             'foo': source_runtime_params.DynamicRuntimeParams(
-                prescribed_values=jnp.ones(self.geo.rho.shape)
+                prescribed_values=(jnp.ones(self.geo.rho.shape),)
             )
         },
     )