Delete formula, affected_core_profiles and check_mode from `Sou…

…rce` API and runtime_params.FORMULA_BASED.

PiperOrigin-RevId: 707889575

docs/configuration.rst

@@ -803,10 +803,6 @@ The configurable runtime parameters of each source are as follows:
     Source values come from a model in code. Specific model selection is not yet available in TORAX since there are no source components with more than one
     physics model. However, this will be straightforward to develop when that occurs.
 
-* ``'FORMULA'``
-    Source values come from a prescribed (possibly time-dependent) formula that is not dependent on the state of the system. The formula type (Gaussian, exponential)
-    is set by ``formula_type``.
-
 * ``'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`.
@@ -840,9 +836,6 @@ and can be set to anything convenient.
   beginning of a time step, or do not have any dependance on state. Implicit sources depend on updated states as the iterative solvers evolve the state through the
   course of a time step. If a source model is complex but evolves over slow timescales compared to the state, it may be beneficial to set it as explicit.
 
-* ``'default'``
-    Some sources have default implementations which use the above formulas under the hood with intuitive parameter names for c1 and c2.
-    Consult the list below for further details.
 
 generic_ion_el_heat_source
 ^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -890,12 +883,9 @@ Fusion power assuming a 50-50 D-T ion distribution.
 gas_puff_source
 ^^^^^^^^^^^^^^^
 
-Formula based exponential gas puff source. No first-principle-based model is yet implemented in TORAX.
-
-``mode`` (str = 'formula')
+Exponential based gas puff source. No first-principle-based model is yet implemented in TORAX.
 
-``formula_type`` (str = 'default')
-  Uses the exponential formula with ``c1=1``.
+``mode`` (str = 'model')
 
 ``puff_decay_length`` (float = 0.05), **time-varying-scalar**
   Gas puff decay length from edge in units of :math:`\hat{\rho}`.
@@ -908,10 +898,7 @@ pellet_source
 
 Time dependent Gaussian pellet source. No first-principle-based model is yet implemented in TORAX.
 
-``mode`` (str = 'formula')
-
-``formula_type`` (str = 'default')
-  Uses the Gaussian formula.
+``mode`` (str = 'model')
 
 ``pellet_deposition_location`` (float = 0.85), **time-varying-scalar**
   Gaussian center of source profile in units of :math:`\hat{\rho}`.
@@ -927,10 +914,7 @@ generic_particle_source
 
 Time dependent Gaussian particle source. No first-principle-based model is yet implemented in TORAX.
 
-``mode`` (str = 'formula')
-
-``formula_type`` (str = 'default')
-  Uses the Gaussian formula with.
+``mode`` (str = 'model')
 
 ``deposition_location`` (float = 0.0), **time-varying-scalar**
   Gaussian center of source profile in units of :math:`\hat{\rho}`.
@@ -956,10 +940,7 @@ generic_current_source
 
 Generic external current profile, parameterized as a Gaussian.
 
-``mode`` (str = 'formula')
-
-``formula_type`` (str = 'default')
-  Uses the Gaussian formula.
+``mode`` (str = 'model')
 
 ``rext`` (float = 0.4), **time-varying-scalar**
   Gaussian center of current profile in units of :math:`\hat{\rho}`.

torax/sources/bootstrap_current_source.py

