Add dW/dt terms to P_SOL and P_loss calculations.

Introduces a dW_dt_smoothing_time_scale parameter in the numerics config to control an exponential moving average for dW/dt. The smoothed dW/dt is now used in the calculation of P_SOL terms and the energy confinement time (tauE). The raw dW/dt is still available as a separate output, for backwards compatibility and also helping with judging the impact of the smoothing.

Sim tests regenerated due to modified post processed outputs.

P_SOL and P_loss terms with dW/dt verified against RAPTOR.

PiperOrigin-RevId: 852692319

Author: jcitrin

torax/_src/config/numerics.py

@@ -45,6 +45,7 @@ class RuntimeParams:
   resistivity_multiplier: array_typing.FloatScalar
   adaptive_T_source_prefactor: float
   adaptive_n_source_prefactor: float
+  dW_dt_smoothing_time_scale: float
   evolve_ion_heat: bool = dataclasses.field(metadata={'static': True})
   evolve_electron_heat: bool = dataclasses.field(metadata={'static': True})
   evolve_current: bool = dataclasses.field(metadata={'static': True})
@@ -102,6 +103,9 @@ class Numerics(torax_pydantic.BaseModelFrozen):
       temperature internal boundary conditions.
     adaptive_n_source_prefactor: Prefactor for adaptive source term for setting
       density internal boundary conditions.
