diff --git a/docs/configuration.rst b/docs/configuration.rst index 6e26b3dd..87c7f182 100644 --- a/docs/configuration.rst +++ b/docs/configuration.rst @@ -1074,6 +1074,27 @@ Bremsstrahlung model from Wesson, with an optional correction for relativistic e ``use_relativistic_correction`` (bool = False) +impurity_radiation_heat_sink +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Various models for impurity radiation. Runtime params for each available model are listed separately + +``mode`` (str = 'model') + +``model_func`` (str = 'impurity_radiation_mavrin_fit') + +The following models are available: + +* ``'impurity_radiation_mavrin_fit'`` + Polynomial fits to ADAS data from `Mavrin, 2018. `_ + + ``radiation_multiplier`` (float = 1.0). Multiplication factor for radiation term for testing sensitivities. + +* ``'radially_constant_fraction_of_Pin'`` + Sets impurity radiation to be a constant fraction of the total external input power. + + ``fraction_of_total_power_density`` (float = 1.0). Fraction of total external input power to use for impurity radiation. + cyclotron_radiation_heat_sink ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/docs/physics_models.rst b/docs/physics_models.rst index 45499dfe..ed8a3ace 100644 --- a/docs/physics_models.rst +++ b/docs/physics_models.rst @@ -422,8 +422,6 @@ Presently, TORAX provides three built-in formula-based particle sources for the Radiation --------- -Currently, TORAX has dedicated models for Bremsstrahlung and cyclotron radiation. -Models for line radiation are left for future work. Bremsstrahlung ^^^^^^^^^^^^^^ @@ -439,8 +437,31 @@ with a deposition profile from `Artaud NF 2018 `_ + Provides radiative cooling rates for the following impurity species: + - Helium + - Lithium + - Beryllium + - Carbon + - Nitrogen + - Oxygen + - Neon + - Argon + - Krypton + - Xenon + - Tungsten + These cooling curves are multiplied by the electron density and impurity densities to obtain the impurity radiation power density. + The valid temperature range of the fit is [0.1-100] keV. + Ion Cyclotron Resonance Heating -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +------------------------------- Presently this source is implemented for a SPARC specific ICRH scenario. diff --git a/docs/roadmap.rst b/docs/roadmap.rst index c33e76eb..6e2a0512 100644 --- a/docs/roadmap.rst +++ b/docs/roadmap.rst @@ -3,23 +3,17 @@ Development Roadmap ################### -Short term development plans include: +Ongoing developments include: * Implementation of forward sensitivity calculations w.r.t. control inputs and parameters * More extensive tutorials -* Extended sinks - - * Cyclotron radiation - * Impurity line radiation - * Sawtooth model (Porcelli + reconnection) -* Multi-ion and impurity prescribed profiles * IMAS coupling Longer term desired features include: -* Neoclassical tearing modes (modified Rutherford equation) * Neoclassical transport + multi-ion transport, with a focus on heavy impurities +* Neoclassical tearing modes (modified Rutherford equation) * Stationary-state solver * Momentum transport diff --git a/torax/sources/impurity_radiation_heat_sink/impurity_radiation_heat_sink.py b/torax/sources/impurity_radiation_heat_sink/impurity_radiation_heat_sink.py index 57218629..809dd062 100644 --- a/torax/sources/impurity_radiation_heat_sink/impurity_radiation_heat_sink.py +++ b/torax/sources/impurity_radiation_heat_sink/impurity_radiation_heat_sink.py @@ -23,7 +23,7 @@ from torax.sources import source as source_lib from torax.sources import source_models as source_models_lib -from torax.sources.impurity_radiation_heat_sink import impurity_radiation_constant_fraction +from torax.sources.impurity_radiation_heat_sink import impurity_radiation_mavrin_fit @dataclasses.dataclass(kw_only=True, frozen=True, eq=True) @@ -32,10 +32,10 @@ class ImpurityRadiationHeatSink(source_lib.Source): SOURCE_NAME = "impurity_radiation_heat_sink" DEFAULT_MODEL_FUNCTION_NAME: ClassVar[str] = ( - impurity_radiation_constant_fraction.MODEL_FUNCTION_NAME + impurity_radiation_mavrin_fit.MODEL_FUNCTION_NAME ) model_func: source_lib.SourceProfileFunction - source_models: source_models_lib.SourceModels + source_models: source_models_lib.SourceModels | None = None @property def source_name(self) -> str: diff --git a/torax/sources/impurity_radiation_heat_sink/impurity_radiation_mavrin_fit.py b/torax/sources/impurity_radiation_heat_sink/impurity_radiation_mavrin_fit.py new file mode 100644 index 00000000..da50a572 --- /dev/null +++ b/torax/sources/impurity_radiation_heat_sink/impurity_radiation_mavrin_fit.py @@ -0,0 +1,246 @@ +# Copyright 2024 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Routines for calculating impurity radiation based on a polynomial fit.""" + +import dataclasses +import functools +from typing import Final, Mapping, Sequence +import chex +import immutabledict +import jax +import jax.numpy as jnp +import numpy as np +from torax import array_typing +from torax import constants +from torax import jax_utils +from torax import state +from torax.config import plasma_composition +from torax.config import runtime_params_slice +from torax.geometry import geometry +from torax.sources import runtime_params as runtime_params_lib +from torax.sources import source_models as source_models_lib + +MODEL_FUNCTION_NAME = 'impurity_radiation_mavrin_fit' + +# Polynomial fit coefficients from A. A. Mavrin (2018): +# Improved fits of coronal radiative cooling rates for high-temperature plasmas, +# Radiation Effects and Defects in Solids, 173:5-6, 388-398, +# DOI: 10.1080/10420150.2018.1462361 +_MAVRIN_L_COEFFS: Final[Mapping[str, array_typing.ArrayFloat]] = ( + immutabledict.immutabledict({ + 'He3': np.array([ + [2.5020e-02, -9.3730e-02, 1.0156e-01, 3.1469e-01, -3.5551e01], + ]), + 'He4': np.array([ + [2.5020e-02, -9.3730e-02, 1.0156e-01, 3.1469e-01, -3.5551e01], + ]), + 'Li': np.array([ + [3.5190e-02, -1.6070e-01, 2.5082e-01, 1.9475e-01, -3.5115e01], + ]), + 'Be': np.array([ + [4.1690e-02, -2.1384e-01, 3.8363e-01, 3.7270e-02, -3.4765e01], + ]), + 'C': np.array([ + [-7.2904e00, -1.6637e01, -1.2788e01, -5.0085e00, -3.4738e01], + [4.4470e-02, -2.9191e-01, 6.8856e-01, -3.6687e-01, -3.4174e01], + ]), + 'N': np.array([ + [-6.9621e00, -1.1570e01, -6.0605e00, -2.3614e00, -3.4065e01], + [5.8770e-02, -1.7160e-01, 7.6272e-01, -5.9668e-01, -3.3899e01], + [2.8350e-02, -2.2790e-01, 7.0047e-01, -5.2628e-01, -3.3913e01], + ]), + 'O': np.array([ + [0.0000e00, -5.3765e00, -1.7141e01, -1.5635e01, -3.7257e01], + [1.4360e-02, -2.0850e-01, 7.9655e-01, -7.6211e-01, -3.3640e01], + ]), + 'Ne': np.array([ + [1.5648e01, 2.8939e01, 1.5230e01, 1.7309e00, -3.3132e01], + [1.7244e-01, -3.9544e-01, 8.6842e-01, -8.7750e-01, -3.3290e01], + [-2.6930e-02, 4.3960e-02, 2.9731e-01, -4.5345e-01, -3.3410e01], + ]), + 'Ar': np.array([ + [1.5353e01, 3.9161e01, 3.0769e01, 6.5221e00, -3.2155e01], + [4.9806e00, -7.6887e00, 1.5389e00, 5.4490e-01, -3.2530e01], + [-8.2260e-02, 1.7480e-01, 6.1339e-01, -1.6674e00, -3.1853e01], + ]), + 'Kr': np.array([ + [-1.3564e01, -4.0133e01, -4.4723e01, -2.1484e01, -3.4512e01], + [-5.2704e00, -2.5865e00, 1.9148e00, -5.0091e-01, -3.1399e01], + [4.8356e-01, -2.9674e00, 6.6831e00, -6.3683e00, -2.9954e01], + ]), + 'Xe': np.array([ + [2.5615e01, 5.9580e01, 4.7081e01, 1.4351e01, -2.9303e01], + [1.0748e01, -1.1628e01, 1.2808e00, 5.9339e-01, -3.1113e01], + [1.0069e01, -3.6885e01, 4.8614e01, -2.7526e01, -2.5813e01], + [1.0858e00, -7.5181e00, 1.9619e01, -2.2592e01, -2.2138e01], + ]), + 'W': np.array([ + [-1.0103e-01, -1.0311e00, -9.5126e-01, 3.8304e-01, -3.0374e01], + [5.1849e01, -6.3303e01, 2.2824e01, -2.9208e00, -3.0238e01], + [-3.6759e-01, 2.6627e00, -6.2740e00, 5.2499e00, -3.2153e01], + ]), + }) +) + +# Temperature boundaries in keV, separating the rows for the fit coefficients. +_TEMPERATURE_INTERVALS: Final[Mapping[str, array_typing.ArrayFloat]] = ( + immutabledict.immutabledict({ + 'C': np.array([0.5]), + 'N': np.array([0.5, 2.0]), + 'O': np.array([0.3]), + 'Ne': np.array([0.7, 5.0]), + 'Ar': np.array([0.6, 3.0]), + 'Kr': np.array([0.447, 2.364]), + 'Xe': np.array([0.5, 2.5, 10.0]), + 'W': np.array([1.5, 4.0]), + }) +) + + +# pylint: disable=invalid-name +def _calculate_impurity_radiation_single_species( + Te: array_typing.ArrayFloat, + ion_symbol: str, +) -> array_typing.ArrayFloat: + """Calculates the line radiation for single impurity species. + + Polynomial fit range is 0.1-100 keV, which is well within the typical + bounds of core tokamak modelling. For safety, inputs are clipped to avoid + extrapolation outside this range. + + Args: + Te: Electron temperature [keV]. + ion_symbol: Species to calculate line radiation for. + + Returns: + LZ: Radiative cooling rate in units of Wm^3. + """ + + if ion_symbol not in constants.ION_SYMBOLS: + raise ValueError( + f'Invalid ion symbol: {ion_symbol}. Allowed symbols are :' + f' {constants.ION_SYMBOLS}' + ) + + # Avoid extrapolating fitted polynomial out of bounds. + Te = jnp.clip(Te, 0.1, 100.0) + + # Gather coefficients for each temperature + if ion_symbol in {'He3', 'He4', 'Be', 'Li'}: + interval_indices = 0 + else: + interval_indices = jnp.searchsorted(_TEMPERATURE_INTERVALS[ion_symbol], Te) + + L_coeffs_in_range = jnp.take( + _MAVRIN_L_COEFFS[ion_symbol], interval_indices, axis=0 + ).transpose() + + X = jnp.log10(Te) + log10_LZ = jnp.polyval(L_coeffs_in_range, X) + return 10**log10_LZ + + +@functools.partial( + jax_utils.jit, + static_argnames=[ + 'ion_symbols', + ], +) +def calculate_total_impurity_radiation( + ion_symbols: Sequence[str], + ion_mixture: plasma_composition.DynamicIonMixture, + Te: array_typing.ArrayFloat, +) -> array_typing.ArrayFloat: + """Calculates impurity line radiation profile (JAX-compatible). + + Args: + ion_symbols: Ion symbols of the impurity species. + ion_mixture: DynamicIonMixture object containing impurity information. + Te: Electron temperature [keV]. Can be any sized array, e.g. on cell grid, + face grid, or a single scalar. + + Returns: + effective_LZ: Total effective radiative cooling rate in units of Wm^3, + summed over all species in the mixture. The shape of LZ is the same as Te. + """ + + effective_LZ = jnp.zeros_like(Te) + for ion_symbol, fraction in zip(ion_symbols, ion_mixture.fractions): + effective_LZ += fraction * _calculate_impurity_radiation_single_species( + Te, ion_symbol + ) + return effective_LZ + + +def impurity_radiation_mavrin_fit( + static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice, + dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice, + geo: geometry.Geometry, + source_name: str, + core_profiles: state.CoreProfiles, + unused_source_models: source_models_lib.SourceModels | None = None, +) -> jax.Array: + """Model function for impurity radiation heat sink.""" + del (geo, unused_source_models) + effective_LZ = calculate_total_impurity_radiation( + ion_symbols=static_runtime_params_slice.impurity_names, + ion_mixture=dynamic_runtime_params_slice.plasma_composition.impurity, + Te=core_profiles.temp_el.value, + ) + dynamic_source_runtime_params = dynamic_runtime_params_slice.sources[ + source_name + ] + assert isinstance(dynamic_source_runtime_params, DynamicRuntimeParams) + radiation_profile = ( + effective_LZ + * core_profiles.ne.value + * core_profiles.nimp.value + * dynamic_source_runtime_params.radiation_multiplier + * dynamic_runtime_params_slice.numerics.nref**2 + ) + + # The impurity radiation heat sink is a negative source, so we return a + # negative profile. + return -radiation_profile + + +@dataclasses.dataclass(kw_only=True) +class RuntimeParams(runtime_params_lib.RuntimeParams): + radiation_multiplier: float = 1.0 + mode: runtime_params_lib.Mode = runtime_params_lib.Mode.MODEL_BASED + + def make_provider( + self, + torax_mesh: geometry.Grid1D | None = None, + ) -> 'RuntimeParamsProvider': + return RuntimeParamsProvider(**self.get_provider_kwargs(torax_mesh)) + + +@chex.dataclass +class RuntimeParamsProvider(runtime_params_lib.RuntimeParamsProvider): + """Provides runtime parameters for a given time and geometry.""" + + runtime_params_config: RuntimeParams + + def build_dynamic_params( + self, + t: chex.Numeric, + ) -> 'DynamicRuntimeParams': + return DynamicRuntimeParams(**self.get_dynamic_params_kwargs(t)) + + +@chex.dataclass(frozen=True) +class DynamicRuntimeParams(runtime_params_lib.DynamicRuntimeParams): + radiation_multiplier: array_typing.ScalarFloat diff --git a/torax/sources/register_source.py b/torax/sources/register_source.py index 3fc25932..3f4b330a 100644 --- a/torax/sources/register_source.py +++ b/torax/sources/register_source.py @@ -54,6 +54,7 @@ class to build, the runtime associated with that source and (optionally) the from torax.sources import source from torax.sources.impurity_radiation_heat_sink import impurity_radiation_constant_fraction from torax.sources.impurity_radiation_heat_sink import impurity_radiation_heat_sink +from torax.sources.impurity_radiation_heat_sink import impurity_radiation_mavrin_fit @dataclasses.dataclass(frozen=True) @@ -198,10 +199,14 @@ class SupportedSource: source_class=impurity_radiation_heat_sink.ImpurityRadiationHeatSink, model_functions={ impurity_radiation_heat_sink.ImpurityRadiationHeatSink.DEFAULT_MODEL_FUNCTION_NAME: ModelFunction( + source_profile_function=impurity_radiation_mavrin_fit.impurity_radiation_mavrin_fit, + runtime_params_class=impurity_radiation_mavrin_fit.RuntimeParams, + ), + impurity_radiation_constant_fraction.MODEL_FUNCTION_NAME: ModelFunction( source_profile_function=impurity_radiation_constant_fraction.radially_constant_fraction_of_Pin, runtime_params_class=impurity_radiation_constant_fraction.RuntimeParams, links_back=True, - ) + ), }, ), } diff --git a/torax/sources/tests/impurity_radiation_heat_sink_test.py b/torax/sources/tests/constant_fraction_impurity_radiation_heat_sink_test.py similarity index 98% rename from torax/sources/tests/impurity_radiation_heat_sink_test.py rename to torax/sources/tests/constant_fraction_impurity_radiation_heat_sink_test.py index 3d3b43d2..29806bc4 100644 --- a/torax/sources/tests/impurity_radiation_heat_sink_test.py +++ b/torax/sources/tests/constant_fraction_impurity_radiation_heat_sink_test.py @@ -52,7 +52,7 @@ def test_source_value(self): runtime_params_lib.Mode.MODEL_BASED ) if not source_lib.is_source_builder(impurity_radiation_sink_builder): - raise TypeError(f"{type(self)} has a bad _source_class_builder") + raise TypeError(f'{type(self)} has a bad _source_class_builder') # Source builder for generic_ion_el_heat_source # We don't test this class, as that should be done in its own test @@ -180,5 +180,5 @@ def test_extraction_of_relevant_profile_from_output(self): ) -if __name__ == "__main__": +if __name__ == '__main__': absltest.main() diff --git a/torax/sources/tests/mavrin_impurity_radiation_heat_sink_test.py b/torax/sources/tests/mavrin_impurity_radiation_heat_sink_test.py new file mode 100644 index 00000000..d774f8b2 --- /dev/null +++ b/torax/sources/tests/mavrin_impurity_radiation_heat_sink_test.py @@ -0,0 +1,359 @@ + +# Copyright 2024 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from absl.testing import absltest +from absl.testing import parameterized +import jax.numpy as jnp +import numpy as np +from torax.config import plasma_composition +from torax.config import runtime_params as general_runtime_params +from torax.config import runtime_params_slice +from torax.geometry import circular_geometry +from torax.sources import generic_ion_el_heat_source +from torax.sources import runtime_params as runtime_params_lib +from torax.sources import source as source_lib +from torax.sources import source_models as source_models_lib +from torax.sources.impurity_radiation_heat_sink import ( + impurity_radiation_heat_sink as impurity_radiation_heat_sink_lib, +) +from torax.sources.impurity_radiation_heat_sink import impurity_radiation_mavrin_fit +from torax.sources.tests import test_lib + + +class ImpurityRadiationMavrinFitTest(test_lib.SourceTestCase): + """Tests impurity_radiation_mavrin_fit implementation of ImpurityRadiationHeatSink.""" + + @classmethod + def setUpClass(cls): + super().setUpClass( + source_class=impurity_radiation_heat_sink_lib.ImpurityRadiationHeatSink, + runtime_params_class=impurity_radiation_mavrin_fit.RuntimeParams, + source_name=impurity_radiation_heat_sink_lib.ImpurityRadiationHeatSink.SOURCE_NAME, + model_func=impurity_radiation_mavrin_fit.impurity_radiation_mavrin_fit, + ) + + def test_source_value(self): + """Tests that the source value is correct.""" + # Source builder for this class + impurity_radiation_sink_builder = self._source_class_builder() + impurity_radiation_sink_builder.runtime_params.mode = ( + runtime_params_lib.Mode.MODEL_BASED + ) + if not source_lib.is_source_builder(impurity_radiation_sink_builder): + raise TypeError(f'{type(self)} has a bad _source_class_builder') + + # Source builder for generic_ion_el_heat_source + # We don't test this class, as that should be done in its own test + heat_source_builder_builder = source_lib.make_source_builder( + source_type=generic_ion_el_heat_source.GenericIonElectronHeatSource, + runtime_params_type=generic_ion_el_heat_source.RuntimeParams, + model_func=generic_ion_el_heat_source.default_formula, + ) + heat_source_builder = heat_source_builder_builder( + model_func=generic_ion_el_heat_source.default_formula + ) + + # Runtime params + runtime_params = general_runtime_params.GeneralRuntimeParams() + + # Source models + source_models_builder = source_models_lib.SourceModelsBuilder( + { + self._source_name: impurity_radiation_sink_builder, + generic_ion_el_heat_source.GenericIonElectronHeatSource.SOURCE_NAME: ( + heat_source_builder + ), + }, + ) + source_models = source_models_builder() + + # Extract the source we're testing and check that it's been built correctly + impurity_radiation_sink = source_models.sources[self._source_name] + self.assertIsInstance(impurity_radiation_sink, source_lib.Source) + + # Geometry, profiles, and dynamic runtime params + geo = circular_geometry.build_circular_geometry() + + dynamic_runtime_params_slice = ( + runtime_params_slice.DynamicRuntimeParamsSliceProvider( + runtime_params=runtime_params, + sources=source_models_builder.runtime_params, + torax_mesh=geo.torax_mesh, + )( + t=runtime_params.numerics.t_initial, + ) + ) + impurity_radiation_sink_dynamic_runtime_params_slice = ( + dynamic_runtime_params_slice.sources[self._source_name] + ) + + heat_source_dynamic_runtime_params_slice = ( + dynamic_runtime_params_slice.sources[ + generic_ion_el_heat_source.GenericIonElectronHeatSource.SOURCE_NAME + ] + ) + + assert isinstance( + impurity_radiation_sink_dynamic_runtime_params_slice, + impurity_radiation_mavrin_fit.DynamicRuntimeParams, + ) + assert isinstance( + heat_source_dynamic_runtime_params_slice, + generic_ion_el_heat_source.DynamicRuntimeParams, + ) + + def test_extraction_of_relevant_profile_from_output(self): + """Tests that the relevant profile is extracted from the output.""" + geo = circular_geometry.build_circular_geometry() + source_builder = self._source_class_builder() + source_models_builder = source_models_lib.SourceModelsBuilder( + {self._source_name: source_builder}, + ) + source_models = source_models_builder() + source = source_models.sources[self._source_name] + self.assertIsInstance(source, source_lib.Source) + cell = source_lib.get_cell_profile_shape(geo) + fake_profile = jnp.ones(cell) + # Check TEMP_EL is modified + np.testing.assert_allclose( + source.get_source_profile_for_affected_core_profile( + fake_profile, + source_lib.AffectedCoreProfile.TEMP_EL.value, + geo, + ), + jnp.ones(cell), + ) + # For unrelated states, this should just return all 0s. + np.testing.assert_allclose( + source.get_source_profile_for_affected_core_profile( + fake_profile, + source_lib.AffectedCoreProfile.NE.value, + geo, + ), + jnp.zeros(cell), + ) + + # pylint: disable=invalid-name + @parameterized.product( + ion_symbol=[ + ('C',), + ('N',), + ('O',), + ('Ne',), + ('Ar',), + ('Kr',), + ('Xe',), + ('W',), + ], + temperature=[0.1, 1.0, [10.0, 20.0], 90.0], + ) + def test_calculate_total_impurity_radiation_sanity( + self, ion_symbol, temperature + ): + """Test with valid ions and within temperature range.""" + Te = np.array(temperature) + ion_mixture = plasma_composition.DynamicIonMixture( + fractions=np.array([1.0]), + avg_A=2.0, # unused + ) + LZ_calculated = ( + impurity_radiation_mavrin_fit.calculate_total_impurity_radiation( + ion_symbol, + ion_mixture, + Te, + ) + ) + np.testing.assert_equal( + LZ_calculated.shape, + Te.shape, + err_msg=( + f'LZ and T shapes unequal for {ion_symbol} at temperature {Te}. LZ' + f' = {LZ_calculated}, LZ.shape = {LZ_calculated.shape}, Te.shape =' + f' {Te.shape}.' + ), + ) + # Physical sanity checking + np.testing.assert_array_less( + 0.0, + LZ_calculated, + err_msg=( + f'Unphysical negative LZ for {ion_symbol} at temperature {Te}. ' + f'LZ = {LZ_calculated}.' + ), + ) + + @parameterized.named_parameters( + ('Te_low_input', 0.05, 0.1), + ('Te_high_input', 150.0, 100.0), + ) + def test_temperature_clipping(self, Te_input, Te_clipped): + """Test with valid ions and within temperature range.""" + ion_symbol = ('W',) + ion_mixture = plasma_composition.DynamicIonMixture( + fractions=np.array([1.0]), + avg_A=2.0, # unused + ) + LZ_calculated = ( + impurity_radiation_mavrin_fit.calculate_total_impurity_radiation( + ion_symbol, + ion_mixture, + Te_input, + ) + ) + LZ_expected = ( + impurity_radiation_mavrin_fit.calculate_total_impurity_radiation( + ion_symbol, + ion_mixture, + Te_clipped, + ) + ) + + np.testing.assert_allclose( + LZ_calculated, + LZ_expected, + err_msg=( + f'Te clipping not working as expected for Te_input={Te_input},' + f' LZ_calculated = {LZ_calculated}, Z_expected={LZ_expected}' + ), + ) + + @parameterized.named_parameters( + ( + 'Helium-3', + {'He3': 1.0}, + [0.1, 2, 10], + [2.26267051e-36, 3.55291080e-36, 6.25952387e-36], + ), + ( + 'Helium-4', + {'He4': 1.0}, + [0.1, 2, 10], + [2.26267051e-36, 3.55291080e-36, 6.25952387e-36], + ), + ( + 'Lithium', + {'Li': 1.0}, + [0.1, 2, 10], + [1.37075024e-35, 9.16765402e-36, 1.60346076e-35], + ), + ( + 'Beryllium', + {'Be': 1.0}, + [0.1, 2, 10], + [6.86895406e-35, 1.88578938e-35, 3.04614535e-35], + ), + ( + 'Carbon', + {'C': 1.0}, + [0.1, 2, 10], + [6.74683566e-34, 5.89332177e-35, 7.94786067e-35], + ), + ( + 'Nitrogen', + {'N': 1.0}, + [0.1, 2, 10], + [6.97912189e-34, 9.68950644e-35, 1.15250226e-34], + ), + ( + 'Oxygen', + {'O': 1.0}, + [0.1, 2, 10], + [4.10676301e-34, 1.57469152e-34, 1.58599054e-34], + ), + ( + 'Neon', + {'Ne': 1.0}, + [0.1, 2, 10], + [1.19151664e-33, 3.27468464e-34, 2.82416557e-34], + ), + ( + 'Argon', + {'Ar': 1.0}, + [0.1, 2, 10], + [1.92265224e-32, 4.02388371e-33, 1.53295491e-33], + ), + ( + 'Krypton', + {'Kr': 1.0}, + [0.1, 2, 10], + [6.57654706e-32, 3.23512795e-32, 7.53285680e-33], + ), + ( + 'Xenon', + {'Xe': 1.0}, + [0.1, 2, 10], + [2.89734288e-31, 8.96916315e-32, 2.87740863e-32], + ), + ( + 'Tungsten', + {'W': 1.0}, + [0.1, 2, 10], + [1.66636258e-31, 4.46651033e-31, 1.31222935e-31], + ), + ( + 'Mixture', + { + 'He4': 0.2, + 'Li': 0.1, + 'Be': 0.197, + 'C': 0.1, + 'N': 0.1, + 'O': 0.1, + 'Ne': 0.1, + 'Ar': 0.1, + 'Kr': 0.001, + 'Xe': 0.001, + 'W': 0.001, + }, + [0.1, 2, 10], + [2.75762225e-33, 1.04050126e-33, 3.93256085e-34], + ), + ) + def test_calculate_total_impurity_radiation( + self, + species, + Te, + expected_LZ, + ): + """Test calculate_total_impurity_radiation. + + Args: + species: A dictionary of ion symbols and their fractions. + Te: The temperature in KeV. + expected_LZ: The expected effective cooling curve value. + + expected_LZ references were verified against plots in the Mavrin 2018 paper. + """ + Te = np.array(Te) + expected_LZ = np.array(expected_LZ) + avg_A = 2.0 # arbitrary, not used. + ion_symbols = tuple(species.keys()) + fractions = np.array(tuple(species.values())) + ion_mixture = plasma_composition.DynamicIonMixture( + fractions=fractions, + avg_A=avg_A, + ) + LZ_calculated = ( + impurity_radiation_mavrin_fit.calculate_total_impurity_radiation( + ion_symbols, + ion_mixture, + Te, + ) + ) + + np.testing.assert_allclose(LZ_calculated, expected_LZ, rtol=1e-5) + + +if __name__ == '__main__': + absltest.main() diff --git a/torax/tests/sim.py b/torax/tests/sim.py index cbf4489d..61ebc4f0 100644 --- a/torax/tests/sim.py +++ b/torax/tests/sim.py @@ -392,10 +392,10 @@ class SimTest(sim_test_case.SimTestCase): _ALL_PROFILES, 0, ), - # Predictor-corrector solver with simple impurity radiation + # Predictor-corrector solver with constant fraction of Pin radiation ( - 'test_iterhybrid_predictor_corrector_impurity_radiation', - 'test_iterhybrid_predictor_corrector_impurity_radiation.py', + 'test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation', + 'test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.py', _ALL_PROFILES, 0, ), diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_impurity_radiation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc similarity index 100% rename from torax/tests/test_data/test_iterhybrid_predictor_corrector_impurity_radiation.nc rename to torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_impurity_radiation.py b/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.py similarity index 94% rename from torax/tests/test_data/test_iterhybrid_predictor_corrector_impurity_radiation.py rename to torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.py index 3370991c..d193dd05 100644 --- a/torax/tests/test_data/test_iterhybrid_predictor_corrector_impurity_radiation.py +++ b/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.py @@ -21,5 +21,6 @@ CONFIG = copy.deepcopy(test_iterhybrid_predictor_corrector.CONFIG) CONFIG['sources']['impurity_radiation_heat_sink'] = { + 'model_func': 'radially_constant_fraction_of_Pin', 'fraction_of_total_power_density': 0.1 } diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc new file mode 100644 index 00000000..47321f75 Binary files /dev/null and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.py b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.py new file mode 100644 index 00000000..ddb7735c --- /dev/null +++ b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.py @@ -0,0 +1,32 @@ +# Copyright 2024 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Identical to test_iterhybrid_predictor_corrector but with radiation from Mavrin.""" + +import copy + +from torax.tests.test_data import test_iterhybrid_predictor_corrector + +# Set W to Neon ratio +# pylint: disable=invalid-name +W_frac = 1e-3 +CONFIG = copy.deepcopy(test_iterhybrid_predictor_corrector.CONFIG) +CONFIG['runtime_params']['plasma_composition']['impurity'] = { + 'Ne': 1 - W_frac, + 'W': W_frac, +} +CONFIG['runtime_params']['plasma_composition']['Zeff'] = 3.0 +CONFIG['sources']['impurity_radiation_heat_sink'] = { + 'model_func': 'impurity_radiation_mavrin_fit', +}