-
Notifications
You must be signed in to change notification settings - Fork 39
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Impurity radiation from Mavrin polynomial fits.
Implementation of ADAS polynomial fit of impurity line radiation from A. A. Mavrin (2018). PiperOrigin-RevId: 721109765
- Loading branch information
Showing
13 changed files
with
699 additions
and
20 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
246 changes: 246 additions & 0 deletions
246
torax/sources/impurity_radiation_heat_sink/impurity_radiation_mavrin_fit.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.