Remove the Exponential and Gaussian formulas from configuration.

The formulas are kept for general use but all the configuration objects are removed as they could be directly input by a user if needed (pending addition of a utility to register new model functions).

PiperOrigin-RevId: 707580250

docs/configuration.rst

@@ -840,25 +840,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.
 
-``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.
@@ -868,10 +849,7 @@ 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}`.

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,10 +25,7 @@
 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
 from torax.sources import source_models as source_models_lib
 from torax.stepper import linear_theta_method
@@ -353,48 +350,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 +371,33 @@ 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()
-  # 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}')
+  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()
   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()