Add simple impurity radiation model (P_imp as a fraction of P_tot_in)

torax/examples/prad.py

@@ -0,0 +1,27 @@
+CONFIG = {
+    'runtime_params': {},
+    'geometry': {
+        'geometry_type': 'circular',
+    },
+    'sources': {
+        'j_bootstrap': {},
+        'generic_current_source': {},
+        'generic_particle_source': {},
+        'gas_puff_source': {},
+        'pellet_source': {},
+        'generic_ion_el_heat_source': {},
+        'fusion_heat_source': {},
+        'qei_source': {},
+        'ohmic_heat_source': {},
+        'radiation_heat_sink': {},
+    },
+    'transport': {
+        'transport_model': 'constant',
+    },
+    'stepper': {
+        'stepper_type': 'linear',
+    },
+    'time_step_calculator': {
+        'calculator_type': 'chi',
+    },
+}

torax/plotting/configs/default_plot_config.py

@@ -145,8 +145,8 @@
             suppress_zero_values=True,  # Do not plot all-zero data
         ),
         plotruns_lib.PlotProperties(
-            attrs=('q_brems',),
-            labels=(r'$Q_\mathrm{brems}$',),
+            attrs=('q_brems', 'q_imp'),
+            labels=(r'$Q_\mathrm{brems}$', r'$Q_\mathrm{imp}$'),
             ylabel=r'Heat sink density $[MW~m^{-3}]$',
             suppress_zero_values=True,  # Do not plot all-zero data
         ),

torax/plotting/configs/sources_plot_config.py

@@ -99,8 +99,8 @@
             suppress_zero_values=True,  # Do not plot all-zero data
         ),
         plotruns_lib.PlotProperties(
-            attrs=('q_brems',),
-            labels=(r'$Q_\mathrm{brems}$',),
+            attrs=('q_brems', 'q_imp'),
+            labels=(r'$Q_\mathrm{brems}$', r'$Q_\mathrm{imp}$'),
             ylabel=r'Heat sink density $[MW~m^{-3}]$',
             suppress_zero_values=True,  # Do not plot all-zero data
         ),

torax/plotting/plotruns_lib.py

@@ -111,6 +111,7 @@ class PlotData:
   q_ohmic: np.ndarray  # [MW/m^3]
   q_brems: np.ndarray  # [MW/m^3]
   q_ei: np.ndarray  # [MW/m^3]
+  q_imp: np.ndarray  # [MW/m^3]
   Q_fusion: np.ndarray  # pylint: disable=invalid-name  # Dimensionless
   s_puff: np.ndarray  # [10^20 m^-3 s^-1]
   s_generic: np.ndarray  # [10^20 m^-3 s^-1]
@@ -171,12 +172,14 @@ def _transform_data(ds: xr.Dataset):
         'fusion_heat_source_el': 1e6,  # W/m^3 to MW/m^3
         'ohmic_heat_source': 1e6,  # W/m^3 to MW/m^3
         'bremsstrahlung_heat_sink': 1e6,  # W/m^3 to MW/m^3
+        'impurity_radiation_heat_sink': 1e6,  # W/m^3 to MW/m^3
         'qei_source': 1e6,  # W/m^3 to MW/m^3
         'P_ohmic': 1e6,  # W to MW
         'P_external_tot': 1e6,  # W to MW
         'P_alpha_tot': 1e6,  # W to MW
         'P_brems': 1e6,  # W to MW
         'P_ecrh': 1e6,  # W to MW
+        'P_imp': 1e6,  # W to MW
         'I_ecrh': 1e6,  # A to MA
         'I_generic': 1e6,  # A to MA
     }
@@ -246,6 +249,9 @@ def _transform_data(ds: xr.Dataset):
       q_brems=get_optional_data(
           core_sources_dataset, 'bremsstrahlung_heat_sink', 'cell'
       ),
+      q_imp=get_optional_data(
+          core_sources_dataset, 'impurity_radiation_heat_sink', 'cell'
+      ),
       q_ei=core_sources_dataset['qei_source'].to_numpy(),  # ion heating/sink
       Q_fusion=post_processed_outputs_dataset['Q_fusion'].to_numpy(),  # pylint: disable=invalid-name
       s_puff=get_optional_data(core_sources_dataset, 'gas_puff_source', 'cell'),
