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,38 +836,13 @@ 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.
 
-``formula_type`` (str='default')
-  Sets the formula type if ``mode=='formula'``. The current options are:
-
-* ``'exponential'`` takes the following arguments:
-  * c1 (float): Offset location
-  * c2 (float): Exponential decay parameter
-  * total (float): integral
-
-  The profile is parameterized as follows :math:`Q = C e^{-(r - c1) / c2}` , where ``C`` is calculated to be consistent with ``total``. If ``use_normalized_r==True``,
-  then c1 and c2 are interpreted as being in normalized toroidal flux units.
-
-* ``'gaussian'`` takes the following arguments:
-  * c1 (float): Gaussian peak Location
-  * c2 (float): Gaussian width
-  * total (float): integral
-
-  The profile is parameterized as follows :math:`Q = C e^{-((r - c1)^2) / (2 c2^2)}` , where ``C`` is calculated to be consistent with ``total``. If ``use_normalized_r==True``,
-  then c1 and c2 are interpreted as being in normalized toroidal flux units.
-
-* ``'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
 ^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 A utility source module that allows for a time dependent Gaussian ion and electron heat source.
 
-``mode`` (str = 'formula')
-
-``formula_type`` (str = 'default')
-  Uses the Gaussian formula.
+``mode`` (str = 'model')
 
 ``rsource`` (float = 0.0), **time-varying-scalar**
   Gaussian center of source profile in units of :math:`\hat{\rho}`.
@@ -912,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.
+Exponential based gas puff source. No first-principle-based model is yet implemented in TORAX.
 
-``mode`` (str = 'formula')
-
-``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}`.
@@ -930,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}`.
@@ -949,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}`.
@@ -978,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}`.

docs/physics_models.rst

@@ -284,11 +284,8 @@ not need to be JAX-compatible, since explicit sources are an input into the PDE
 and do not require JIT compilation. Conversely, implicit treatment can be important for accurately
 resolving the impact of fast-evolving source terms.
 
-All sources can optionally be set to zero, prescribed with non-physics-based formulas
-(currently Gaussian or exponential) with user-configurable time-dependent parameters like
-amplitude, width, and location, or calculated with a dedicated physics-based model. Not
-all sources currently have a model implementation. However, the code modular structure
-facilitates easy coupling of additional source models in future work. Specifics of source models
+All sources can optionally be set to zero, prescribed with explicit values or calculated with a dedicated physics-based model.
+However, the code modular structure facilitates easy coupling of additional source models in future work. Specifics of source models
 currently implemented in TORAX follow:
 
 Ion-electron heat exchange

torax/__init__.py

@@ -73,6 +73,7 @@ def set_jax_precision():
     'geo_t',
     'geo_t_plus_dt',
     'geometry_provider',
+    'source_name',
     'x_old',
     'state',
     'unused_state',
@@ -88,6 +89,7 @@ def set_jax_precision():
     'source_profiles',
     'source_profile',
     'explicit_source_profiles',
+    'model_func',
     'source_models',
     'pedestal_model',
     'time_step_calculator',

torax/config/build_sim.py

@@ -25,8 +25,6 @@
 from torax.geometry import geometry_provider
 from torax.pedestal_model import pedestal_model as pedestal_model_lib
 from torax.pedestal_model import set_tped_nped
-from torax.sources import formula_config
-from torax.sources import formulas
 from torax.sources import register_source
 from torax.sources import runtime_params as source_runtime_params_lib
 from torax.sources import source as source_lib
@@ -353,48 +351,12 @@ def build_sources_builder_from_config(
         },
     }
 
-  If the `mode` is set to `formula_based`, then the you can provide a
-  `formula_type` key which may have the following values:
-
-  -  `default`: Uses the default impl (if the source has one) (default)
-
-    -  The other config args are based on the source's RuntimeParams object
-       outlined above.
-
-  -  `exponential`: Exponential profile.
-
-    - The other config args are from `sources.formula_config.Exponential`.
-
-  -  `gaussian`: Gaussian profile.
-
-    - The other config args are from `sources.formula_config.Gaussian`.
-
-  E.g. for an example heat source:
-
-  .. code-block:: python
-
-    {
-        mode: 'formula',
-        formula_type: 'gaussian',
-        total: 120e6,  # total heating
-        c1: 0.0,  # Source Gaussian central location (in normalized r)
-        c2: 0.25,  # Gaussian width in normalized radial coordinates
-    }
-
-  If you have custom source implementations, you may update this funtion to
-  handle those new sources and keys, or you may use the "advanced" configuration
-  method and build your `SourceModel` object directly.
-
   Args:
     source_configs: Input config dict defining all sources, with a structure as
       described above.
 
   Returns:
     A `SourceModelsBuilder`.
