diff --git a/torax/_src/output_tools/post_processing.py b/torax/_src/output_tools/post_processing.py index fc109118d..44e7c1522 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,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, diff --git a/torax/_src/output_tools/tests/post_processing_test.py b/torax/_src/output_tools/tests/post_processing_test.py index 43e59f0f7..6ee0b85fe 100644 --- a/torax/_src/output_tools/tests/post_processing_test.py +++ b/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.""" diff --git a/torax/_src/physics/tests/collisions_test.py b/torax/_src/physics/tests/collisions_test.py index 2820b19d9..1d6f1916e 100644 --- a/torax/_src/physics/tests/collisions_test.py +++ b/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()