@@ -126,8 +126,6 @@ def get_value(
     dynamic_source_runtime_params = dynamic_runtime_params_slice.sources[
         self.source_name
     ]
-    # Make sure the input mode requested is supported.
-    self.check_mode(static_source_runtime_params.mode)
     # Make sure the input params are the correct type.
     if not isinstance(dynamic_source_runtime_params, DynamicRuntimeParams):
       raise TypeError(

torax/sources/bremsstrahlung_heat_sink.py

@@ -160,15 +160,6 @@ class BremsstrahlungHeatSink(source.Source):
   def source_name(self) -> str:
     return self.SOURCE_NAME
 
-  @property
-  def supported_modes(self) -> tuple[runtime_params_lib.Mode, ...]:
-    """Returns the modes supported by this source."""
-    return (
-        runtime_params_lib.Mode.ZERO,
-        runtime_params_lib.Mode.MODEL_BASED,
-        runtime_params_lib.Mode.PRESCRIBED,
-    )
-
   @property
   def affected_core_profiles(self) -> tuple[source.AffectedCoreProfile, ...]:
     return (source.AffectedCoreProfile.TEMP_EL,)

torax/sources/generic_current_source.py

@@ -284,7 +284,6 @@ def generic_current_source_hires(
         self.source_name
     ]
     assert isinstance(dynamic_source_runtime_params, DynamicRuntimeParams)
-    self.check_mode(static_source_runtime_params.mode)
 
     # Interpolate prescribed values onto the hires grid
     hires_prescribed_values = jnp.where(
@@ -305,7 +304,6 @@ def generic_current_source_hires(
         core_profiles=None,
         # There is no model for this source.
         model_func=self.hires_formula,
-        formula=None,
         output_shape=geo.rho_hires_norm.shape,
         prescribed_values=hires_prescribed_values,
         source_models=getattr(self, 'source_models', None),

torax/sources/impurity_radiation_heat_sink.py

@@ -152,15 +152,6 @@ class ImpurityRadiationHeatSink(source_lib.Source):
   def source_name(self) -> str:
     return self.SOURCE_NAME
 
-  @property
-  def supported_modes(self) -> tuple[runtime_params_lib.Mode, ...]:
-    """Returns the modes supported by this source."""
-    return (
-        runtime_params_lib.Mode.ZERO,
-        runtime_params_lib.Mode.MODEL_BASED,
-        runtime_params_lib.Mode.PRESCRIBED,
-    )
-
   @property
   def affected_core_profiles(
       self,

torax/sources/ion_cyclotron_source.py

@@ -499,14 +499,6 @@ class IonCyclotronSource(source.Source):
   def source_name(self) -> str:
     return self.SOURCE_NAME
 
-  @property
-  def supported_modes(self) -> tuple[runtime_params_lib.Mode, ...]:
-    return (
-        runtime_params_lib.Mode.ZERO,
-        runtime_params_lib.Mode.MODEL_BASED,
-        runtime_params_lib.Mode.PRESCRIBED,
-    )
-
   @property
   def affected_core_profiles(self) -> tuple[source.AffectedCoreProfile, ...]:
     return (
@@ -538,7 +530,8 @@ def __post_init__(self):
 
   def __call__(
       self,
-      formula: source.SourceProfileFunction | None = None,
   ) -> IonCyclotronSource:
 
-    return IonCyclotronSource(formula=formula, model_func=self.model_func,)
+    return IonCyclotronSource(
+        model_func=self.model_func,
+    )

torax/sources/ohmic_heat_source.py

@@ -199,14 +199,6 @@ class OhmicHeatSource(source_lib.Source):
   def source_name(self) -> str:
     return self.SOURCE_NAME
 
-  @property
-  def supported_modes(self) -> tuple[runtime_params_lib.Mode, ...]:
-    return (
-        runtime_params_lib.Mode.ZERO,
-        runtime_params_lib.Mode.MODEL_BASED,
-        runtime_params_lib.Mode.PRESCRIBED,
-    )
-
   @property
   def affected_core_profiles(
       self,

torax/sources/qei_source.py

@@ -78,14 +78,6 @@ class QeiSource(source.Source):
   def source_name(self) -> str:
     return self.SOURCE_NAME
 
-  @property
-  def supported_modes(self) -> tuple[runtime_params_lib.Mode, ...]:
-    return (
-        runtime_params_lib.Mode.ZERO,
-        runtime_params_lib.Mode.MODEL_BASED,
-        runtime_params_lib.Mode.PRESCRIBED,
-    )
-
   @property
   def affected_core_profiles(self) -> tuple[source.AffectedCoreProfile, ...]:
     return (
@@ -101,7 +93,6 @@ def get_qei(
       core_profiles: state.CoreProfiles,
   ) -> source_profiles.QeiInfo:
     """Computes the value of the source."""
-    self.check_mode(static_runtime_params_slice.sources[self.source_name].mode)
     dynamic_source_runtime_params = dynamic_runtime_params_slice.sources[
         self.source_name
     ]

torax/sources/runtime_params.py

@@ -41,15 +41,10 @@ class Mode(enum.Enum):
   # explicit depending on the model implementation.
   MODEL_BASED = 1
 
-  # Source values come from a prescribed (possibly time-dependent) formula that
-  # is not dependent on the state of the system. These formulas may be dependent
-  # on the config and geometry of the system.
-  FORMULA_BASED = 2
-
   # Source values come from a pre-determined set of values, that may evolve in
   # time. Values can be drawn from a file or an array. These sources are always
   # explicit.
-  PRESCRIBED = 3
+  PRESCRIBED = 2
 
 
 @dataclasses.dataclass

torax/sources/source.py

@@ -118,11 +118,6 @@ class Source(abc.ABC):
       By default, the number of affected core profiles should equal the rank of
       the output shape returned by output_shape_getter. Subclasses may override
       this requirement.
-    supported_modes: Defines how the source computes its profile. Can be set to
-      zero, model-based, etc. At runtime, the input config (the RuntimeParams or
-      the DynamicRuntimeParams) will specify which supported type the Source is
-      running with. If the runtime config specifies an unsupported type, an
-      error will raise.
     output_shape_getter: Callable which returns the shape of the profiles given
       by this source.
     model_func: The function used when the the runtime type is set to
@@ -134,7 +129,6 @@ class Source(abc.ABC):
   """
   DEFAULT_MODEL_FUNCTION_NAME: ClassVar[str] = 'default'
   model_func: SourceProfileFunction | None = None
-  formula: SourceProfileFunction | None = None
 
   @property
   @abc.abstractmethod
@@ -151,31 +145,10 @@ def output_shape_getter(self) -> SourceOutputShapeFunction:
     """Returns a function which returns the shape of the source's output."""
     return get_cell_profile_shape
 
-  @property
-  def supported_modes(self) -> tuple[runtime_params_lib.Mode, ...]:
-    """Returns the modes supported by this source."""
-    return (
-        runtime_params_lib.Mode.ZERO,
-        runtime_params_lib.Mode.FORMULA_BASED,
-        runtime_params_lib.Mode.MODEL_BASED,
-        runtime_params_lib.Mode.PRESCRIBED,
-    )
-
   @property
   def affected_core_profiles_ints(self) -> tuple[int, ...]:
     return tuple([int(cp) for cp in self.affected_core_profiles])
 
-  def check_mode(
-      self,
-      mode: int,
-  ):
-    """Raises an error if the source type is not supported."""
-    if runtime_params_lib.Mode(mode) not in self.supported_modes:
-      raise ValueError(
-          f'This source supports the following modes: {self.supported_modes}.'
-          f' Unsupported mode provided: {mode}.'
-      )
-
   def get_value(
       self,
       static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
@@ -200,13 +173,9 @@ def get_value(
     Returns:
       Array, arrays, or nested dataclass/dict of arrays for the source profile.
     """
-    static_source_runtime_params = static_runtime_params_slice.sources[
-        self.source_name
-    ]
     dynamic_source_runtime_params = dynamic_runtime_params_slice.sources[
         self.source_name
     ]
-    self.check_mode(static_source_runtime_params.mode)
     output_shape = self.output_shape_getter(geo)
 
     return get_source_profiles(
@@ -215,7 +184,6 @@ def get_value(
         geo=geo,
         core_profiles=core_profiles,
         model_func=self.model_func,
-        formula=self.formula,
         prescribed_values=dynamic_source_runtime_params.prescribed_values,
         output_shape=output_shape,
         source_models=getattr(self, 'source_models', None),
@@ -306,7 +274,6 @@ def get_source_profiles(
     source_name: str,
     core_profiles: state.CoreProfiles,
     model_func: SourceProfileFunction | None,
-    formula: SourceProfileFunction | None,
     prescribed_values: chex.Array,
     output_shape: tuple[int, ...],
     source_models: Optional['source_models.SourceModels'],
@@ -327,7 +294,6 @@ def get_source_profiles(
     core_profiles: Core plasma profiles. Used as input to the source profile
       functions.
     model_func: Model function.
-    formula: Formula implementation.
     prescribed_values: Array of values for this timeslice, interpolated onto the
       grid (ie with shape output_shape)
     output_shape: Expected shape of the output array.
@@ -352,19 +318,6 @@ def get_source_profiles(
           core_profiles,
           source_models,
       )
-    case runtime_params_lib.Mode.FORMULA_BASED.value:
-      if formula is None:
-        raise ValueError(
-            'Source is in FORMULA_BASED mode but has no formula function.'
-        )
-      return formula(
-          static_runtime_params_slice,
-          dynamic_runtime_params_slice,
-          geo,
-          source_name,
-          core_profiles,
-          source_models,
-      )
     case runtime_params_lib.Mode.PRESCRIBED.value:
       return prescribed_values
     case _:

torax/sources/tests/bootstrap_current_source.py

@@ -19,7 +19,6 @@
 import numpy as np
 from torax.geometry import geometry
 from torax.sources import bootstrap_current_source
-from torax.sources import runtime_params as runtime_params_lib
 from torax.sources import source as source_lib
 from torax.sources import source_profiles
 from torax.sources.tests import test_lib
@@ -33,9 +32,6 @@ def setUpClass(cls):
     super().setUpClass(
         source_class=bootstrap_current_source.BootstrapCurrentSource,
         runtime_params_class=bootstrap_current_source.RuntimeParams,
-        unsupported_modes=[
-            runtime_params_lib.Mode.FORMULA_BASED,
-        ],
         source_name=bootstrap_current_source.BootstrapCurrentSource.SOURCE_NAME,
         model_func=None,
     )

torax/sources/tests/bremsstrahlung_heat_sink.py

@@ -23,7 +23,6 @@
 from torax import core_profile_setters
 from torax.config import runtime_params_slice
 from torax.sources import bremsstrahlung_heat_sink
-from torax.sources import runtime_params as runtime_params_lib
 from torax.sources import source_models as source_models_lib
 from torax.sources.tests import test_lib
 from torax.stepper import runtime_params as stepper_runtime_params
@@ -38,9 +37,6 @@ def setUpClass(cls):
     super().setUpClass(
         source_class=bremsstrahlung_heat_sink.BremsstrahlungHeatSink,
         runtime_params_class=bremsstrahlung_heat_sink.RuntimeParams,
-        unsupported_modes=[
-            runtime_params_lib.Mode.FORMULA_BASED,
-        ],
         source_name=bremsstrahlung_heat_sink.BremsstrahlungHeatSink.SOURCE_NAME,
         model_func=bremsstrahlung_heat_sink.bremsstrahlung_model_func,
     )