-
-  Raises:
-    ValueError if an input key doesn't match one of the source names defined
-      above.
   """
 
   source_builders = {
@@ -410,42 +372,38 @@ def _build_single_source_builder_from_config(
     source_config: dict[str, Any],
 ) -> source_lib.SourceBuilderProtocol:
   """Builds a source builder from the input config."""
-  registered_source = register_source.get_registered_source(source_name)
-  runtime_params = registered_source.default_runtime_params_class()
+  supported_source = register_source.get_supported_source(source_name)
+  if 'model_func' in source_config:
+    # If the user has specified a model function, try to retrive that from the
+    # registered source model functions.
+    model_func = source_config.pop('model_func')
+    model_function = supported_source.model_functions[model_func]
+  else:
+    # Otherwise, use the default model function.
+    model_function = supported_source.model_functions[
+        supported_source.source_class.DEFAULT_MODEL_FUNCTION_NAME
+    ]
+  runtime_params = model_function.runtime_params_class()
   # Update the defaults with the config provided.
   source_config = copy.copy(source_config)
   if 'mode' in source_config:
     mode = source_runtime_params_lib.Mode[source_config.pop('mode').upper()]
     runtime_params.mode = mode
-  formula = None
-  if 'formula_type' in source_config:
-    func = source_config.pop('formula_type').lower()
-    if func == 'default':
-      pass  # Nothing to do here.
-    elif func == 'exponential':
-      runtime_params.formula = config_args.recursive_replace(
-          formula_config.Exponential(),
-          ignore_extra_kwargs=True,
-          **source_config,
-      )
-      formula = formulas.Exponential()
-    elif func == 'gaussian':
-      runtime_params.formula = config_args.recursive_replace(
-          formula_config.Gaussian(),
-          ignore_extra_kwargs=True,
-          **source_config,
-      )
-      formula = formulas.Gaussian()
-    else:
-      raise ValueError(f'Unknown formula_type for source {source_name}: {func}')
   runtime_params = config_args.recursive_replace(
       runtime_params, ignore_extra_kwargs=True, **source_config
   )
   kwargs = {'runtime_params': runtime_params}
-  if formula is not None:
-    kwargs['formula'] = formula
 
-  return registered_source.source_builder_class(**kwargs)
+  source_builder_class = model_function.source_builder_class
+  if source_builder_class is None:
+    source_builder_class = source_lib.make_source_builder(
+        supported_source.source_class,
+        runtime_params_type=model_function.runtime_params_class,
+        links_back=model_function.links_back,
+        model_func=model_function.source_profile_function,
+    )
+
+  return source_builder_class(**kwargs)
 
 
 def build_transport_model_builder_from_config(

torax/config/tests/build_sim.py

@@ -23,8 +23,6 @@
 from torax.geometry import geometry
 from torax.geometry import geometry_provider
 from torax.pedestal_model import set_tped_nped
-from torax.sources import formula_config
-from torax.sources import formulas
 from torax.sources import runtime_params as source_runtime_params_lib
 from torax.stepper import linear_theta_method
 from torax.stepper import nonlinear_theta_method
@@ -366,39 +364,13 @@ def test_adding_standard_source_via_config(self):
     # pytype: enable=attribute-error
     self.assertEqual(
         source_models_builder.runtime_params['gas_puff_source'].mode,
-        source_runtime_params_lib.Mode.FORMULA_BASED,  # On by default.
+        source_runtime_params_lib.Mode.MODEL_BASED,  # On by default.
     )
     self.assertEqual(
         source_models_builder.runtime_params['ohmic_heat_source'].mode,
         source_runtime_params_lib.Mode.ZERO,
     )
 
-  def test_updating_formula_via_source_config(self):
-    """Tests that we can set the formula type and params via the config."""
-    source_models_builder = build_sim.build_sources_builder_from_config({
-        'gas_puff_source': {
-            'formula_type': 'gaussian',
-            'total': 1,
-            'c1': 2,
-            'c2': 3,
-        }
-    })
-    source_models = source_models_builder()
-    gas_source = source_models.sources['gas_puff_source']
-    self.assertIsInstance(gas_source.formula, formulas.Gaussian)
-    gas_source_runtime_params = source_models_builder.runtime_params[
-        'gas_puff_source'
-    ]
-    self.assertIsInstance(
-        gas_source_runtime_params.formula,
-        formula_config.Gaussian,
-    )
-    # pytype: disable=attribute-error
-    self.assertEqual(gas_source_runtime_params.formula.total, 1)
-    self.assertEqual(gas_source_runtime_params.formula.c1, 2)
-    self.assertEqual(gas_source_runtime_params.formula.c2, 3)
-    # pytype: enable=attribute-error
-
   def test_missing_transport_model_raises_error(self):
     with self.assertRaises(ValueError):
       build_sim.build_transport_model_builder_from_config({})

torax/config/tests/runtime_params_slice.py

@@ -26,9 +26,7 @@
 from torax.geometry import geometry
 from torax.pedestal_model import set_tped_nped
 from torax.sources import electron_density_sources
-from torax.sources import formula_config
 from torax.sources import generic_current_source
-from torax.sources import runtime_params as sources_params_lib
 from torax.stepper import runtime_params as stepper_params_lib
 from torax.tests.test_lib import default_sources
 from torax.transport_model import runtime_params as transport_params_lib
@@ -248,63 +246,6 @@ def test_source_formula_config_has_time_dependent_params(self):
       )
       np.testing.assert_allclose(generic_particle_source.S_tot, 4.0)
 
-    with self.subTest('exponential_formula'):
-      runtime_params = general_runtime_params.GeneralRuntimeParams()
-      dcs = runtime_params_slice_lib.DynamicRuntimeParamsSliceProvider(
-          runtime_params=runtime_params,
-          sources={
-              electron_density_sources.GasPuffSource.SOURCE_NAME: (
-                  sources_params_lib.RuntimeParams(
-                      formula=formula_config.Exponential(
-                          total={0.0: 0.0, 1.0: 1.0},
-                          c1={0.0: 0.0, 1.0: 2.0},
-                          c2={0.0: 0.0, 1.0: 3.0},
-                      )
-                  )
-              ),
-          },
-          torax_mesh=self._geo.torax_mesh,
-      )(
-          t=0.25,
-      )
-      gas_puff_source = dcs.sources[
-          electron_density_sources.GasPuffSource.SOURCE_NAME
-      ]
-      assert isinstance(
-          gas_puff_source.formula,
-          formula_config.DynamicExponential,
-      )
-      np.testing.assert_allclose(gas_puff_source.formula.total, 0.25)
-      np.testing.assert_allclose(gas_puff_source.formula.c1, 0.5)
-      np.testing.assert_allclose(gas_puff_source.formula.c2, 0.75)
-
-    with self.subTest('gaussian_formula'):
-      runtime_params = general_runtime_params.GeneralRuntimeParams()
-      dcs = runtime_params_slice_lib.DynamicRuntimeParamsSliceProvider(
-          runtime_params=runtime_params,
-          sources={
-              electron_density_sources.GasPuffSource.SOURCE_NAME: (
-                  sources_params_lib.RuntimeParams(
-                      formula=formula_config.Gaussian(
-                          total={0.0: 0.0, 1.0: 1.0},
-                          c1={0.0: 0.0, 1.0: 2.0},
-                          c2={0.0: 0.0, 1.0: 3.0},
-                      )
-                  )
-              ),
-          },
-          torax_mesh=self._geo.torax_mesh,
-      )(
-          t=0.25,
-      )
-      gas_puff_source = dcs.sources[
-          electron_density_sources.GasPuffSource.SOURCE_NAME
-      ]
-      assert isinstance(gas_puff_source.formula, formula_config.DynamicGaussian)
-      np.testing.assert_allclose(gas_puff_source.formula.total, 0.25)
-      np.testing.assert_allclose(gas_puff_source.formula.c1, 0.5)
-      np.testing.assert_allclose(gas_puff_source.formula.c2, 0.75)
-
   def test_wext_in_dynamic_runtime_params_cannot_be_negative(self):
     """Tests that wext cannot be negative."""
     runtime_params = general_runtime_params.GeneralRuntimeParams()