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

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_rad'),
+            labels=(r'$Q_\mathrm{brems}$', r'$Q_\mathrm{rad}$'),
             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_rad'),
+            labels=(r'$Q_\mathrm{brems}$', r'$Q_\mathrm{rad}$'),
             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_rad: 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_rad': 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_rad=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_rad'].to_numpy(),
       t=time,
   )
 

torax/post_processing.py

@@ -40,6 +40,7 @@
     'ohmic_heat_source': 'P_ohmic',
     'bremsstrahlung_heat_sink': 'P_brems',
     'electron_cyclotron_source': 'P_ecrh',
+    'impurity_radiation_heat_sink': 'P_rad',
 }
 EXTERNAL_HEATING_SOURCES = [
     'generic_ion_el_heat_source',

torax/sources/impurity_radiation_heat_sink.py

@@ -0,0 +1,128 @@
+"""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(
+    static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
+    static_source_runtime_params: runtime_params_lib.StaticRuntimeParams,
+    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."""
+    del (
+        static_source_runtime_params,
+    ) # Unused
+
+    # 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,
+            dynamic_source_runtime_params=dynamic_runtime_params_slice.sources[source_name],
+            static_runtime_params_slice=static_runtime_params_slice,
+            static_source_runtime_params=static_runtime_params_slice.sources[source_name],
+            geo=geo,
+            core_profiles=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.')
-