@@ -263,7 +269,7 @@ def _transform_data(ds: xr.Dataset):
           - post_processed_outputs_dataset['P_ohmic']
       ).to_numpy(),
       p_alpha=post_processed_outputs_dataset['P_alpha_tot'].to_numpy(),
-      p_sink=post_processed_outputs_dataset['P_brems'].to_numpy(),
+      p_sink=post_processed_outputs_dataset['P_brems'].to_numpy() + post_processed_outputs_dataset['P_imp'].to_numpy(),
       t=time,
   )
 

torax/post_processing.py

@@ -17,8 +17,10 @@
 # In torax/post_processing.py
 
 import dataclasses
+
 import jax
 from jax import numpy as jnp
+
 from torax import array_typing
 from torax import constants
 from torax import geometry
@@ -40,6 +42,7 @@
     'ohmic_heat_source': 'P_ohmic',
     'bremsstrahlung_heat_sink': 'P_brems',
     'electron_cyclotron_source': 'P_ecrh',
+    'impurity_radiation_heat_sink': 'P_imp',
 }
 EXTERNAL_HEATING_SOURCES = [
     'generic_ion_el_heat_source',

torax/sources/impurity_radiation_heat_sink.py

@@ -0,0 +1,120 @@
+"""Basic impurity radiation heat sink for electron heat equation.."""
+
+import dataclasses
+
+import chex
+import jax
+import jax.numpy as jnp
+
+from torax import array_typing
+from torax import geometry
+from torax import math_utils
+from torax import state
+from torax.config import runtime_params_slice
+from torax.sources import runtime_params as runtime_params_lib
+from torax.sources import source as source_lib
+from torax.sources import source_models as source_models_lib
+
+SOURCE_NAME = "impurity_radiation_heat_sink"
+
+
+def _radially_constant_fraction_of_Pin(
+    dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
+    dynamic_source_runtime_params: runtime_params_lib.DynamicRuntimeParams,
+    geo: geometry.Geometry,
+    core_profiles: state.CoreProfiles,
+    source_models: source_models_lib.SourceModels,
+) -> jax.Array:
+    """Model function for radiation heat sink from impurities.
+
+    This model represents a sink in the temp_el equation, whose value is a fixed
+    % of the total heating power input."""
+    # Based on source_models.sum_sources_temp_el and source_models.calc_and_sum_sources_psi,
+    # but only summing over heating *input* sources (Pohm + Paux + Palpha + ...)
+    # and summing over *both* ion and electron heating
+
+    def get_heat_source_profile(source_name: str, source: source_lib.Source) -> jax.Array:
+        # TODO: Currently this recomputes the profile for each source, which is inefficient
+        # (and will be a problem if sources are slow/non-jittable)
+        # A similar TODO is noted in source_models.calc_and_sum_sources_psi
+        profile = source.get_value(
+            dynamic_runtime_params_slice,
+            dynamic_runtime_params_slice.sources[source_name],
+            geo,
+            core_profiles,
+        )
+        return source.get_source_profile_for_affected_core_profile(
+            profile, source_lib.AffectedCoreProfile.TEMP_EL.value, geo
+        ) + source.get_source_profile_for_affected_core_profile(
+            profile, source_lib.AffectedCoreProfile.TEMP_ION.value, geo
+        )
+
+    # Calculate the total power input to the heat equations
+    heat_sources_and_sinks = source_models.temp_el_sources | source_models.temp_ion_sources
+    heat_sources = {k: v for k, v in heat_sources_and_sinks.items() if not "sink" in k}
+    source_profiles = jax.tree.map(
+        get_heat_source_profile,
+        list(heat_sources.keys()),
+        list(heat_sources.values()),
+    )
+    Qtot_in = jnp.sum(jnp.stack(source_profiles), axis=0)
+    Ptot_in = math_utils.cell_integration(Qtot_in * geo.vpr, geo)
+    Vtot = geo.volume_face[-1]
+
+    # Calculate the heat sink as a fraction of the total power input
+    return (
+        -dynamic_source_runtime_params.fraction_of_total_power_density
+        * Ptot_in / Vtot
+        * jnp.ones_like(geo.rho)
+    )
+
+
+@dataclasses.dataclass(kw_only=True)
+class RuntimeParams(runtime_params_lib.RuntimeParams):
+    fraction_of_total_power_density: runtime_params_lib.TimeInterpolatedInput = 0.1
+    mode: runtime_params_lib.Mode = runtime_params_lib.Mode.MODEL_BASED
+
+    def make_provider(
+        self,
+        torax_mesh: geometry.Grid1D | None = None,
+    ) -> "RuntimeParamsProvider":
+        return RuntimeParamsProvider(**self.get_provider_kwargs(torax_mesh))
+
+
+@chex.dataclass
+class RuntimeParamsProvider(runtime_params_lib.RuntimeParamsProvider):
+    """Provides runtime parameters for a given time and geometry."""
+
+    runtime_params_config: RuntimeParams
+
+    def build_dynamic_params(
+        self,
+        t: chex.Numeric,
+    ) -> "DynamicRuntimeParams":
+        return DynamicRuntimeParams(**self.get_dynamic_params_kwargs(t))
+
+
+@chex.dataclass(frozen=True)
+class DynamicRuntimeParams(runtime_params_lib.DynamicRuntimeParams):
+    fraction_of_total_power_density: array_typing.ScalarFloat
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True, eq=True)
+class ImpurityRadiationHeatSink(source_lib.Source):
+    """Impurity radiation heat sink for electron heat equation."""
+
+    source_models: source_models_lib.SourceModels
+    model_func: source_lib.SourceProfileFunction = _radially_constant_fraction_of_Pin
+
+    @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_lib.AffectedCoreProfile, ...]:
+        return (source_lib.AffectedCoreProfile.TEMP_EL,)

