From 4c42cee8bcd9c538cbb36d842b201475b18f86e3 Mon Sep 17 00:00:00 2001
From: dylan tirandaz <93934418+dylantirandaz@users.noreply.github.com>
Date: Mon, 29 Jun 2026 14:42:41 -0500
Subject: [PATCH] Add tau_ei to TORAX outputs

---
 docs/output.rst                               |  4 +-
 torax/_src/output_tools/output.py             |  2 +
 torax/_src/output_tools/post_processing.py    | 10 ++++
 torax/_src/output_tools/tests/output_test.py  | 15 ++++++
 .../tests/post_processing_test.py             | 37 ++++++++++++++
 torax/_src/physics/collisions.py              | 49 ++++++++++++-------
 torax/_src/physics/tests/collisions_test.py   |  8 +++
 7 files changed, 105 insertions(+), 20 deletions(-)

diff --git a/docs/output.rst b/docs/output.rst
index eb549e0c0..902deca5a 100644
--- a/docs/output.rst
+++ b/docs/output.rst
@@ -468,6 +468,9 @@ These are called out in the list of profiles below, and generally relate to:
   Derivative of plasma surface area enclosed by each flux surface, with respect
   to the normalized toroidal flux coordinate rho_norm [:math:`m^2`].
 
+``tau_ei`` (time, rho_face_norm)
+  Electron-ion collision time [:math:`s`].
+
 ``T_e`` (time, rho_norm)
   Electron temperature [:math:`keV`].
 
@@ -1003,4 +1006,3 @@ purposes or to rerun the simulation.
   # We can also use ToraxConfig to run the simulation again.
   torax_config = torax.ToraxConfig.from_dict(config_dict)
   new_output = torax.run_simulation(torax_config)
