Add a data structure for the energy state info to core profiles

And update/initialise this as expected

PiperOrigin-RevId: 876290487

Author: tamaranorman

torax/_src/core_profiles/initialization.py

@@ -32,6 +32,7 @@
 from torax._src.geometry import standard_geometry
 from torax._src.neoclassical import neoclassical_models as neoclassical_models_lib
 from torax._src.neoclassical.bootstrap_current import base as bootstrap_current_base
+from torax._src.physics import formulas
 from torax._src.physics import psi_calculations
 from torax._src.sources import source_models as source_models_lib
 from torax._src.sources import source_profile_builders
@@ -129,6 +130,13 @@ def initial_core_profiles(
       charge_state_info_face=ions.charge_state_info_face,
   )
 
+  # TODO(b/398816463): Clean up energy state initialization as part of V2
+  # core profiles refactor.
+  core_profiles = dataclasses.replace(
+      core_profiles,
+      internal_plasma_energy=_initialise_internal_energy(core_profiles, geo),
+  )
+
   return _init_psi_and_psi_derived(
       runtime_params,
       geo,
@@ -138,6 +146,30 @@ def initial_core_profiles(
   )
 
 
+def _initialise_internal_energy(
+    core_profiles: state.CoreProfiles,
+    geo: geometry.Geometry,
+) -> state.PlasmaInternalEnergy:
+  """Initializes the energy stored in the plasma."""
+  W_thermal_e, W_thermal_i, W_thermal_total = (
+      formulas.calculate_stored_thermal_energy(
+          core_profiles.pressure_thermal_e,
+          core_profiles.pressure_thermal_i,
+          core_profiles.pressure_thermal_total,
+          geo,
+      )
+  )
+  return state.PlasmaInternalEnergy(
+      W_thermal_i=W_thermal_i,
+      W_thermal_e=W_thermal_e,
+      W_thermal_total=W_thermal_total,
+      dW_thermal_i_dt=jnp.array(0.0, dtype=jax_utils.get_dtype()),
+      dW_thermal_e_dt=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()),
+  )
+
+
 def update_psi_from_j(
     Ip: array_typing.FloatScalar,
     geo: geometry.Geometry,

torax/_src/core_profiles/updaters.py

@@ -31,14 +31,17 @@
 import dataclasses
 
 import jax
+import jax.numpy as jnp
 from torax._src import array_typing
+from torax._src import jax_utils
 from torax._src import state
 from torax._src.config import runtime_params as runtime_params_lib
 from torax._src.core_profiles import convertors
 from torax._src.core_profiles import getters
 from torax._src.fvm import cell_variable
 from torax._src.geometry import geometry
 from torax._src.neoclassical import neoclassical_models as neoclassical_models_lib
+from torax._src.physics import formulas
 from torax._src.physics import psi_calculations
 from torax._src.sources import source_models as source_models_lib
 from torax._src.sources import source_profile_builders
@@ -68,15 +71,12 @@ def update_core_profiles_during_step(
     x_new: The new values of the evolving variables.
     runtime_params: The runtime params slice.
     geo: Magnetic geometry.
-    core_profiles: The old set of core plasma profiles.
+    core_profiles: The old set of core plasma profiles for this timestep.
     prev_core_profiles: Core plasma profiles from the previous timestep if
       available, used to update the energy state.
     dt: The size of the last timestep, used to update the energy state.
     evolving_names: The names of the evolving variables.
   """
-  # Currently unused but will be used to update the energy state soon
-  del prev_core_profiles, dt
-
   updated_core_profiles = convertors.solver_x_tuple_to_core_profiles(
       x_new, evolving_names, core_profiles
   )
@@ -88,7 +88,7 @@ def update_core_profiles_during_step(
       updated_core_profiles.T_e,
   )
 
-  return dataclasses.replace(
+  updated_core_profiles = dataclasses.replace(
       updated_core_profiles,
       n_i=ions.n_i,
       n_impurity=ions.n_impurity,
@@ -109,6 +109,24 @@ def update_core_profiles_during_step(
       charge_state_info_face=ions.charge_state_info_face,
   )
 
+  if prev_core_profiles is not None:
+    if dt is None:
+      raise ValueError('dt must be provided when updating the energy state.')
+    energy_state = _update_energy_state(
+        runtime_params,
+        geo,
+        updated_core_profiles,
+        prev_core_profiles.internal_plasma_energy,
+        dt,
+    )
+  else:
+    energy_state = core_profiles.internal_plasma_energy
+
+  return dataclasses.replace(
+      updated_core_profiles,
+      internal_plasma_energy=energy_state,
+  )
+
 
 def update_core_and_source_profiles_after_step(
     dt: array_typing.FloatScalar,
@@ -208,6 +226,13 @@ def update_core_and_source_profiles_after_step(
       charge_state_info=ions.charge_state_info,
       charge_state_info_face=ions.charge_state_info_face,
   )
+  energy_state = _update_energy_state(
+      runtime_params_t_plus_dt,
+      geo,
+      intermediate_core_profiles,
+      core_profiles_t.internal_plasma_energy,
+      dt,
+  )
 
   conductivity = neoclassical_models.conductivity.calculate_conductivity(
       geo, intermediate_core_profiles
@@ -217,6 +242,7 @@ def update_core_and_source_profiles_after_step(
       intermediate_core_profiles,
       sigma=conductivity.sigma,
       sigma_face=conductivity.sigma_face,
+      internal_plasma_energy=energy_state,
   )
 
   # build_source_profiles calculates the union with explicit + implicit
@@ -316,3 +342,55 @@ def provide_core_profiles_t_plus_dt(
       toroidal_angular_velocity=toroidal_angular_velocity,
   )
   return core_profiles_t_plus_dt
+
+
+def _update_energy_state(
+    runtime_params: runtime_params_lib.RuntimeParams,
+    geo: geometry.Geometry,
+    core_profiles: state.CoreProfiles,
+    prev_energy_state: state.PlasmaInternalEnergy,
+    dt: array_typing.FloatScalar,
+) -> state.PlasmaInternalEnergy:
+  """Updates the energy state."""
+  W_thermal_e, W_thermal_i, W_thermal_total = (
+      formulas.calculate_stored_thermal_energy(
+          core_profiles.pressure_thermal_e,
+          core_profiles.pressure_thermal_i,
+          core_profiles.pressure_thermal_total,
+          geo,
+      )
+  )
+  dW_i_dt_raw = (W_thermal_i - prev_energy_state.W_thermal_i) / dt
+  dW_e_dt_raw = (W_thermal_e - prev_energy_state.W_thermal_e) / dt
+
+  exponential_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(-dt / runtime_params.numerics.dW_dt_smoothing_time_scale),
+      lambda: jnp.array(1.0, dtype=jax_utils.get_dtype()),
+  )
+  dW_i_dt_smoothed = _exponential_smoothing(
+      dW_i_dt_raw,
+      prev_energy_state.dW_thermal_i_dt_smoothed,
+      exponential_smoothing_alpha,
+  )
+  dW_e_dt_smoothed = _exponential_smoothing(
+      dW_e_dt_raw,
+      prev_energy_state.dW_thermal_e_dt_smoothed,
+      exponential_smoothing_alpha,
+  )
+
+  return state.PlasmaInternalEnergy(
+      W_thermal_i=W_thermal_i,
+      W_thermal_e=W_thermal_e,
+      W_thermal_total=W_thermal_total,
+      dW_thermal_i_dt=dW_i_dt_raw,
+      dW_thermal_e_dt=dW_e_dt_raw,
+      dW_thermal_i_dt_smoothed=dW_i_dt_smoothed,
+      dW_thermal_e_dt_smoothed=dW_e_dt_smoothed,
+  )
+
+
+def _exponential_smoothing(new_raw, old_smoothed, alpha):
+  """Exponential moving average (EMA)."""
+  return (1.0 - alpha) * old_smoothed + alpha * new_raw

torax/_src/orchestration/initial_state.py

@@ -259,9 +259,20 @@ def get_initial_state_and_post_processed_outputs_from_file(
           'dW_thermal_e_dt_smoothed'
       ].to_numpy(),
   )
+  energy_state = initial_state.core_profiles.internal_plasma_energy
+  energy_state = dataclasses.replace(
+      energy_state,
+      dW_thermal_i_dt_smoothed=scalars_dataset.data_vars[
+          'dW_thermal_i_dt_smoothed'
+      ].to_numpy(),
+      dW_thermal_e_dt_smoothed=scalars_dataset.data_vars[
+          'dW_thermal_e_dt_smoothed'
+      ].to_numpy(),
+  )
   core_profiles = dataclasses.replace(
       initial_state.core_profiles,
       v_loop_lcfs=scalars_dataset.v_loop_lcfs.values,
+      internal_plasma_energy=energy_state,
   )
   numerics_dataset = data_tree.children[output.NUMERICS].dataset
   numerics_dataset = numerics_dataset.squeeze()

torax/_src/output_tools/output.py

@@ -581,6 +581,11 @@ def _save_core_profiles(
       if attr_name == "main_ion_fractions":
         continue
 
+      # Skip internal_plasma_energy as it is in post_processed_outputs.
+      # TODO(b/434175938): Remove once we move to V2.
+      if attr_name == "internal_plasma_energy":
+        continue
+
       attr_value = getattr(stacked_core_profiles, attr_name)
 
       output_key = output_name_map.get(attr_name, attr_name)

torax/_src/state.py

@@ -32,6 +32,41 @@
 
 
 # pylint: disable=invalid-name
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True)
+class PlasmaInternalEnergy:
+  """Internal energy stored in the plasma and its time derivatives.
+
+  Attributes:
+    W_thermal_i: Ion thermal stored energy [J]
+    W_thermal_e: Electron thermal stored energy [J]
+    W_thermal_total: Total thermal stored energy [J]
+    dW_thermal_i_dt: Time derivative of ion thermal stored energy [W]
+    dW_thermal_e_dt: Time derivative of electron thermal stored energy [W]
+    dW_thermal_i_dt_smoothed: Smoothed time derivative of ion thermal stored
+      energy [W]
+    dW_thermal_e_dt_smoothed: Smoothed time derivative of electron thermal
+      stored energy [W]
+  """
+  W_thermal_i: array_typing.FloatScalar
+  W_thermal_e: array_typing.FloatScalar
+  W_thermal_total: array_typing.FloatScalar
+  dW_thermal_i_dt: array_typing.FloatScalar
+  dW_thermal_e_dt: array_typing.FloatScalar
+  dW_thermal_i_dt_smoothed: array_typing.FloatScalar
+  dW_thermal_e_dt_smoothed: array_typing.FloatScalar
+
+  @property
+  def dW_thermal_dt(self) -> array_typing.FloatScalar:
+    """Total thermal stored energy time derivative [W], raw unsmoothed."""
+    return self.dW_thermal_i_dt + self.dW_thermal_e_dt
+
+  @property
+  def dW_thermal_dt_smoothed(self) -> array_typing.FloatScalar:
+    """Smoothed total thermal stored energy time derivative [W]."""
+    return self.dW_thermal_i_dt_smoothed + self.dW_thermal_e_dt_smoothed
+
+
 @jax.tree_util.register_dataclass
 @dataclasses.dataclass(frozen=True, eq=False)
 class CoreProfiles:
@@ -78,6 +113,8 @@ class CoreProfiles:
         state information. See `charge_states.ChargeStateInfo`. Cell grid.
       charge_state_info_face: Container with averaged and per-species ion charge
         state information. See `charge_states.ChargeStateInfo`. Face grid.
+      internal_plasma_energy: Container with energy variables. See
+        `PlasmaInternalEnergy`.
   """
 
   T_i: cell_variable.CellVariable
@@ -109,6 +146,7 @@ class CoreProfiles:
   toroidal_angular_velocity: cell_variable.CellVariable
   charge_state_info: charge_states.ChargeStateInfo
   charge_state_info_face: charge_states.ChargeStateInfo
+  internal_plasma_energy: PlasmaInternalEnergy | None = None
 
   @functools.cached_property
   def impurity_density_scaling(self) -> jax.Array: