Make all sources discriminated on model function name.

PiperOrigin-RevId: 742220027

Author: Nush395

torax/sources/base.py

@@ -23,6 +23,12 @@
 class SourceModelBase(torax_pydantic.BaseModelFrozen, abc.ABC):
   """Base model holding parameters common to all source models.
 
+  Subclasses should define the `model_function_name` attribute as a `Literal`
+  string. This string should match the name of the function that calculates the
+  source profile. This is used as an identifier for the model function in the
+  source config for Pydantic to "discriminate" against. This should be given a
+  unique value for each source model function implementation.
+
   Attributes:
     mode: Defines how the source values are computed (from a model, from a file,
       etc.)
@@ -39,7 +45,6 @@ class SourceModelBase(torax_pydantic.BaseModelFrozen, abc.ABC):
       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/sources/bootstrap_current_source.py

@@ -14,7 +14,7 @@
 
 """bootstrap current source profile."""
 import dataclasses
-from typing import ClassVar, Literal
+from typing import ClassVar
 
 import chex
 import jax
@@ -51,7 +51,6 @@ class BootstrapCurrentSource(source.Source):
   """
 
   SOURCE_NAME: ClassVar[str] = 'j_bootstrap'
-  DEFAULT_MODEL_FUNCTION_NAME: ClassVar[str] = 'calc_neoclassical'
 
   @property
   def source_name(self) -> str:
@@ -136,8 +135,6 @@ class BootstrapCurrentSourceConfig(base.SourceModelBase):
   Attributes:
     bootstrap_mult: Multiplication factor for bootstrap current.
   """
-
-  source_name: Literal['j_bootstrap'] = 'j_bootstrap'
   bootstrap_mult: float = 1.0
   mode: runtime_params_lib.Mode = runtime_params_lib.Mode.MODEL_BASED
 

torax/sources/bremsstrahlung_heat_sink.py

@@ -16,7 +16,7 @@
 
 """Bremsstrahlung heat sink for electron heat equation.."""
 import dataclasses
-from typing import ClassVar, Literal
+from typing import ClassVar, Final, Literal
 
 import chex
 import jax
@@ -31,6 +31,12 @@
 from torax.sources import source_profiles
 
 
+# Default value for the model function to be used for the Bremsstrahlung heat
+# sink. This is also used as an identifier for the model function in the default
+# source config for Pydantic to "discriminate" against.
+DEFAULT_MODEL_FUNCTION_NAME: Final[str] = 'bremsstrahlung_model_func'
+
+
 @chex.dataclass(frozen=True)
 class DynamicRuntimeParams(runtime_params_lib.DynamicRuntimeParams):
   use_relativistic_correction: bool
@@ -121,7 +127,6 @@ class BremsstrahlungHeatSink(source.Source):
   """Brehmsstrahlung heat sink for electron heat equation."""
 
   SOURCE_NAME: ClassVar[str] = 'bremsstrahlung_heat_sink'
-  DEFAULT_MODEL_FUNCTION_NAME: ClassVar[str] = 'bremsstrahlung_model_func'
   model_func: source.SourceProfileFunction = bremsstrahlung_model_func
 
   @property
@@ -139,8 +144,9 @@ class BremsstrahlungHeatSinkConfig(base.SourceModelBase):
   Attributes:
     use_relativistic_correction: Whether to use relativistic correction.
   """
-
-  source_name: Literal['bremsstrahlung_heat_sink'] = 'bremsstrahlung_heat_sink'
+  model_function_name: Literal['bremsstrahlung_model_func'] = (
+      'bremsstrahlung_model_func'
+  )
   use_relativistic_correction: bool = False
   mode: runtime_params_lib.Mode = runtime_params_lib.Mode.MODEL_BASED
 

torax/sources/cyclotron_radiation_heat_sink.py

@@ -35,6 +35,12 @@
 import typing_extensions
 
 
+# Default value for the model function to be used for the Cyclotron radiation
+# heat sink source. This is also used as an identifier for the model function in
+# the default source config for Pydantic to "discriminate" against.
+DEFAULT_MODEL_FUNCTION_NAME: str = 'cyclotron_radiation_albajar'
+
+
 @chex.dataclass(frozen=True)
 class StaticRuntimeParams(runtime_params_lib.StaticRuntimeParams):
   beta_min: float
@@ -359,7 +365,6 @@ class CyclotronRadiationHeatSink(source.Source):
   """Cyclotron radiation heat sink for electron heat equation."""
 
   SOURCE_NAME: ClassVar[str] = 'cyclotron_radiation_heat_sink'