torax/sources/register_source.py

@@ -45,6 +45,7 @@ class to build, the runtime associated with that source and (optionally) the
 from torax.sources import fusion_heat_source
 from torax.sources import generic_current_source
 from torax.sources import generic_ion_el_heat_source as ion_el_heat
+from torax.sources import impurity_radiation_heat_sink
 from torax.sources import ion_cyclotron_source
 from torax.sources import ohmic_heat_source
 from torax.sources import qei_source
@@ -146,6 +147,11 @@ def _register_new_source(
         source_class=ion_cyclotron_source.IonCyclotronSource,
         default_runtime_params_class=ion_cyclotron_source.RuntimeParams,
     ),
+    impurity_radiation_heat_sink.SOURCE_NAME: _register_new_source(
+      source_class=impurity_radiation_heat_sink.ImpurityRadiationHeatSink,
+      default_runtime_params_class=impurity_radiation_heat_sink.RuntimeParams,
+      links_back=True,
+    )
 }
 
 
@@ -155,4 +161,3 @@ def get_registered_source(source_name: str) -> RegisteredSource:
     return _REGISTERED_SOURCES[source_name]
   else:
     raise RuntimeError(f'Source:{source_name} has not been registered.')
-

torax/sources/tests/bootstrap_current_source.py

@@ -39,6 +39,15 @@ def setUpClass(cls):
         expected_affected_core_profiles=(source_lib.AffectedCoreProfile.PSI,),
     )
 
+  def test_expected_mesh_states(self):
+    # This function is reimplemented here as BootstrapCurrentSource does not
+    # appear in source_models, which the parent class uses to build the source
+    source = bootstrap_current_source.BootstrapCurrentSource()
+    self.assertSameElements(
+        source.affected_core_profiles,
+        self._expected_affected_core_profiles,
+    )
+
   def test_extraction_of_relevant_profile_from_output(self):
     """Tests that the relevant profile is extracted from the output."""
     source = bootstrap_current_source.BootstrapCurrentSource()

torax/sources/tests/electron_cyclotron_source.py

@@ -143,7 +143,15 @@ def test_invalid_source_types_raise_errors(self):
   def test_extraction_of_relevant_profile_from_output(self):
     """Tests that the relevant profile is extracted from the output."""
     geo = geometry.build_circular_geometry()
-    source = self._source_class()
+    # pylint: disable=missing-kwoa
+    source_builder = self._source_class_builder()  # pytype: disable=missing-parameter
+    # pylint: enable=missing-kwoa
+    source_models_builder = source_models_lib.SourceModelsBuilder(
+        {'foo': source_builder},
+    )
+    source_models = source_models_builder()
+    source = source_models.sources['foo']
+    self.assertIsInstance(source, source_lib.Source)
     cell = source_lib.ProfileType.CELL.get_profile_shape(geo)
     fake_profile = jnp.stack((jnp.ones(cell), 2 * jnp.ones(cell)))
     # Check TEMP_EL and PSI are modified