Add normalized collisionality nu_star to 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,11 @@ 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: Normalized collisionality on the face grid [dimensionless]: the
+      electron-ion collision frequency normalized by the bounce frequency.
+      Reported with a collisionality multiplier of 1.0 (the physical
+      collisionality); the multiplier is a QLKNN sensitivity-testing knob and is
+      intentionally not applied to this diagnostic output.
     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 +225,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 +344,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 +685,15 @@ 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])
+  # Normalized collisionality (electron-ion collision frequency normalized by
+  # the bounce frequency) on the face grid. A collisionality multiplier of 1.0
+  # is used so the reported value is the physical collisionality; the multiplier
+  # is a QLKNN sensitivity-testing knob and is intentionally not applied here.
+  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 +973,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/post_processing_test.py

@@ -326,6 +326,39 @@ def test_current_outputs(self):
         rtol=1e-5,
     )
 
+  def test_nu_star_output(self):
+    """Checks nu_star is exposed as a physical, face-grid output."""
+    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
+        ),
+    )
+    # nu_star is defined on the face grid.
+    self.assertEqual(outputs.nu_star.shape, self.geo.rho_face_norm.shape)
+    # Collisionality is physically finite and positive.
+    self.assertTrue(np.all(np.isfinite(outputs.nu_star)))
+    self.assertTrue(np.all(outputs.nu_star > 0.0))
+
 
 class PostProcessingSimTest(sim_test_case.SimTestCase):
   """Tests for the cumulative outputs."""

torax/_src/physics/tests/collisions_test.py

@@ -19,8 +19,13 @@
 from jax import numpy as jnp
 import numpy as np
 from torax._src import state
+from torax._src.config import build_runtime_params
+from torax._src.core_profiles import initialization
 from torax._src.fvm import cell_variable
 from torax._src.physics import collisions
+from torax._src.test_utils import default_configs
+from torax._src.test_utils import default_sources
+from torax._src.torax_pydantic import model_config
 
 
 # pylint: disable=invalid-name
@@ -151,5 +156,58 @@ def test_calculate_weighted_Z_eff(
     )
 
 
+class CalcNuStarTest(parameterized.TestCase):
+  """Tests for `calc_nu_star` on a realistic core_profiles + geometry."""
+
+  def setUp(self):
+    super().setUp()
+    config = default_configs.get_default_config_dict()
+    config['sources'] = default_sources.get_default_source_config()
+    torax_config = model_config.ToraxConfig.from_dict(config)
+    self.runtime_params = (
+        build_runtime_params.RuntimeParamsProvider.from_config(torax_config)(
+            t=0.0
+        )
+    )
+    self.models = torax_config.build_models()
+    self.geo = torax_config.geometry.build_provider(t=0.0)
+    self.core_profiles = initialization.initial_core_profiles(
+        runtime_params=self.runtime_params,
+        geo=self.geo,
+        source_models=self.models.source_models,
+        neoclassical_models=self.models.neoclassical_models,
+    )
+
+  def test_nu_star_is_on_face_grid_positive_and_finite(self):
+    """nu_star is a physically positive, finite quantity on the face grid."""
+    nu_star = collisions.calc_nu_star(
+        geo=self.geo,
+        core_profiles=self.core_profiles,
+        collisionality_multiplier=1.0,
+    )
+    self.assertEqual(nu_star.shape, self.geo.rho_face.shape)
+    self.assertTrue(np.all(np.isfinite(nu_star)))
+    self.assertTrue(np.all(nu_star > 0.0))
+
+  @parameterized.parameters([2.0, 3.0, 0.5])
+  def test_nu_star_scales_linearly_with_collisionality_multiplier(
+      self, multiplier
+  ):
+    """nu_star is proportional to the collisionality multiplier."""
+    nu_star_unit = collisions.calc_nu_star(
+        geo=self.geo,
+        core_profiles=self.core_profiles,
+        collisionality_multiplier=1.0,
+    )
+    nu_star_scaled = collisions.calc_nu_star(
+        geo=self.geo,
+        core_profiles=self.core_profiles,
+        collisionality_multiplier=multiplier,
+    )
+    np.testing.assert_allclose(
+        nu_star_scaled, multiplier * nu_star_unit, rtol=1e-6
+    )
+
+
 if __name__ == '__main__':
   absltest.main()