Add Prad to post_processing and plots

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
+        '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, '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

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

torax/sources/radiation_heat_sink.py

@@ -70,7 +70,7 @@ def get_temp_el_profile(source_name: str, source: source_lib.Source) -> jax.Arra
 
 @dataclasses.dataclass(kw_only=True)
 class RuntimeParams(runtime_params_lib.RuntimeParams):
-    fraction_of_total_power_density: runtime_params_lib.TimeInterpolatedInput = 0.
+    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(

torax/state.py

@@ -350,6 +350,7 @@ class PostProcessedOutputs:
   P_ohmic: array_typing.ScalarFloat
   P_brems: array_typing.ScalarFloat
   P_ecrh: array_typing.ScalarFloat
+  P_rad: array_typing.ScalarFloat
   I_ecrh: array_typing.ScalarFloat
   I_generic: array_typing.ScalarFloat
   Q_fusion: array_typing.ScalarFloat
@@ -391,6 +392,7 @@ def zeros(cls, geo: geometry.Geometry) -> PostProcessedOutputs:
         P_ohmic=jnp.array(0.0),
         P_brems=jnp.array(0.0),
         P_ecrh=jnp.array(0.0),
+        P_rad=jnp.array(0.0),
         I_ecrh=jnp.array(0.0),
         I_generic=jnp.array(0.0),
         Q_fusion=jnp.array(0.0),