Add nu_star to TORAX outputs

docs/output.rst

@@ -346,6 +346,10 @@ These are called out in the list of profiles below, and generally relate to:
 ``n_impurity_species`` (impurity_symbol, time, rho_cell_norm)
   True impurity density per species [:math:`m^{-3}`].
 
+``nu_star`` (time, rho_face_norm)
+  Electron-ion collisionality normalized by the bounce frequency
+  [dimensionless].
+
 ``p_alpha_e`` (time, rho_cell_norm)
   Fusion alpha heating power density to electrons [:math:`W/m^3`]. Only output
   if ``fusion`` source is active.
@@ -1003,4 +1007,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)
-

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
@@ -111,6 +112,7 @@
 TIME = "time"
 
 # Post processed outputs
+NU_STAR = "nu_star"
 Q_FUSION = "Q_fusion"
 
 # Edge model outputs

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,8 @@ 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]
+    nu_star: Electron-ion collisionality normalized by bounce frequency on the
+      face grid [dimensionless]
     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 +222,7 @@ class PostProcessedOutputs:
   H20: array_typing.FloatScalar
   FFprime: array_typing.FloatVector
   psi_norm: array_typing.FloatVector
+  nu_star: array_typing.FloatVector
   # Integrated heat sources
   P_SOL_i: array_typing.FloatScalar
   P_SOL_e: array_typing.FloatScalar
@@ -337,6 +341,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),
+        nu_star=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 +682,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])
+  nu_star_face = collisions.calc_nu_star(
+      geo=sim_state.geometry,
+      core_profiles=sim_state.core_profiles,
+      collisionality_multiplier=1.0,
+  )
   integrated_sources = _calculate_integrated_sources(
       sim_state.geometry,
       sim_state.core_profiles,
@@ -956,6 +966,7 @@ def cumulative_values():
       H20=H20,
       FFprime=FFprime_face,
       psi_norm=psi_norm_face,
+      nu_star=nu_star_face,
       **integrated_sources,
       Q_fusion=Q_fusion,
       P_LH=P_LH_martin,

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_nu_star_is_saved(self):
+    """Tests that normalized collisionality 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.NU_STAR, profiles_dataset.data_vars)
+    self.assertEqual(
+        profiles_dataset[output.NU_STAR].dims,
+        (output.TIME, output.RHO_FACE_NORM),
+    )
+    np.testing.assert_allclose(
+        profiles_dataset[output.NU_STAR].values[0],
+        self._output_state.nu_star,
+    )
+
   def test_geometry_is_saved(self):
     """Tests that the geometry is saved correctly."""
     # Construct a second state with a slightly different geometry.

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_nu_star_output(self):
+    """Checks normalized collisionality 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_nu_star = collisions.calc_nu_star(
+        geo=self.geo,
+        core_profiles=self.core_profiles,
+        collisionality_multiplier=1.0,
+    )
+    np.testing.assert_allclose(outputs.nu_star, expected_nu_star)
+
   def test_current_outputs(self):
     """Checks calculation of current-related outputs."""
     # Setup non-zero bootstrap current