Extend return value of charge_states.get_average_charge_state.

Now outputs a dataclass including the previous average charge state (now called "Z_mixture", as well as various intermediate calculated quantities.

This is useful for the upcoming extended impurity API, as well as reducing some code duplication elsewhere, like in the Mavrin radiation calculation.

PiperOrigin-RevId: 794141009

Author: jcitrin

torax/_src/config/tests/plasma_composition_test.py

@@ -223,7 +223,7 @@ def test_ion_mixture_averaging(self, species, time, expected_Z, expected_A):
         ion_symbols=tuple(species.keys()),  # pytype: disable=attribute-error
         ion_mixture=dynamic_mixture,
         T_e=np.array(10.0),  # Ensure that all ions in test are fully ionized
-    )
+    ).Z_mixture
     np.testing.assert_allclose(calculated_Z, expected_Z)
     np.testing.assert_allclose(dynamic_mixture.avg_A, expected_A)
 
@@ -248,7 +248,7 @@ def test_ion_mixture_override(self, Z_override, A_override, Z, A):
         ion_symbols=tuple(mixture.species.keys()),
         ion_mixture=dynamic_mixture,
         T_e=np.array(1.0),  # arbitrary temperature, won't be used for D
-    )
+    ).Z_mixture
     Z_expected = Z if Z_override is None else Z_override
     A_expected = A if A_override is None else A_override
     np.testing.assert_allclose(calculated_Z, Z_expected)

torax/_src/core_profiles/getters.py

@@ -293,23 +293,23 @@ def _get_charge_states(
       ion_symbols=static_runtime_params_slice.main_ion_names,
       ion_mixture=dynamic_runtime_params_slice.plasma_composition.main_ion,
       T_e=T_e.value,
-  )
+  ).Z_mixture
   Z_i_face = charge_states.get_average_charge_state(
       ion_symbols=static_runtime_params_slice.main_ion_names,
       ion_mixture=dynamic_runtime_params_slice.plasma_composition.main_ion,
       T_e=T_e.face_value(),
-  )
+  ).Z_mixture
 
   Z_impurity = charge_states.get_average_charge_state(
       ion_symbols=static_runtime_params_slice.impurity_names,
       ion_mixture=dynamic_runtime_params_slice.plasma_composition.impurity,
       T_e=T_e.value,
-  )
+  ).Z_mixture
   Z_impurity_face = charge_states.get_average_charge_state(
       ion_symbols=static_runtime_params_slice.impurity_names,
       ion_mixture=dynamic_runtime_params_slice.plasma_composition.impurity,
       T_e=T_e.face_value(),
-  )
+  ).Z_mixture
 
   return Z_i, Z_i_face, Z_impurity, Z_impurity_face
 

torax/_src/core_profiles/updaters.py

@@ -383,12 +383,12 @@ def compute_boundary_conditions_for_t_plus_dt(
       static_runtime_params_slice.main_ion_names,
       ion_mixture=dynamic_runtime_params_slice_t_plus_dt.plasma_composition.main_ion,
       T_e=profile_conditions_t_plus_dt.T_e_right_bc,
-  )
+  ).Z_mixture
   Z_impurity_edge = charge_states.get_average_charge_state(
       static_runtime_params_slice.impurity_names,
       ion_mixture=dynamic_runtime_params_slice_t_plus_dt.plasma_composition.impurity,
       T_e=profile_conditions_t_plus_dt.T_e_right_bc,
-  )
+  ).Z_mixture
 
   dilution_factor_edge = formulas.calculate_main_ion_dilution_factor(
       Z_i_edge,

torax/_src/physics/charge_states.py

@@ -14,15 +14,19 @@
 
 """Routines for calculating impurity charge states."""
 
+import dataclasses
 from typing import Final, Mapping, Sequence
 
 import immutabledict
+import jax
 from jax import numpy as jnp
 import numpy as np
 from torax._src import array_typing
 from torax._src import constants
 from torax._src.config import plasma_composition
 
+# pylint: disable=invalid-name
+
 # Polynomial fit coefficients from A. A. Mavrin (2018):
 # Improved fits of coronal radiative cooling rates for high-temperature plasmas,
 # Radiation Effects and Defects in Solids, 173:5-6, 388-398,
@@ -86,6 +90,33 @@
 )
 
 
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True)
+class ChargeStateInfo:
+  """Container for average charge state calculations.
+
+  Attributes:
+    Z_avg: Average charge of the mixture, weighted by ion fractions. <Z> =
+      sum(fraction_i * Z_i).
+    Z2_avg: Average squared charge of the mixture, weighted by ion fractions.
+      <Z^2> = sum(fraction_i * Z_i^2).
+    Z_per_species: Charge state for each individual ion species in the mixture.
+      For impurities, this is the outcome of a temperature dependent charge
+      state calculation.
+    Z_mixture: Effective charge of the mixture, defined as <Z^2> / <Z>. This is
+      the charge used in quasineutrality calculations when treating the mixture
+      as a single effective species.
+  """
+
+  Z_avg: array_typing.ArrayFloat
+  Z2_avg: array_typing.ArrayFloat
+  Z_per_species: array_typing.ArrayFloat
+
+  @property
+  def Z_mixture(self) -> array_typing.ArrayFloat:
+    return self.Z2_avg / self.Z_avg
+
+
 # pylint: disable=invalid-name
 def calculate_average_charge_state_single_species(
     T_e: array_typing.ArrayFloat,
@@ -135,7 +166,7 @@ def get_average_charge_state(
     ion_symbols: Sequence[str],
     ion_mixture: plasma_composition.DynamicIonMixture,
     T_e: array_typing.ArrayFloat,
-) -> array_typing.ArrayFloat:
+) -> ChargeStateInfo:
   """Calculates or prescribes average impurity charge state profile (JAX-compatible).
 
   Equations for quasineutrality and Zeff are the following:
@@ -171,18 +202,37 @@ def get_average_charge_state(
       face grid, or a single scalar.
 
   Returns:
-    avg_Z: Average charge state profile [amu].
-      The shape of avg_Z is the same as T_e.
+    AverageChargeState: dataclass with average charge state info.
   """
 
   if ion_mixture.Z_override is not None:
-    return jnp.ones_like(T_e) * ion_mixture.Z_override
-
-  avg_Z = jnp.zeros_like(T_e)
-  avg_Z2 = jnp.zeros_like(T_e)
-  for ion_symbol, fraction in zip(ion_symbols, ion_mixture.fractions):
-    Z_species = calculate_average_charge_state_single_species(T_e, ion_symbol)
-    avg_Z += fraction * Z_species
-    avg_Z2 += fraction * Z_species**2
+    override_val = jnp.ones_like(T_e) * ion_mixture.Z_override
+    return ChargeStateInfo(
+        Z_avg=override_val,
+        Z2_avg=override_val**2,
+        Z_per_species=jnp.stack([override_val for _ in ion_symbols]),
+    )
 
-  return avg_Z2 / avg_Z
+  Z_per_species = jnp.stack([
+      calculate_average_charge_state_single_species(T_e, ion_symbol)
+      for ion_symbol in ion_symbols
+  ])
+
+  # ion_mixture.fractions has shape (n_species,).
+  # Z_per_species has shape (n_species,) if T_e is a scalar, or
+  # (n_species, n_grid) if T_e is an array.
+  # We need to broadcast fractions for element-wise multiplication.
+  # Reshape fractions to be broadcastable with Z_per_species.
+  fractions_reshaped = jnp.reshape(
+      ion_mixture.fractions,
+      ion_mixture.fractions.shape + (1,) * (Z_per_species.ndim - 1),
+  )
+
+  Z_avg = jnp.sum(fractions_reshaped * Z_per_species, axis=0)
+  Z2_avg = jnp.sum(fractions_reshaped * Z_per_species**2, axis=0)
+
+  return ChargeStateInfo(
+      Z_avg=Z_avg,
+      Z2_avg=Z2_avg,
+      Z_per_species=Z_per_species,
+  )

torax/_src/physics/tests/charge_states_test.py

@@ -179,7 +179,7 @@ def test_get_average_charge_state(
     )
     Z_calculated = charge_states.get_average_charge_state(
         ion_symbols, ion_mixture, T_e
-    )
+    ).Z_mixture
 
     np.testing.assert_allclose(Z_calculated, expected_Z, rtol=1e-5)
 
@@ -196,7 +196,7 @@ def test_Z_override_in_get_average_charge_state(self):
     )
     Z_calculated = charge_states.get_average_charge_state(
         ion_symbols, ion_mixture, T_e
-    )
+    ).Z_mixture
     np.testing.assert_allclose(Z_calculated, Z_override)
 
 

torax/_src/sources/impurity_radiation_heat_sink/impurity_radiation_mavrin_fit.py

@@ -225,15 +225,13 @@ def impurity_radiation_mavrin_fit(
   # impurity density, not the effective one.
 
   # ion_symbols is a static argument so can use the for loop under jit
-  Z_per_species = jnp.stack([
-      charge_states.calculate_average_charge_state_single_species(
-          core_profiles.T_e.value, ion_symbol
-      )
-      for ion_symbol in ion_symbols
-  ])
-
-  avg_Z = jnp.sum(ion_mixture.fractions[:, jnp.newaxis] * Z_per_species, axis=0)
-  impurity_density_scaling = core_profiles.Z_impurity / avg_Z
+  charge_state_info = charge_states.get_average_charge_state(
+      ion_symbols=ion_symbols,
+      ion_mixture=ion_mixture,
+      T_e=core_profiles.T_e.value,
+  )
+  Z_avg = charge_state_info.Z_avg
+  impurity_density_scaling = core_profiles.Z_impurity / Z_avg
 
   dynamic_source_runtime_params = dynamic_runtime_params_slice.sources[
       source_name