+    dW_dt_smoothing_time_scale: Time scale [s] for the exponential moving
+      average smoothing of dW/dt terms used in P_SOL and confinement time
+      calculations. If 0.0, no smoothing is applied and raw dW/dt is used.
   """
 
   t_initial: torax_pydantic.Second = 0.0
@@ -125,6 +129,7 @@ class Numerics(torax_pydantic.BaseModelFrozen):
   )
   adaptive_T_source_prefactor: pydantic.PositiveFloat = 2.0e10
   adaptive_n_source_prefactor: pydantic.PositiveFloat = 2.0e8
+  dW_dt_smoothing_time_scale: pydantic.NonNegativeFloat = 0.3
 
   T_minimum_eV: pydantic.PositiveFloat = 5.0
 
@@ -170,6 +175,7 @@ def build_runtime_params(self, t: chex.Numeric) -> RuntimeParams:
         resistivity_multiplier=self.resistivity_multiplier.get_value(t),
         adaptive_T_source_prefactor=self.adaptive_T_source_prefactor,
         adaptive_n_source_prefactor=self.adaptive_n_source_prefactor,
+        dW_dt_smoothing_time_scale=self.dW_dt_smoothing_time_scale,
         evolve_ion_heat=self.evolve_ion_heat,
         evolve_electron_heat=self.evolve_electron_heat,
         evolve_current=self.evolve_current,

torax/_src/output_tools/post_processing.py

@@ -66,12 +66,13 @@ class PostProcessedOutputs:
     FFprime: FF' on the face grid, where F is the toroidal flux function
     psi_norm: Normalized poloidal flux on the face grid [Wb]
     P_SOL_i: Total ion heating power exiting the plasma with all sources:
-      auxiliary heating + ion-electron exchange + fusion [W]
+      auxiliary heating + ion-electron exchange + fusion [W]. Includes smoothed
+      dW/dt correction.
     P_SOL_e: Total electron heating power exiting the plasma with all sources
       and sinks: auxiliary heating + ion-electron exchange + Ohmic + fusion +
-      radiation sinks [W]
+      radiation sinks [W]. Includes smoothed dW/dt correction.
     P_SOL_total: Total heating power exiting the plasma with all sources and
-      sinks
+      sinks. Includes smoothed dW/dt correction.
     P_aux_i: Total auxiliary ion heating power [W]
     P_aux_e: Total auxiliary electron heating power [W]
     P_aux_total: Total auxiliary heating power [W]
@@ -115,7 +116,6 @@ class PostProcessedOutputs:
     T_e_volume_avg: Volume average electron temperature [keV]
     T_i_volume_avg: Volume average ion temperature [keV]
     n_e_volume_avg: Volume average electron density [m^-3]
-    n_e_volume_avg: Volume average electron density [m^-3]
     n_i_volume_avg: Volume average main ion density [m^-3]
     n_e_line_avg: Line averaged electron density [m^-3]
     n_i_line_avg: Line averaged main ion density [m^-3]
@@ -126,7 +126,14 @@ class PostProcessedOutputs:
     q95: q at 95% of the normalized poloidal flux
     W_pol: Total magnetic energy [J]
     li3: Normalized plasma internal inductance, ITER convention [dimensionless]
-    dW_thermal_dt: Time derivative of the total stored thermal energy [W]
+    dW_thermal_dt: Time derivative of the total stored thermal energy [W], raw
+      unsmoothed value.
+    dW_thermal_dt_smoothed: Smoothed time derivative of total stored thermal
+      energy [W].
+    dW_thermal_i_dt_smoothed: Smoothed time derivative of ion stored thermal
+      energy [W].
+    dW_thermal_e_dt_smoothed: Smoothed time derivative of electron stored
+      thermal energy [W].
     q_min: Minimum q value
     rho_q_min: rho_norm at the minimum q
     rho_q_3_2_first: First outermost rho_norm value that intercepts the q=3/2
@@ -235,6 +242,9 @@ class PostProcessedOutputs:
   W_pol: array_typing.FloatScalar
   li3: array_typing.FloatScalar
   dW_thermal_dt: array_typing.FloatScalar
+  dW_thermal_dt_smoothed: array_typing.FloatScalar
+  dW_thermal_i_dt_smoothed: array_typing.FloatScalar
+  dW_thermal_e_dt_smoothed: array_typing.FloatScalar
   rho_q_min: array_typing.FloatScalar
   q_min: array_typing.FloatScalar
   rho_q_3_2_first: array_typing.FloatScalar
@@ -332,6 +342,9 @@ def zeros(cls, geo: geometry.Geometry) -> typing_extensions.Self:
         W_pol=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         li3=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         dW_thermal_dt=jnp.array(0.0, dtype=jax_utils.get_dtype()),
+        dW_thermal_dt_smoothed=jnp.array(0.0, dtype=jax_utils.get_dtype()),
+        dW_thermal_i_dt_smoothed=jnp.array(0.0, dtype=jax_utils.get_dtype()),
+        dW_thermal_e_dt_smoothed=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         rho_q_min=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         q_min=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         rho_q_3_2_first=jnp.array(0.0, dtype=jax_utils.get_dtype()),
@@ -460,12 +473,7 @@ def _calculate_integrated_sources(
   integrated['P_ei_exchange_i'] = math_utils.volume_integration(qei, geo)
   integrated['P_ei_exchange_e'] = -integrated['P_ei_exchange_i']
 
-  # Initialize total electron and ion powers
-  # TODO(b/380848256): P_sol is now correct for stationary state. However,
-  # for generality need to add dWth/dt to the equation (time dependence of
-  # stored energy).
-  integrated['P_SOL_i'] = integrated['P_ei_exchange_i']
-  integrated['P_SOL_e'] = integrated['P_ei_exchange_e']
+  # Initialize total electron and ion auxiliary powers.
   integrated['P_aux_i'] = jnp.array(0.0, dtype=jax_utils.get_dtype())
   integrated['P_aux_e'] = jnp.array(0.0, dtype=jax_utils.get_dtype())
   integrated['P_external_injected'] = jnp.array(
@@ -492,8 +500,7 @@ def _calculate_integrated_sources(
     integrated[f'{value}_total'] = (
         integrated[f'{value}_i'] + integrated[f'{value}_e']
     )
-    integrated['P_SOL_i'] += integrated[f'{value}_i']
-    integrated['P_SOL_e'] += integrated[f'{value}_e']
+
     if key in EXTERNAL_HEATING_SOURCES:
       integrated['P_aux_i'] += integrated[f'{value}_i']
       integrated['P_aux_e'] += integrated[f'{value}_e']
@@ -521,7 +528,6 @@ def _calculate_integrated_sources(
     integrated[f'{value}'] = _get_integrated_source_value(
         core_sources.T_e, key, geo, math_utils.volume_integration
     )
-    integrated['P_SOL_e'] += integrated[f'{value}']
     if key in EXTERNAL_HEATING_SOURCES:
       integrated['P_aux_e'] += integrated[f'{value}']
       integrated['P_external_injected'] += integrated[f'{value}']
@@ -543,7 +549,6 @@ def _calculate_integrated_sources(
     )
     integrated['S_total'] += integrated[f'{value}']
 
-  integrated['P_SOL_total'] = integrated['P_SOL_i'] + integrated['P_SOL_e']
   integrated['P_aux_total'] = integrated['P_aux_i'] + integrated['P_aux_e']
   integrated['P_fusion'] = 5 * integrated['P_alpha_total']
   integrated['P_external_total'] = (
@@ -614,24 +619,100 @@ def make_post_processed_outputs(
       )
   )
 
-  # Thermal energy confinement time is the stored energy divided by the total
-  # input power into the plasma.
+  # Calculate dW/dt.
+  # We perform raw calculation and smoothing inside a conditional block to
+  # prevent division by zero on the first step (when dt=0) and to avoid
+  # large transients (since previous W is initialized to 0).
+  def _calculate_dW_dt_terms():
+    # Raw values
+    dW_i_dt_raw = (
+        W_thermal_ion - previous_post_processed_outputs.W_thermal_i
+    ) / sim_state.dt
+    dW_e_dt_raw = (
+        W_thermal_el - previous_post_processed_outputs.W_thermal_e
+    ) / sim_state.dt
+    dW_total_dt_raw = dW_i_dt_raw + dW_e_dt_raw
+
+    # Smoothing
+    alpha = jax.lax.cond(
+        runtime_params.numerics.dW_dt_smoothing_time_scale > 0.0,
+        lambda: jnp.array(1.0, dtype=jax_utils.get_dtype())
+        - jnp.exp(
+            -sim_state.dt / runtime_params.numerics.dW_dt_smoothing_time_scale
+        ),
+        lambda: jnp.array(1.0, dtype=jax_utils.get_dtype()),
+    )
+
+    def _smooth(new_raw, old_smoothed):
+      return (1.0 - alpha) * old_smoothed + alpha * new_raw
+
+    dW_i_dt_smoothed = _smooth(
+        dW_i_dt_raw,
+        previous_post_processed_outputs.dW_thermal_i_dt_smoothed,
+    )
+    dW_e_dt_smoothed = _smooth(
+        dW_e_dt_raw,
+        previous_post_processed_outputs.dW_thermal_e_dt_smoothed,
+    )
+    dW_total_dt_smoothed = dW_i_dt_smoothed + dW_e_dt_smoothed
+
+    return (
+        dW_total_dt_raw,
+        dW_total_dt_smoothed,
+        dW_i_dt_smoothed,
+        dW_e_dt_smoothed,
+    )
+
+  (
+      dW_thermal_total_dt_raw,
+      dW_thermal_total_dt_smoothed,
+      dW_thermal_i_dt_smoothed,
+      dW_thermal_e_dt_smoothed,
+  ) = jax.lax.cond(
+      previous_post_processed_outputs.first_step,
+      lambda: (0.0, 0.0, 0.0, 0.0),
+      _calculate_dW_dt_terms,
+  )
+
+  # Calculate P_SOL (Power crossing separatrix) = P_sources - P_sinks - dW/dt
+  integrated_sources['P_SOL_i'] = (
+      integrated_sources['P_aux_i']
+      + integrated_sources['P_alpha_i']
+      + integrated_sources['P_ei_exchange_i']
+      - dW_thermal_i_dt_smoothed
+  )
+
+  # Note: P_bremsstrahlung_e, P_cyclotron_e, P_radiation_e are sink terms
+  # (negative) in integrated_sources, so we add them.
+  integrated_sources['P_SOL_e'] = (
+      integrated_sources['P_aux_e']
+      + integrated_sources['P_alpha_e']
+      + integrated_sources['P_ei_exchange_e']
+      + integrated_sources['P_ohmic_e']
+      + integrated_sources['P_bremsstrahlung_e']
+      + integrated_sources['P_cyclotron_e']
+      + integrated_sources['P_radiation_e']
+      - dW_thermal_e_dt_smoothed
+  )
+
+  integrated_sources['P_SOL_total'] = (
+      integrated_sources['P_SOL_i'] + integrated_sources['P_SOL_e']
+  )
 
-  # Ploss term here does not include the reduction of radiated power. Most
-  # analysis of confinement times from databases have not included this term.
+  # Calculate P_loss term used for confinement time calculations.
+  # As per standard definitions, P_loss does not include radiation terms.
   # Therefore highly radiative scenarios can lead to skewed results.
 
-  Ploss = (
+  P_loss = (
       integrated_sources['P_alpha_total']
       + integrated_sources['P_aux_total']
       + integrated_sources['P_ohmic_e']
+      - dW_thermal_total_dt_smoothed
       + constants.CONSTANTS.eps  # Division guard.
   )
 
   def cumulative_values():
-    dW_th_dt = (
-        W_thermal_tot - previous_post_processed_outputs.W_thermal_total
-    ) / sim_state.dt
+
     E_fusion = (
         previous_post_processed_outputs.E_fusion
         + sim_state.dt
@@ -678,7 +759,6 @@ def cumulative_values():
         / 2.0
     )
     return (
-        dW_th_dt,
         E_fusion,
         E_aux_total,
         E_ohmic_e,
@@ -687,37 +767,36 @@ def cumulative_values():
     )
 
   (
-      dW_th_dt,
       E_fusion,
       E_aux_total,
       E_ohmic_e,
       E_external_injected,
       E_external_total,
   ) = jax.lax.cond(
       previous_post_processed_outputs.first_step,
-      lambda: (0.0,) * 6,
+      lambda: (0.0,) * 5,
       cumulative_values,
   )
 
-  tauE = W_thermal_tot / Ploss
+  tau_E = W_thermal_tot / P_loss
 
   tauH89P = scaling_laws.calculate_scaling_law_confinement_time(
-      sim_state.geometry, sim_state.core_profiles, Ploss, 'H89P'
+      sim_state.geometry, sim_state.core_profiles, P_loss, 'H89P'
   )
   tauH98 = scaling_laws.calculate_scaling_law_confinement_time(
-      sim_state.geometry, sim_state.core_profiles, Ploss, 'H98'
+      sim_state.geometry, sim_state.core_profiles, P_loss, 'H98'
   )
   tauH97L = scaling_laws.calculate_scaling_law_confinement_time(
-      sim_state.geometry, sim_state.core_profiles, Ploss, 'H97L'
+      sim_state.geometry, sim_state.core_profiles, P_loss, 'H97L'
   )
   tauH20 = scaling_laws.calculate_scaling_law_confinement_time(
-      sim_state.geometry, sim_state.core_profiles, Ploss, 'H20'
+      sim_state.geometry, sim_state.core_profiles, P_loss, 'H20'
   )
 
-  H89P = tauE / tauH89P
-  H98 = tauE / tauH98
-  H97L = tauE / tauH97L
-  H20 = tauE / tauH20
+  H89P = tau_E / tauH89P
+  H98 = tau_E / tauH98
+  H97L = tau_E / tauH97L
+  H20 = tau_E / tauH20
 
   # Calculate q at 95% of the normalized poloidal flux
   q95 = psi_calculations.calc_q95(psi_norm_face, sim_state.core_profiles.q_face)
@@ -749,12 +828,12 @@ def cumulative_values():
   fgw_n_e_line_avg = formulas.calculate_greenwald_fraction(
       n_e_line_avg, sim_state.core_profiles, sim_state.geometry
   )
-  Wpol = psi_calculations.calc_Wpol(
+  W_pol = psi_calculations.calc_Wpol(
       sim_state.geometry, sim_state.core_profiles.psi
   )
   li3 = psi_calculations.calc_li3(
       sim_state.geometry.R_major,
-      Wpol,
+      W_pol,
       sim_state.core_profiles.Ip_profile_face[-1],
   )
 
@@ -841,7 +920,7 @@ def cumulative_values():
       W_thermal_i=W_thermal_ion,
       W_thermal_e=W_thermal_el,
       W_thermal_total=W_thermal_tot,
-      tau_E=tauE,
+      tau_E=tau_E,
       H89P=H89P,
       H98=H98,
       H97L=H97L,
@@ -868,9 +947,12 @@ def cumulative_values():
       fgw_n_e_volume_avg=fgw_n_e_volume_avg,
       fgw_n_e_line_avg=fgw_n_e_line_avg,
       q95=q95,
-      W_pol=Wpol,
+      W_pol=W_pol,
       li3=li3,
-      dW_thermal_dt=dW_th_dt,
+      dW_thermal_dt=dW_thermal_total_dt_raw,
+      dW_thermal_dt_smoothed=dW_thermal_total_dt_smoothed,
+      dW_thermal_i_dt_smoothed=dW_thermal_i_dt_smoothed,
+      dW_thermal_e_dt_smoothed=dW_thermal_e_dt_smoothed,
       rho_q_min=safety_factor_fit_outputs.rho_q_min,
       q_min=safety_factor_fit_outputs.q_min,
       rho_q_3_2_first=safety_factor_fit_outputs.rho_q_3_2_first,
@@ -924,4 +1006,3 @@ def _convert_j_parallel_face_to_j_toroidal_face(
       j_parallel_to_j_toroidal_factor_cell,
   )
   return j_parallel_to_j_toroidal_factor_face * j_parallel_face
-

torax/_src/physics/scaling_laws.py

@@ -105,15 +105,15 @@ def calculate_plh_scaling_factor(
 def calculate_scaling_law_confinement_time(
     geo: geometry.Geometry,
     core_profiles: state.CoreProfiles,
-    Ploss: jax.Array,
+    P_loss: jax.Array,
     scaling_law: str,
 ) -> jax.Array:
   """Calculates the thermal energy confinement time for a given empirical scaling law.
 
   Args:
     geo: Torus geometry.
     core_profiles: Core plasma profiles.
-    Ploss: Plasma power loss in W.
+    P_loss: Plasma power loss in W.
     scaling_law: Scaling law to use.
 
   Returns:
@@ -184,8 +184,14 @@ def calculate_scaling_law_confinement_time(
 
   params = scaling_params[scaling_law]
 
+  # Ensure P_loss is positive to avoid NaNs in power laws (e.g. x^-0.69).
+  # Physically, scaling laws are derived for steady state (P_loss ~ P_heat).
+  # During transients where dW/dt > P_heat, P_loss can be negative.
+  # We clamp it to a small positive value.
+  P_loss = jnp.maximum(P_loss, 1.0)
+
   scaled_Ip = core_profiles.Ip_profile_face[-1] / 1e6  # convert to MA
-  scaled_Ploss = Ploss / 1e6  # convert to MW
+  scaled_Ploss = P_loss / 1e6  # convert to MW
   B = geo.B_0
   line_avg_n_e = (  # convert to 10^19 m^-3
       math_utils.line_average(core_profiles.n_e.value, geo) / 1e19