Merge pull request #572 from google-deepmind:theo-brown/add-prad

PiperOrigin-RevId: 704205875

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,8 @@ 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,169 @@
+# 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.
+
+
+"""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
+
+
+def _radially_constant_fraction_of_Pin(  # pylint: disable=invalid-name
+    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.
+
+  Args:
+    static_runtime_params_slice: Static runtime parameters.
+    static_source_runtime_params: Static source runtime parameters.
+    dynamic_runtime_params_slice: Dynamic runtime parameters.
+    dynamic_source_runtime_params: Dynamic source runtime parameters.
+    geo: Geometry object.
+    core_profiles: Core profiles object.
+    source_models: Source models object.
+
+  Returns:
+    The heat sink profile.
+  """
+  del (static_source_runtime_params,)  # Unused
+  assert isinstance(dynamic_source_runtime_params, DynamicRuntimeParams)
+
+  # 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 + electron heating
+
+  def get_heat_source_profile(
+      source_name: str, source: source_lib.Source
+  ) -> jax.Array:
+    # TODO(b/381543891): Currently this recomputes the profile for each source,
+    # which is inefficient. Refactor to avoid this.
+    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 "sink" not 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_NAME = "impurity_radiation_heat_sink"
+  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
@@ -152,6 +153,11 @@ def _register_new_source(
         source_class=ion_cyclotron_source.IonCyclotronSource,
         default_runtime_params_class=ion_cyclotron_source.RuntimeParams,
     ),
+    impurity_radiation_heat_sink.ImpurityRadiationHeatSink.SOURCE_NAME: _register_new_source(
+        source_class=impurity_radiation_heat_sink.ImpurityRadiationHeatSink,
+        default_runtime_params_class=impurity_radiation_heat_sink.RuntimeParams,
+        links_back=True,
+    ),
 }