-  DEFAULT_MODEL_FUNCTION_NAME: ClassVar[str] = 'cyclotron_radiation_albajar'
   model_func: source.SourceProfileFunction = cyclotron_radiation_albajar
 
   @property
@@ -386,9 +391,8 @@ class CyclotronRadiationHeatSinkConfig(base.SourceModelBase):
     beta_grid_size: The number of points to use in the grid search for the best
       fit of the temperature function.
   """
-
-  source_name: Literal['cyclotron_radiation_heat_sink'] = (
-      'cyclotron_radiation_heat_sink'
+  model_function_name: Literal['cyclotron_radiation_albajar'] = (
+      'cyclotron_radiation_albajar'
   )
   mode: runtime_params_lib.Mode = runtime_params_lib.Mode.MODEL_BASED
   wall_reflection_coeff: float = 0.9

torax/sources/electron_cyclotron_source.py

@@ -30,9 +30,11 @@
 from torax.sources import source_profiles
 from torax.torax_pydantic import torax_pydantic
 
-InterpolatedVarTimeRhoInput = (
-    runtime_params_lib.interpolated_param.InterpolatedVarTimeRhoInput
-)
+
+# Default value for the model function to be used for the electron cyclotron
+# source. This is also used as an identifier for the model function in
+# the default source config for Pydantic to "discriminate" against.
+DEFAULT_MODEL_FUNCTION_NAME: str = "calc_heating_and_current"
 
 
 @chex.dataclass(frozen=True)
@@ -122,7 +124,6 @@ class ElectronCyclotronSource(source.Source):
   """Electron cyclotron source for the Te and Psi equations."""
 
   SOURCE_NAME: ClassVar[str] = "electron_cyclotron_source"
-  DEFAULT_MODEL_FUNCTION_NAME: ClassVar[str] = "calc_heating_and_current"
   model_func: source.SourceProfileFunction = calc_heating_and_current
 
   @property
@@ -146,9 +147,8 @@ class ElectronCyclotronSourceConfig(base.SourceModelBase):
       location
     gaussian_ec_total_power: Gaussian EC total power
   """