-
diff --git a/torax/_src/output_tools/output.py b/torax/_src/output_tools/output.py
index 7bbe86aec..0dd72d372 100644
--- a/torax/_src/output_tools/output.py
+++ b/torax/_src/output_tools/output.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Module containing functions for saving and loading simulation output."""
+
 from collections.abc import Mapping, Sequence
 import dataclasses
 import functools
@@ -112,6 +113,7 @@
 
 # Post processed outputs
 Q_FUSION = "Q_fusion"
+TAU_EI = "tau_ei"
 
 # Edge model outputs
 SEED_IMPURITY_CONCENTRATIONS = "seed_impurity_concentrations"
diff --git a/torax/_src/output_tools/post_processing.py b/torax/_src/output_tools/post_processing.py
index fc109118d..f7119f7cc 100644
--- a/torax/_src/output_tools/post_processing.py
+++ b/torax/_src/output_tools/post_processing.py
@@ -31,6 +31,7 @@
 from torax._src.orchestration import sim_state as sim_state_lib
 from torax._src.output_tools import impurity_radiation
 from torax._src.output_tools import safety_factor_fit
+from torax._src.physics import collisions
 from torax._src.physics import formulas
 from torax._src.physics import psi_calculations
 from torax._src.physics import rotation
@@ -65,6 +66,7 @@ class PostProcessedOutputs:
       law derived from the updated (2020) ITER H-mode confinement database
     FFprime: FF' on the face grid, where F is the toroidal flux function
     psi_norm: Normalized poloidal flux on the face grid [Wb]
+    tau_ei: Electron-ion collision time on the face grid [s]
     P_heat_i: Total ion heating power: all sources - sinks. i.e. auxiliary
       heating + ion-electron exchange + fusion + (negative) radiation sinks [W].
     P_heat_e: Total electron heating power: all sources - sinks. i.e. auxiliary
@@ -219,6 +221,7 @@ class PostProcessedOutputs:
   H20: array_typing.FloatScalar
   FFprime: array_typing.FloatVector
   psi_norm: array_typing.FloatVector
+  tau_ei: array_typing.FloatVector
   # Integrated heat sources
   P_SOL_i: array_typing.FloatScalar
   P_SOL_e: array_typing.FloatScalar
@@ -337,6 +340,7 @@ def zeros(cls, geo: geometry.Geometry) -> typing_extensions.Self:
         H20=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         FFprime=jnp.zeros(geo.rho_face.shape),
         psi_norm=jnp.zeros(geo.rho_face.shape),
+        tau_ei=jnp.zeros(geo.rho_face.shape),
         P_SOL_i=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         P_SOL_e=jnp.array(0.0, dtype=jax_utils.get_dtype()),
         P_SOL_total=jnp.array(0.0, dtype=jax_utils.get_dtype()),
@@ -677,6 +681,11 @@ def make_post_processed_outputs(
   # Calculate normalized poloidal flux.
   psi_face = sim_state.core_profiles.psi.face_value()
   psi_norm_face = (psi_face - psi_face[0]) / (psi_face[-1] - psi_face[0])
+  tau_ei_face = collisions.calculate_tau_ei(
+      T_e=sim_state.core_profiles.T_e.face_value(),
+      n_e=sim_state.core_profiles.n_e.face_value(),
+      Z_eff=sim_state.core_profiles.Z_eff_face,
+  )
   integrated_sources = _calculate_integrated_sources(
       sim_state.geometry,
       sim_state.core_profiles,
@@ -956,6 +965,7 @@ def cumulative_values():
       H20=H20,
       FFprime=FFprime_face,
       psi_norm=psi_norm_face,
+      tau_ei=tau_ei_face,
       **integrated_sources,
       Q_fusion=Q_fusion,
       P_LH=P_LH_martin,
diff --git a/torax/_src/output_tools/tests/output_test.py b/torax/_src/output_tools/tests/output_test.py
index 373242db0..93dd2bb5b 100644
--- a/torax/_src/output_tools/tests/output_test.py
+++ b/torax/_src/output_tools/tests/output_test.py
@@ -155,6 +155,21 @@ def test_core_transport_is_saved(self):
         (1, len(self.geo.rho_face_norm)),
     )
 
+  def test_tau_ei_is_saved(self):
+    """Tests that electron-ion collision time is saved as a face profile."""
+    output_xr = self.history.simulation_output_to_xr()
+    profiles_dataset = output_xr.children[output.PROFILES].dataset
+
+    self.assertIn(output.TAU_EI, profiles_dataset.data_vars)
+    self.assertEqual(
+        profiles_dataset[output.TAU_EI].dims,
+        (output.TIME, output.RHO_FACE_NORM),
+    )
+    np.testing.assert_allclose(
+        profiles_dataset[output.TAU_EI].values[0],
+        self._output_state.tau_ei,
+    )
+
   def test_geometry_is_saved(self):
     """Tests that the geometry is saved correctly."""
     # Construct a second state with a slightly different geometry.
diff --git a/torax/_src/output_tools/tests/post_processing_test.py b/torax/_src/output_tools/tests/post_processing_test.py
index 43e59f0f7..f65bbf0eb 100644
--- a/torax/_src/output_tools/tests/post_processing_test.py
+++ b/torax/_src/output_tools/tests/post_processing_test.py
@@ -27,6 +27,7 @@
 from torax._src.orchestration import run_simulation
 from torax._src.orchestration import sim_state
 from torax._src.output_tools import post_processing
+from torax._src.physics import collisions
 from torax._src.sources import source_profiles as source_profiles_lib
 from torax._src.test_utils import default_configs
 from torax._src.test_utils import default_sources
@@ -240,6 +241,42 @@ def test_zero_sources_do_not_make_nans(self):
         post_processed_outputs.check_for_errors(), state.SimError.NO_ERROR
     )
 
+  def test_tau_ei_output(self):
+    """Checks electron-ion collision time is calculated on the face grid."""
+    input_state = sim_state.SimState(
+        t=jnp.array(0.0),
+        dt=jnp.array(1e-3),
+        core_profiles=self.core_profiles,
+        core_transport=state.CoreTransport.zeros(self.geo),
+        core_sources=self.source_profiles,
+        geometry=self.geo,
+        solver_numeric_outputs=state.SolverNumericOutputs(
+            solver_error_state=np.array(0, jax_utils.get_int_dtype()),
+            outer_solver_iterations=np.array(0, jax_utils.get_int_dtype()),
+            inner_solver_iterations=np.array(0, jax_utils.get_int_dtype()),
+            sawtooth_crash=False,
+        ),
+        edge_outputs=None,
+        time_step_calculator_state=(
+            self.models.time_step_calculator.initial_state(self.runtime_params)
+        ),
+    )
+
+    outputs = post_processing.make_post_processed_outputs(
+        sim_state=input_state,
+        runtime_params=self.runtime_params,
+        previous_post_processed_outputs=post_processing.PostProcessedOutputs.zeros(
+            self.geo
+        ),
+    )
+
+    expected_tau_ei = collisions.calculate_tau_ei(
+        T_e=self.core_profiles.T_e.face_value(),
+        n_e=self.core_profiles.n_e.face_value(),
+        Z_eff=self.core_profiles.Z_eff_face,
+    )
+    np.testing.assert_allclose(outputs.tau_ei, expected_tau_ei)
+
   def test_current_outputs(self):
     """Checks calculation of current-related outputs."""
     # Setup non-zero bootstrap current
diff --git a/torax/_src/physics/collisions.py b/torax/_src/physics/collisions.py
index 2d8ef9745..dd27ea1b6 100644
--- a/torax/_src/physics/collisions.py
+++ b/torax/_src/physics/collisions.py
@@ -27,6 +27,7 @@
       electron-ion collisions.
     - calculate_log_lambda_ii: Calculates the Coulomb logarithm for ion-ion
       collisions.
+    - calculate_tau_ei: Calculates the electron-ion collision time.
     - calculate_tau_ii: Calculates the ion-ion collision time.
     - _calculate_weighted_Z_eff: Calculates ion mass weighted Z_eff used in
       the equipartion calculation.
@@ -101,26 +102,12 @@ def calc_nu_star(
   Returns:
     nu_star: on face grid.
   """
