Create functionality to register sources and remove default source se…

…tting code.

This change enables functionality for sources to be registered without touching TORAX internals.

PiperOrigin-RevId: 681949209

torax/config/build_sim.py

@@ -22,9 +22,9 @@
 from torax import sim as sim_lib
 from torax.config import config_args
 from torax.config import runtime_params as runtime_params_lib
-from torax.sources import default_sources
 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
@@ -360,9 +360,8 @@ def _build_single_source_builder_from_config(
     source_config: dict[str, Any],
 ) -> source_lib.SourceBuilderProtocol:
   """Builds a source builder from the input config."""
-  runtime_params = default_sources.get_default_runtime_params(
-      source_name,
-  )
+  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:
@@ -395,7 +394,8 @@ def _build_single_source_builder_from_config(
   kwargs = {'runtime_params': runtime_params}
   if formula is not None:
     kwargs['formula'] = formula
-  return default_sources.get_source_builder_type(source_name)(**kwargs)
+
+  return registered_source.source_builder_class(**kwargs)
 
 
 def build_transport_model_builder_from_config(

torax/config/tests/runtime_params_slice.py

@@ -24,12 +24,12 @@
 from torax.config import profile_conditions as profile_conditions_lib
 from torax.config import runtime_params as general_runtime_params
 from torax.config import runtime_params_slice as runtime_params_slice_lib
-from torax.sources import default_sources
 from torax.sources import electron_density_sources
 from torax.sources import external_current_source
 from torax.sources import formula_config
 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
 
 

torax/fvm/tests/fvm.py

@@ -32,10 +32,10 @@
 from torax.fvm import cell_variable
 from torax.fvm import implicit_solve_block
 from torax.fvm import residual_and_loss
-from torax.sources import default_sources
 from torax.sources import runtime_params as source_runtime_params
 from torax.sources import source_models as source_models_lib
 from torax.stepper import runtime_params as stepper_runtime_params
+from torax.tests.test_lib import default_sources
 from torax.tests.test_lib import torax_refs
 from torax.transport_model import constant as constant_transport_model
 

torax/sources/register_source.py

@@ -0,0 +1,151 @@
+# Copyright 2024 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.
+"""Utilities for registering sources.
+
+This module contains a set of utilities for registering sources and retrieving
+registered sources.
+
+In TORAX we flexibly support different user provided sources to be active at
+runtime. To do so, we use a registration mechanism such that users can register
+their own sources and TORAX can look them up at runtime in the registry.
+
+To register a new source, use the `register_new_source` function. This function
+takes in a source name, source class, default runtime params class, and an
+optional source builder class. The source name is used to identify the source
+in the registry. The source class is the class of the source itself. The default
+runtime params class is the class of the default runtime params for the source.
+And the source builder class is an optional class which inherits from
+`SourceBuilderProtocol`. If not provided, then a default source builder is
+created which uses the source class and default runtime params class.
+
+Once a source is registered, it can be retrieved using the
+`get_registered_source` function. This function takes in a source name and
+returns a `RegisteredSource` dataclass containing the source class, source
+builder class, and default runtime params class. TORAX uses this dataclass to
+instantiate the source at runtime overriding any default runtime params with
+user provided ones from a config file.
+"""
+import dataclasses
+from typing import Type
+
+from torax.sources import bootstrap_current_source
+from torax.sources import bremsstrahlung_heat_sink
+from torax.sources import electron_density_sources
+from torax.sources import external_current_source
+from torax.sources import fusion_heat_source
+from torax.sources import generic_ion_el_heat_source as ion_el_heat
+from torax.sources import ohmic_heat_source
+from torax.sources import qei_source
+from torax.sources import runtime_params
+from torax.sources import source
+
+_REGISTERED_SOURCES = {}
+
+
+@dataclasses.dataclass(frozen=True)
+class RegisteredSource:
+  source_class: Type[source.Source]
+  source_builder_class: source.SourceBuilderProtocol
+  default_runtime_params_class: Type[runtime_params.RuntimeParams]
+
+
+def register_new_source(
+    source_name: str,
+    source_class: Type[source.Source],
+    default_runtime_params_class: Type[runtime_params.RuntimeParams],
+    source_builder_class: source.SourceBuilderProtocol | None = None,
+):
+  """Register a source class, default runtime params and (optional) builder."""
+  if source_name in _REGISTERED_SOURCES:
+    raise ValueError(f'Source:{source_name} has already been registered.')
+
+  if source_builder_class is None:
+    builder_class = source.make_source_builder(
+        source_class,
+        runtime_params_type=default_runtime_params_class,
+    )
+  else:
+    builder_class = source_builder_class
+
+  _REGISTERED_SOURCES[source_name] = RegisteredSource(
+      source_class=source_class,
+      source_builder_class=builder_class,
+      default_runtime_params_class=default_runtime_params_class,
+  )
+
+
+def get_registered_source(source_name: str) -> RegisteredSource:
+  """Used when building a simulation to get the registered source."""
+  if source_name in _REGISTERED_SOURCES:
+    return _REGISTERED_SOURCES[source_name]
+  else:
+    raise RuntimeError(f'Source:{source_name} has not been registered.')
+
+
+def register_torax_sources():
+  """Register a set of sources commonly used in TORAX."""
+  register_new_source(
+      'j_bootstrap',
+      source_class=bootstrap_current_source.BootstrapCurrentSource,
+      default_runtime_params_class=bootstrap_current_source.RuntimeParams,
+  )
+  register_new_source(
+      'jext',
+      external_current_source.ExternalCurrentSource,
+      default_runtime_params_class=external_current_source.RuntimeParams,
+  )
+  register_new_source(
+      'nbi_particle_source',
+      electron_density_sources.NBIParticleSource,
+      default_runtime_params_class=electron_density_sources.NBIParticleRuntimeParams,
+  )
+  register_new_source(
+      'gas_puff_source',
+      electron_density_sources.GasPuffSource,
+      default_runtime_params_class=electron_density_sources.GasPuffRuntimeParams,
+  )
+  register_new_source(
+      'pellet_source',
+      electron_density_sources.PelletSource,
+      default_runtime_params_class=electron_density_sources.PelletRuntimeParams,
+  )
+  register_new_source(
+      'generic_ion_el_heat_source',
+      ion_el_heat.GenericIonElectronHeatSource,
+      default_runtime_params_class=ion_el_heat.RuntimeParams,
+  )
+  register_new_source(
+      'fusion_heat_source',
+      fusion_heat_source.FusionHeatSource,
+      default_runtime_params_class=fusion_heat_source.FusionHeatSourceRuntimeParams
+  )
+  register_new_source(
+      'qei_source',
+      qei_source.QeiSource,
+      default_runtime_params_class=qei_source.RuntimeParams,
+  )
+  register_new_source(
+      'ohmic_heat_source',
+      ohmic_heat_source.OhmicHeatSource,
+      default_runtime_params_class=ohmic_heat_source.OhmicRuntimeParams,
+      source_builder_class=ohmic_heat_source.OhmicHeatSourceBuilder,
+  )
+  register_new_source(
+      'bremsstrahlung_heat_sink',
+      bremsstrahlung_heat_sink.BremsstrahlungHeatSink,
+      default_runtime_params_class=bremsstrahlung_heat_sink.RuntimeParams,
+  )
+
+
+register_torax_sources()