diff --git a/torax/_src/neoclassical/transport/angioni_sauter.py b/torax/_src/neoclassical/transport/angioni_sauter.py index a827b2e69..61e11bf4b 100644 --- a/torax/_src/neoclassical/transport/angioni_sauter.py +++ b/torax/_src/neoclassical/transport/angioni_sauter.py @@ -21,12 +21,14 @@ https://gitlab.epfl.ch/spc/public/neos [O. Sauter et al] """ +import dataclasses from typing import Annotated, Literal +import jax from jax import numpy as jnp -from torax._src import array_typing -from torax._src import constants -from torax._src import state +from typing_extensions import override + +from torax._src import array_typing, constants, state from torax._src.config import runtime_params as runtime_params_lib from torax._src.geometry import geometry as geometry_lib from torax._src.neoclassical import formulas @@ -34,25 +36,46 @@ from torax._src.neoclassical.transport import runtime_params as transport_runtime_params from torax._src.physics import collisions from torax._src.torax_pydantic import torax_pydantic -from typing_extensions import override # pylint: disable=invalid-name +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class RuntimeParams(transport_runtime_params.RuntimeParams): + """RuntimeParams for the Angioni-Sauter neoclassical transport model.""" + + use_shaing_ion_correction: array_typing.BoolScalar + shaing_ion_multiplier: array_typing.FloatScalar + shaing_blend_start: array_typing.FloatScalar + shaing_blend_rate: array_typing.FloatScalar + + class AngioniSauterModelConfig(base.NeoclassicalTransportModelConfig): """Pydantic model for the Angioni-Sauter neoclassical transport model.""" model_name: Annotated[ Literal['angioni_sauter'], torax_pydantic.JAX_STATIC ] = 'angioni_sauter' + use_shaing_ion_correction: bool = False + shaing_ion_multiplier: float = 1.8 + shaing_blend_start: float = 0.2 + shaing_blend_rate: float = 5.0 @override def build_model(self) -> 'AngioniSauterModel': return AngioniSauterModel() @override - def build_runtime_params(self) -> transport_runtime_params.RuntimeParams: - return super().build_runtime_params() + def build_runtime_params(self) -> RuntimeParams: + base_kwargs = dataclasses.asdict(super().build_runtime_params()) + return RuntimeParams( + use_shaing_ion_correction=self.use_shaing_ion_correction, + shaing_ion_multiplier=self.shaing_ion_multiplier, + shaing_blend_start=self.shaing_blend_start, + shaing_blend_rate=self.shaing_blend_rate, + **base_kwargs + ) class AngioniSauterModel(base.NeoclassicalTransportModel): @@ -65,12 +88,47 @@ def _call_implementation( geometry: geometry_lib.Geometry, core_profiles: state.CoreProfiles, ) -> base.NeoclassicalTransport: - """Calculates neoclassical transport coefficients.""" - return _calculate_angioni_sauter_transport( + """Calculates neoclassical transport coefficients with smooth blend. + + When use_shaing_ion_correction is enabled, chi_ion is smoothly blended + between Shaing (near axis) and Angioni-Sauter (far from axis) models using + an exponential transition function. + """ + angioni_sauter = _calculate_angioni_sauter_transport( runtime_params=runtime_params, geometry=geometry, core_profiles=core_profiles, ) + shaing = _calculate_shaing_transport( + runtime_params=runtime_params, + geometry=geometry, + core_profiles=core_profiles, + ) + + + # Calculate sigmoid blend weight for Angioni-Sauter (alpha) + # If correction disabled: alpha = 1 (pure Angioni-Sauter) + # If correction enabled: alpha varies smoothly with rho_norm + alpha = jnp.where( + runtime_params.neoclassical.transport.use_shaing_ion_correction, + _calculate_blend_alpha( + rho_norm=geometry.rho_face_norm, + start=runtime_params.neoclassical.transport.shaing_blend_start, + rate=runtime_params.neoclassical.transport.shaing_blend_rate, + ), + 1.0, # Pure Angioni-Sauter when correction disabled + ) + + return base.NeoclassicalTransport( + # Ion transport blend: (1-alpha)*Shaing + alpha*Angioni-Sauter + chi_neo_i=(1.0 - alpha) * shaing.chi_neo_i + + alpha * angioni_sauter.chi_neo_i, + # Electron transport: pure Angioni-Sauter + chi_neo_e=angioni_sauter.chi_neo_e, + D_neo_e=angioni_sauter.D_neo_e, + V_neo_e=angioni_sauter.V_neo_e, + V_neo_ware_e=angioni_sauter.V_neo_ware_e, + ) def __hash__(self) -> int: return hash(self.__class__.__name__) @@ -587,3 +645,111 @@ def _calculate_Lmn( ) return Lmn_e, Lmn_i + + +def _calculate_shaing_transport( + runtime_params: runtime_params_lib.RuntimeParams, + geometry: geometry_lib.Geometry, + core_profiles: state.CoreProfiles, +) -> base.NeoclassicalTransport: + """JIT-compatible implementation of the Shaing transport model. + + Currently only implements ion thermal transport. Other contributions are + negligible. + + Args: + runtime_params: Runtime parameters. + geometry: Geometry object. + core_profiles: Core profiles object. + + Returns: + - Neoclassical transport coefficients. + - Radius of validity (in terms of rho_norm) of Shaing model for electrons. + - Radius of validity (in terms of rho_norm) of Shaing model for ions. + """ + # Aliases for readability + m_ion = core_profiles.A_i * constants.CONSTANTS.m_amu + q = core_profiles.q_face + kappa = geometry.elongation_face # Note: denoted delta in Shaing + F = geometry.F_face # Note: denoted I in Shaing + R = geometry.R_major_profile_face + T_i_J = core_profiles.T_i.face_value() * constants.CONSTANTS.keV_to_J + + # Collisionality + ln_Lambda_ii = collisions.calculate_log_lambda_ii( + core_profiles.T_i.face_value(), + core_profiles.n_i.face_value(), + core_profiles.Z_i_face, + ) + tau_ii = collisions.calculate_tau_ii( + A_i=core_profiles.A_i, + Z_i=core_profiles.Z_i_face, + T_i=core_profiles.T_i.face_value(), + n_i=core_profiles.n_i.face_value(), + ln_Lambda_ii=ln_Lambda_ii, + ) + nu_ii = 1 / tau_ii # Ion-ion collision frequency + + # Thermal velocity + v_t_ion = jnp.sqrt(2 * T_i_J / m_ion) + + # Larmor (gyro)frequency + Omega_0_ion = ( + constants.CONSTANTS.q_e * core_profiles.Z_i_face * geometry.B_0 / m_ion + ) + + # Large aspect ratio approximation (Equation 3, Shaing March 1997) + C_1 = (2 * q / (kappa * F * R)) ** (1 / 2) + + # Conversion from flux^2/s -> m^2/s + # TODO: make a more informed choice for dpsi_drhon near the axis (currently + # we simply copy the value at i=1). This is ok as chi[0] is never used. + dpsi_drhon = core_profiles.psi.face_grad() + dpsi_drhon = dpsi_drhon.at[0].set(dpsi_drhon[1]) + conversion_factor = 1 / (dpsi_drhon / (2*jnp.pi*geometry.rho_b)) ** 2 + + # Trapped particle fraction (Equation 46, Shaing March 1997) + f_t_ion = (F * v_t_ion * C_1**2 / Omega_0_ion) ** (1 / 3) + + # Orbit width in psi coordinates (Equation 73, Shaing March 1997) + Delta_psi_ion = (F**2 * v_t_ion**2 * C_1 / Omega_0_ion**2) ** (2 / 3) + + # Chi i term (Equation 74, Shaing March 1997) + # psi normalization difference accounted for in conversion_factor + chi_i = (nu_ii * Delta_psi_ion**2 / f_t_ion) * conversion_factor + + return base.NeoclassicalTransport( + chi_neo_i=runtime_params.neoclassical.transport.shaing_ion_multiplier + * chi_i, + chi_neo_e=jnp.zeros_like(geometry.rho_face), + D_neo_e=jnp.zeros_like(geometry.rho_face), + V_neo_e=jnp.zeros_like(geometry.rho_face), + V_neo_ware_e=jnp.zeros_like(geometry.rho_face), + ) + + +def _calculate_blend_alpha( + rho_norm: array_typing.FloatVectorFace, + start: array_typing.FloatScalar, + rate: array_typing.FloatScalar, +) -> array_typing.FloatVectorFace: + """Calculate blending weight between Angioni-Sauter and Shaing models. + + The blend is: + result = (1-alpha)*Shaing + alpha*Angioni-Sauter + where alpha = 1 / (1 + exp(-2*rate*(rho_norm - start))). + + This gives: + - At axis (rho_norm = 0 << start): alpha ~ 0 (pure Shaing) + - At start: alpha = 0.5 (equal blend) + - Far from axis (rho_norm >> start): alpha ~ 1 (pure Angioni-Sauter) + + Args: + rho_norm: Normalized toroidal flux coordinate (face grid) + start: Rho norm value where blend transition is centered + rate: Controls transition steepness (higher = sharper transition) + + Returns: + Blend factor alpha in range [0, 1] + """ + return 1.0 / (1.0 + jnp.exp(-2.0 * rate * (rho_norm - start))) diff --git a/torax/_src/neoclassical/transport/tests/angioni_sauter_test.py b/torax/_src/neoclassical/transport/tests/angioni_sauter_test.py index 254ce67c5..69884631b 100644 --- a/torax/_src/neoclassical/transport/tests/angioni_sauter_test.py +++ b/torax/_src/neoclassical/transport/tests/angioni_sauter_test.py @@ -12,11 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -from absl.testing import absltest +import dataclasses + import numpy as np +from absl.testing import absltest + +from torax._src import state from torax._src.config import build_runtime_params +from torax._src.config import runtime_params as runtime_params_lib from torax._src.core_profiles import initialization -from torax._src.neoclassical.transport import angioni_sauter +from torax._src.geometry import geometry +from torax._src.neoclassical.transport import angioni_sauter, base from torax._src.torax_pydantic import model_config _N_RHO = 10 @@ -25,9 +31,11 @@ class AngioniSauterTest(absltest.TestCase): - - def test_angioni_sauter_against_reference_values(self): - """Reference values generated from Angioni-Sauter with NEOS verification.""" + def _get_reference_runtime_params_geo_and_core_profiles(self) -> tuple[ + runtime_params_lib.RuntimeParams, + geometry.Geometry, + state.CoreProfiles + ]: torax_config = model_config.ToraxConfig.from_dict({ 'profile_conditions': { 'Ip': 15e6, @@ -72,23 +80,101 @@ def test_angioni_sauter_against_reference_values(self): neoclassical_models=neoclassical_models, ) + return runtime_params, geo, core_profiles + + def test_angioni_sauter_against_reference_values(self): + """Reference values generated from running Angioni-Sauter.""" + runtime_params, geo, core_profiles = ( + self._get_reference_runtime_params_geo_and_core_profiles() + ) + + # Test raw Angioni-Sauter values result = angioni_sauter._calculate_angioni_sauter_transport( runtime_params, geo, core_profiles ) np.testing.assert_allclose( - result.chi_neo_i, _EXPECTED_CHI_NEO_I, atol=_A_TOL, rtol=_R_TOL + result.chi_neo_i, + _ANGIONI_SAUTER_REFERENCE_VALUES.chi_neo_i, + atol=_A_TOL, + rtol=_R_TOL, + ) + np.testing.assert_allclose( + result.chi_neo_e, + _ANGIONI_SAUTER_REFERENCE_VALUES.chi_neo_e, + atol=_A_TOL, + rtol=_R_TOL, + ) + np.testing.assert_allclose( + result.D_neo_e, + _ANGIONI_SAUTER_REFERENCE_VALUES.D_neo_e, + atol=_A_TOL, + rtol=_R_TOL, + ) + np.testing.assert_allclose( + result.V_neo_e, + _ANGIONI_SAUTER_REFERENCE_VALUES.V_neo_e, + atol=_A_TOL, + rtol=_R_TOL, + ) + np.testing.assert_allclose( + result.V_neo_ware_e, + _ANGIONI_SAUTER_REFERENCE_VALUES.V_neo_ware_e, + atol=_A_TOL, + rtol=_R_TOL, + ) + + def test_angioni_sauter_with_shaing_against_reference_values(self): + """Reference values generated from Angioni-Sauter + Shaing ion correction""" + runtime_params, geo, core_profiles = ( + self._get_reference_runtime_params_geo_and_core_profiles() + ) + + # Modify runtime params to include settings for Shaing + neoclassical_runtime_params = runtime_params.neoclassical + neoclassical_runtime_params = dataclasses.replace( + neoclassical_runtime_params, + transport=angioni_sauter.AngioniSauterModelConfig( + use_shaing_ion_correction=True + ).build_runtime_params(), + ) + modified_runtime_params = dataclasses.replace( + runtime_params, + neoclassical=neoclassical_runtime_params, + ) + + # Test blended Angioni-Sauter + Shaing values + result = angioni_sauter.AngioniSauterModel()._call_implementation( + modified_runtime_params, geo, core_profiles + ) + np.testing.assert_allclose( + result.chi_neo_i, + _ANGIONI_SAUTER_SHAING_REFERENCE_VALUES.chi_neo_i, + atol=_A_TOL, + rtol=_R_TOL, ) np.testing.assert_allclose( - result.chi_neo_e, _EXPECTED_CHI_NEO_E, atol=_A_TOL, rtol=_R_TOL + result.chi_neo_e, + _ANGIONI_SAUTER_SHAING_REFERENCE_VALUES.chi_neo_e, + atol=_A_TOL, + rtol=_R_TOL, ) np.testing.assert_allclose( - result.D_neo_e, _EXPECTED_D_NEO_E, atol=_A_TOL, rtol=_R_TOL + result.D_neo_e, + _ANGIONI_SAUTER_SHAING_REFERENCE_VALUES.D_neo_e, + atol=_A_TOL, + rtol=_R_TOL, ) np.testing.assert_allclose( - result.V_neo_e, _EXPECTED_V_NEO_E, atol=_A_TOL, rtol=_R_TOL + result.V_neo_e, + _ANGIONI_SAUTER_SHAING_REFERENCE_VALUES.V_neo_e, + atol=_A_TOL, + rtol=_R_TOL, ) np.testing.assert_allclose( - result.V_neo_ware_e, _EXPECTED_V_NEO_WARE_E, atol=_A_TOL, rtol=_R_TOL + result.V_neo_ware_e, + _ANGIONI_SAUTER_SHAING_REFERENCE_VALUES.V_neo_ware_e, + atol=_A_TOL, + rtol=_R_TOL, ) @@ -98,76 +184,94 @@ def test_angioni_sauter_against_reference_values(self): # # The implementation was independently tested against NEOS up to the # generation of the Kmn matrix. -_EXPECTED_CHI_NEO_I = np.array([ - -0.0, - 0.01220085, - 0.02223608, - 0.03117304, - 0.03891618, - 0.04568965, - 0.05179111, - 0.0572006, - 0.06147531, - 0.06320731, - 0.0591895, -]) - -_EXPECTED_CHI_NEO_E = np.array([ - -0.0, - -0.00210023, - -0.0030792, - -0.00388683, - -0.0045548, - -0.00511068, - -0.0056083, - -0.0060884, - -0.00658147, - -0.00717367, - -0.00750323, -]) - -_EXPECTED_D_NEO_E = np.array([ - 0.0, - 0.00011698, - 0.00021105, - 0.00028474, - 0.00033721, - 0.00037529, - 0.00040377, - 0.00042199, - 0.00042404, - 0.00039292, - 0.0002924, -]) - -_EXPECTED_V_NEO_E = np.array([ - 0.0, - 1.07951440e-05, - 1.11015003e-05, - 1.54065751e-05, - 2.65710672e-05, - 4.42853751e-05, - 7.06387381e-05, - 1.12983269e-04, - 1.92360065e-04, - 3.86372126e-04, - 1.18868626e-03, -]) - -_EXPECTED_V_NEO_WARE_E = np.array([ - -0.0, - -0.00038114, - -0.00041759, - -0.00037123, - -0.00032066, - -0.00030646, - -0.0003312, - -0.00038565, - -0.00056229, - -0.00159816, - -0.00178913, -]) +_ANGIONI_SAUTER_REFERENCE_VALUES = base.NeoclassicalTransport( + chi_neo_i=np.array([ + -0.0, + 0.01220085, + 0.02223608, + 0.03117304, + 0.03891618, + 0.04568965, + 0.05179111, + 0.0572006, + 0.06147531, + 0.06320731, + 0.0591895, + ]), + chi_neo_e=np.array([ + -0.0, + -0.00210023, + -0.0030792, + -0.00388683, + -0.0045548, + -0.00511068, + -0.0056083, + -0.0060884, + -0.00658147, + -0.00717367, + -0.00750323, + ]), + D_neo_e=np.array([ + 0.0, + 0.00011698, + 0.00021105, + 0.00028474, + 0.00033721, + 0.00037529, + 0.00040377, + 0.00042199, + 0.00042404, + 0.00039292, + 0.0002924, + ]), + V_neo_e=np.array([ + 0.0, + 1.07951440e-05, + 1.11015003e-05, + 1.54065751e-05, + 2.65710672e-05, + 4.42853751e-05, + 7.06387381e-05, + 1.12983269e-04, + 1.92360065e-04, + 3.86372126e-04, + 1.18868626e-03, + ]), + V_neo_ware_e=np.array([ + -0.0, + -0.00038114, + -0.00041759, + -0.00037123, + -0.00032066, + -0.00030646, + -0.0003312, + -0.00038565, + -0.00056229, + -0.00159816, + -0.00178913, + ]) +) +# Shaing correction only affects ions, so we can reuse the other values +_ANGIONI_SAUTER_SHAING_REFERENCE_VALUES = base.NeoclassicalTransport( + chi_neo_i=np.array([ + 0.20237419, + 0.17130245, + 0.03031974, + 0.02593766, + 0.03523606, + 0.04391418, + 0.05103481, + 0.05690314, + 0.06136931, + 0.06317694, + 0.05918355, + ]), + chi_neo_e=_ANGIONI_SAUTER_REFERENCE_VALUES.chi_neo_e, + D_neo_e=_ANGIONI_SAUTER_REFERENCE_VALUES.D_neo_e, + V_neo_e=_ANGIONI_SAUTER_REFERENCE_VALUES.V_neo_e, + V_neo_ware_e=_ANGIONI_SAUTER_REFERENCE_VALUES.V_neo_ware_e, +) if __name__ == '__main__': absltest.main() diff --git a/torax/_src/physics/collisions.py b/torax/_src/physics/collisions.py index 64c7d93bc..1ccfb3473 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_ii: Calculates the ion-ion collision time. - _calculate_weighted_Z_eff: Calculates ion mass weighted Z_eff used in the equipartion calculation. - _calculate_log_tau_e_Z1: Calculates log of electron-ion collision time for @@ -244,6 +245,41 @@ def calculate_log_lambda_ii( return 30.0 - 0.5 * jnp.log(n_i) + 1.5 * jnp.log(T_i_ev) - 3.0 * jnp.log(Z_i) +def calculate_tau_ii( + A_i: jax.Array, + Z_i: jax.Array, + T_i: jax.Array, + n_i: jax.Array, + ln_Lambda_ii: jax.Array, +) -> jax.Array: + """Calculates ion-ion (self) collision time for a single ion species. + + See Wesson 3rd edition p730. + + Args: + A_i: Ion mass in amu. + Z_i: Ion charge. + T_i: Ion temperature in keV. + n_i: Ion density in m^-3. + ln_Lambda_ii: Coulomb logarithm for ion-ion collisions. + + Returns: + Ion-ion collision time in seconds. + """ + log_tau_ii = ( + jnp.log(12) + + 1.5 * jnp.log(jnp.pi) + + 2 * jnp.log(constants.CONSTANTS.epsilon_0) + + 0.5 * (jnp.log(A_i) + jnp.log(constants.CONSTANTS.m_amu)) + + 1.5 * (jnp.log(T_i) + jnp.log(constants.CONSTANTS.keV_to_J)) + - jnp.log(n_i) + - 4 * jnp.log(Z_i) + - 4 * jnp.log(constants.CONSTANTS.q_e) + - jnp.log(ln_Lambda_ii) + ) + return jnp.exp(log_tau_ii) + + # TODO(b/377225415): generalize to arbitrary number of ions. def _calculate_weighted_Z_eff( core_profiles: state.CoreProfiles, diff --git a/torax/examples/step_flattop_bgb.py b/torax/examples/step_flattop_bgb.py index 2f723536f..0c7683a05 100644 --- a/torax/examples/step_flattop_bgb.py +++ b/torax/examples/step_flattop_bgb.py @@ -32,9 +32,9 @@ import imas import numpy as np + from torax._src import path_utils -from torax._src.imas_tools.input import core_profiles -from torax._src.imas_tools.input import loader +from torax._src.imas_tools.input import core_profiles, loader # Load IDSs path = ( @@ -145,20 +145,16 @@ "D_e_max": 100.0, "V_e_min": -50.0, "V_e_max": 50.0, - # Patching - # Replaces neoclassical in the core, pending potato orbit correction - # https://github.com/google-deepmind/torax/issues/1406 - "apply_inner_patch": True, - "rho_inner": 0.05, - "chi_e_inner": 1.0, - "chi_i_inner": 15.0, # Smoothing "smooth_everywhere": True, "smoothing_width": 0.05, }, "neoclassical": { "bootstrap_current": {"model_name": "sauter"}, - "transport": {"model_name": "angioni_sauter"}, + "transport": { + "model_name": "angioni_sauter", + "use_shaing_ion_correction": True, + }, }, "numerics": { "t_initial": 0.0, diff --git a/torax/tests/test_data/test_step_flattop_bgb.nc b/torax/tests/test_data/test_step_flattop_bgb.nc index c5c7ad2a4..c3a8da0c6 100644 Binary files a/torax/tests/test_data/test_step_flattop_bgb.nc and b/torax/tests/test_data/test_step_flattop_bgb.nc differ