-
-  # Calculate Coulomb logarithm
-  log_lambda_ei_face = calculate_log_lambda_ei(
-      core_profiles.T_e.face_value(),
-      core_profiles.n_e.face_value(),
-  )
-
-  # ion_electron collisionality
-  log_tau_e_Z1 = _calculate_log_tau_e_Z1(
-      core_profiles.T_e.face_value(),
-      core_profiles.n_e.face_value(),
-      log_lambda_ei_face,
-  )
-
-  nu_e = (
-      1
-      / jnp.exp(log_tau_e_Z1)
-      * core_profiles.Z_eff_face
-      * collisionality_multiplier
+  tau_ei = calculate_tau_ei(
+      T_e=core_profiles.T_e.face_value(),
+      n_e=core_profiles.n_e.face_value(),
+      Z_eff=core_profiles.Z_eff_face,
   )
+  nu_e = 1 / tau_ei * collisionality_multiplier
 
   # calculate bounce time
   tau_bounce = (
@@ -144,6 +131,30 @@ def calc_nu_star(
   return nustar
 
 
+def calculate_tau_ei(
+    T_e: jax.Array,
+    n_e: jax.Array,
+    Z_eff: jax.Array,
+) -> jax.Array:
+  """Calculates electron-ion collision time.
+
+  The Z=1 collision time is based on Wesson 3rd edition p729. For multi-species
+  plasmas this returns the effective electron-ion collision time used by
+  `calc_nu_star`, scaling the collision frequency by Z_eff.
+
+  Args:
+    T_e: Electron temperature [keV].
+    n_e: Electron density [m^-3].
+    Z_eff: Effective ion charge [dimensionless].
+
+  Returns:
+    Electron-ion collision time [s].
+  """
+  log_lambda_ei = calculate_log_lambda_ei(T_e, n_e)
+  log_tau_e_Z1 = _calculate_log_tau_e_Z1(T_e, n_e, log_lambda_ei)
+  return jnp.exp(log_tau_e_Z1) / Z_eff
+
+
 def fast_ion_fractional_heating_formula(
     birth_energy: float | array_typing.FloatVector,
     T_e: array_typing.FloatVector,
diff --git a/torax/_src/physics/tests/collisions_test.py b/torax/_src/physics/tests/collisions_test.py
index 2820b19d9..1caf01ef6 100644
--- a/torax/_src/physics/tests/collisions_test.py
+++ b/torax/_src/physics/tests/collisions_test.py
@@ -77,6 +77,14 @@ def test_calculate_log_lambda_ii(self, T_i_ev, n_i, Z_i, expected):
     result = collisions.calculate_log_lambda_ii(T_i_kev, n_i, Z_i)
     np.testing.assert_allclose(result, expected, atol=1e-6)
 
+  def test_calculate_tau_ei_scales_with_zeff(self):
+    T_e = jnp.array([1.0, 2.0])
+    n_e = jnp.array([1e20, 2e20])
+    tau_z1 = collisions.calculate_tau_ei(T_e, n_e, Z_eff=jnp.ones_like(T_e))
+    tau_z2 = collisions.calculate_tau_ei(T_e, n_e, Z_eff=2 * jnp.ones_like(T_e))
+
+    np.testing.assert_allclose(tau_z2, tau_z1 / 2.0)
+
   # TODO(b/377225415): generalize to arbitrary number of ions.
   @parameterized.parameters([
       dict(