-
-  source_name: Literal["electron_cyclotron_source"] = (
-      "electron_cyclotron_source"
+  model_function_name: Literal["calc_heating_and_current"] = (
+      "calc_heating_and_current"
   )
   cd_efficiency: torax_pydantic.TimeVaryingArray = (
       torax_pydantic.ValidatedDefault({0.0: {0.0: 0.2, 1.0: 0.2}})

torax/sources/fusion_heat_source.py

@@ -31,6 +31,12 @@
 from torax.sources import source_profiles
 
 
+# Default value for the model function to be used for the fusion heat
+# source. This is also used as an identifier for the model function in
+# the default source config for Pydantic to "discriminate" against.
+DEFAULT_MODEL_FUNCTION_NAME: str = 'fusion_heat_model_func'
+
+
 def calc_fusion(
     geo: geometry.Geometry,
     core_profiles: state.CoreProfiles,
@@ -162,7 +168,6 @@ class FusionHeatSource(source.Source):
   """Fusion heat source for both ion and electron heat."""
 
   SOURCE_NAME: ClassVar[str] = 'fusion_heat_source'
-  DEFAULT_MODEL_FUNCTION_NAME: ClassVar[str] = 'fusion_heat_model_func'
   model_func: source.SourceProfileFunction = fusion_heat_model_func
 
   @property
@@ -179,7 +184,9 @@ def affected_core_profiles(self) -> tuple[source.AffectedCoreProfile, ...]:
 
 class FusionHeatSourceConfig(base.SourceModelBase):
   """Configuration for the FusionHeatSource."""
-  source_name: Literal['fusion_heat_source'] = 'fusion_heat_source'
+  model_function_name: Literal['fusion_heat_model_func'] = (
+      'fusion_heat_model_func'
+  )
   mode: runtime_params_lib.Mode = runtime_params_lib.Mode.MODEL_BASED
 
   @property

torax/sources/gas_puff_source.py

@@ -29,6 +29,12 @@
 from torax.torax_pydantic import torax_pydantic
 
 
+# Default value for the model function to be used for the gas puff
+# source. This is also used as an identifier for the model function in
+# the default source config for Pydantic to "discriminate" against.
+DEFAULT_MODEL_FUNCTION_NAME: str = 'calc_puff_source'
+
+
 # pylint: disable=invalid-name
 @chex.dataclass(frozen=True)
 class DynamicGasPuffRuntimeParams(runtime_params_lib.DynamicRuntimeParams):
@@ -66,7 +72,6 @@ class GasPuffSource(source.Source):
   """Gas puff source for the ne equation."""
 
   SOURCE_NAME: ClassVar[str] = 'gas_puff_source'
-  DEFAULT_MODEL_FUNCTION_NAME: ClassVar[str] = 'calc_puff_source'
   model_func: source.SourceProfileFunction = calc_puff_source
 
   @property
@@ -87,8 +92,7 @@ class GasPuffSourceConfig(base.SourceModelBase):
       [normalized radial coord]
     S_puff_tot: total gas puff particles/s
   """
-
-  source_name: Literal['gas_puff_source'] = 'gas_puff_source'
+  model_function_name: Literal['calc_puff_source'] = 'calc_puff_source'
   puff_decay_length: torax_pydantic.TimeVaryingScalar = (
       torax_pydantic.ValidatedDefault(0.05)
   )

torax/sources/generic_current_source.py

@@ -31,6 +31,12 @@
 from torax.torax_pydantic import torax_pydantic
 
 
+# Default value for the model function to be used for the generic current
+# source. This is also used as an identifier for the model function in
+# the default source config for Pydantic to "discriminate" against.
+DEFAULT_MODEL_FUNCTION_NAME: str = 'calc_generic_current'
+
+
 # pylint: disable=invalid-name
 @chex.dataclass(frozen=True)
 class DynamicRuntimeParams(runtime_params_lib.DynamicRuntimeParams):
@@ -106,7 +112,6 @@ class GenericCurrentSource(source.Source):
   """A generic current density source profile."""
 
   SOURCE_NAME: ClassVar[str] = 'generic_current_source'
-  DEFAULT_MODEL_FUNCTION_NAME: ClassVar[str] = 'calc_generic_current'
   model_func: source.SourceProfileFunction = calculate_generic_current
 
   @property
@@ -129,8 +134,7 @@ class GenericCurrentSourceConfig(source_base.SourceModelBase):
     use_absolute_current: Toggles if external current is provided absolutely or
       as a fraction of Ip.
   """
-
-  source_name: Literal['generic_current_source'] = 'generic_current_source'
+  model_function_name: Literal['calc_generic_current'] = 'calc_generic_current'
   Iext: torax_pydantic.TimeVaryingScalar = torax_pydantic.ValidatedDefault(3.0)
   fext: torax_pydantic.TimeVaryingScalar = torax_pydantic.ValidatedDefault(0.2)
   wext: torax_pydantic.TimeVaryingScalar = torax_pydantic.ValidatedDefault(0.05)

torax/sources/generic_ion_el_heat_source.py

@@ -29,6 +29,12 @@
 from torax.torax_pydantic import torax_pydantic
 
 
+# Default value for the model function to be used for the electron cyclotron
+# source. This is also used as an identifier for the model function in
+# the default source config for Pydantic to "discriminate" against.
+DEFAULT_MODEL_FUNCTION_NAME: str = 'default_formula'
+
+
 # pylint: disable=invalid-name
 @chex.dataclass(frozen=True)
 class DynamicRuntimeParams(runtime_params_lib.DynamicRuntimeParams):
@@ -103,7 +109,6 @@ class GenericIonElectronHeatSource(source.Source):
   """Generic heat source for both ion and electron heat."""
 
   SOURCE_NAME: ClassVar[str] = 'generic_ion_el_heat_source'
-  DEFAULT_MODEL_FUNCTION_NAME: ClassVar[str] = 'default_formula'
   model_func: source.SourceProfileFunction = default_formula
 
   @property
@@ -128,10 +133,7 @@ class GenericIonElHeatSourceConfig(base.SourceModelBase):
     el_heat_fraction: Electron heating fraction
     absorption_fraction: Fraction of absorbed power
   """
-
-  source_name: Literal['generic_ion_el_heat_source'] = (
-      'generic_ion_el_heat_source'
-  )
+  model_function_name: Literal['default_formula'] = 'default_formula'
   w: torax_pydantic.TimeVaryingScalar = torax_pydantic.ValidatedDefault(0.25)
   rsource: torax_pydantic.TimeVaryingScalar = torax_pydantic.ValidatedDefault(
       0.0

torax/sources/generic_particle_source.py

@@ -28,6 +28,12 @@
 from torax.torax_pydantic import torax_pydantic
 
 
+# Default value for the model function to be used for the generic particle
+# source. This is also used as an identifier for the model function in
+# the default source config for Pydantic to "discriminate" against.
+DEFAULT_MODEL_FUNCTION_NAME: str = 'calc_generic_particle_source'
+
+
 # pylint: disable=invalid-name
 def calc_generic_particle_source(
     unused_static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
@@ -63,7 +69,6 @@ class GenericParticleSource(source.Source):
   """Neutral-beam injection source for the ne equation."""
 
   SOURCE_NAME: ClassVar[str] = 'generic_particle_source'
-  DEFAULT_MODEL_FUNCTION_NAME: ClassVar[str] = 'calc_generic_particle_source'
   model_func: source.SourceProfileFunction = calc_generic_particle_source
 
   @property
@@ -93,8 +98,9 @@ class GenericParticleSourceConfig(base.SourceModelBase):
     mode: Defines how the source values are computed (from a model, from a file,
       etc.)
   """
-
-  source_name: Literal['generic_particle_source'] = 'generic_particle_source'
+  model_function_name: Literal['calc_generic_particle_source'] = (
+      'calc_generic_particle_source'
+  )
   particle_width: torax_pydantic.TimeVaryingScalar = (
       torax_pydantic.ValidatedDefault(0.25)
   )

torax/sources/impurity_radiation_heat_sink/impurity_radiation_constant_fraction.py

@@ -28,8 +28,6 @@
 from torax.sources.impurity_radiation_heat_sink import impurity_radiation_heat_sink
 from torax.torax_pydantic import torax_pydantic
 
-MODEL_FUNCTION_NAME = 'radially_constant_fraction_of_Pin'
-
 
 def radially_constant_fraction_of_Pin(  # pylint: disable=invalid-name
     unused_static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
@@ -101,9 +99,6 @@ class ImpurityRadiationHeatSinkConstantFractionConfig(base.SourceModelBase):
     fraction_of_total_power_density: Fraction of total power density to be
       absorbed by the impurity.
   """
-  source_name: Literal['impurity_radiation_heat_sink'] = (
-      'impurity_radiation_heat_sink'
-  )
   model_function_name: Literal['radially_constant_fraction_of_Pin'] = (
       'radially_constant_fraction_of_Pin'
   )

torax/sources/impurity_radiation_heat_sink/impurity_radiation_mavrin_fit.py

@@ -35,7 +35,10 @@
 from torax.sources.impurity_radiation_heat_sink import impurity_radiation_heat_sink
 
 
-MODEL_FUNCTION_NAME = 'impurity_radiation_mavrin_fit'
+# Default value for the model function to be used for the impurity radiation
+# source. This is also used as an identifier for the model function in
+# the source config for Pydantic to "discriminate" against.
+DEFAULT_MODEL_FUNCTION_NAME: str = 'impurity_radiation_mavrin_fit'
 
 # Polynomial fit coefficients from A. A. Mavrin (2018):
 # Improved fits of coronal radiative cooling rates for high-temperature plasmas,
@@ -224,10 +227,6 @@ class ImpurityRadiationHeatSinkMavrinFitConfig(base.SourceModelBase):
   Attributes:
     radiation_multiplier: Multiplier for the impurity radiation profile.
   """
-
-  source_name: Literal['impurity_radiation_heat_sink'] = (
-      'impurity_radiation_heat_sink'
-  )
   model_function_name: Literal['impurity_radiation_mavrin_fit'] = (
       'impurity_radiation_mavrin_fit'
   )

torax/sources/ion_cyclotron_source.py

@@ -41,6 +41,12 @@
 # Internal import.
 
 
+# Default value for the model function to be used for the ion cyclotron
+# source. This is also used as an identifier for the model function in
+# the default source config for Pydantic to "discriminate" against.
+DEFAULT_MODEL_FUNCTION_NAME: str = 'icrh_model_func'
+
+
 # Environment variable for the TORIC NN model. Used if the model path
 # is not set in the config.
 _MODEL_PATH_ENV_VAR: Final[str] = 'TORIC_NN_MODEL_PATH'
@@ -429,7 +435,6 @@ class IonCyclotronSource(source.Source):
   """Ion cyclotron source with surrogate model."""
 
   SOURCE_NAME: ClassVar[str] = 'ion_cyclotron_source'
-  DEFAULT_MODEL_FUNCTION_NAME: ClassVar[str] = 'icrh_model_func'
 
   @property
   def source_name(self) -> str:
@@ -457,8 +462,7 @@ class IonCyclotronSourceConfig(base.SourceModelBase):
     Ptot: Total heating power [W].
     absorption_fraction: Fraction of absorbed power.
   """
-
-  source_name: Literal['ion_cyclotron_source'] = 'ion_cyclotron_source'
+  model_function_name: Literal['icrh_model_func'] = 'icrh_model_func'
   wall_inner: torax_pydantic.Meter = 1.24
   wall_outer: torax_pydantic.Meter = 2.43
   frequency: torax_pydantic.TimeVaryingScalar = torax_pydantic.ValidatedDefault(