diff --git a/torax/_src/edge/collisional_radiative_models.py b/torax/_src/edge/collisional_radiative_models.py index da3e1729b..5a029917b 100644 --- a/torax/_src/edge/collisional_radiative_models.py +++ b/torax/_src/edge/collisional_radiative_models.py @@ -81,10 +81,11 @@ def calculate_mavrin_2017( ion_symbol_lookup = ion_symbol if ion_symbol_lookup not in cr_module.COEFFS.keys(): - raise ValueError( - f'Invalid ion symbol: {ion_symbol}. Allowed symbols are:' - f' {cr_module.COEFFS.keys()}' - ) + # If the ion is not supported by the edge radiation model, we assume it + # negligibly contributes to the edge physics (radiation or Z_eff/dilution in + # the divertor). This is a good assumption for heavy impurities like W, + # which this case covers. This behaviour is silent to avoid log spam. + return jnp.zeros_like(T_e) # Mavrin 2017 formulas are constructed for [eV] temperature input T_e_ev = T_e * 1e3 diff --git a/torax/_src/edge/pydantic_model.py b/torax/_src/edge/pydantic_model.py index 7efc469c4..12e79d760 100644 --- a/torax/_src/edge/pydantic_model.py +++ b/torax/_src/edge/pydantic_model.py @@ -71,11 +71,6 @@ class ExtendedLengyelConfig(base.EdgeModelConfig): extended_lengyel_defaults.DIVERTOR_BROADENING_FACTOR ) ) - ratio_bpol_omp_to_bpol_avg: torax_pydantic.PositiveTimeVaryingScalar = ( - torax_pydantic.ValidatedDefault( - extended_lengyel_defaults.RATIO_BPOL_OMP_TO_BPOL_AVG - ) - ) sheath_heat_transmission_factor: torax_pydantic.PositiveTimeVaryingScalar = ( torax_pydantic.ValidatedDefault( extended_lengyel_defaults.SHEATH_HEAT_TRANSMISSION_FACTOR @@ -134,16 +129,15 @@ class ExtendedLengyelConfig(base.EdgeModelConfig): connection_length_divertor: ( torax_pydantic.PositiveTimeVaryingScalar | None ) = None - angle_of_incidence_target: torax_pydantic.PositiveTimeVaryingScalar = ( - torax_pydantic.ValidatedDefault( - extended_lengyel_defaults.ANGLE_OF_INCIDENCE_TARGET - ) + angle_of_incidence_target: torax_pydantic.PositiveTimeVaryingScalar | None = ( + None ) - toroidal_flux_expansion: torax_pydantic.PositiveTimeVaryingScalar = ( - torax_pydantic.ValidatedDefault( - extended_lengyel_defaults.TOROIDAL_FLUX_EXPANSION - ) + toroidal_flux_expansion: torax_pydantic.PositiveTimeVaryingScalar | None = ( + None ) + ratio_bpol_omp_to_bpol_avg: ( + torax_pydantic.PositiveTimeVaryingScalar | None + ) = None # Optional input for inverse mode T_e_target: torax_pydantic.PositiveTimeVaryingScalar | None = None @@ -343,7 +337,6 @@ def _get_optional_value( newton_raphson_tol=self.newton_raphson_tol, ne_tau=self.ne_tau.get_value(t), divertor_broadening_factor=self.divertor_broadening_factor.get_value(t), - ratio_bpol_omp_to_bpol_avg=self.ratio_bpol_omp_to_bpol_avg.get_value(t), sheath_heat_transmission_factor=self.sheath_heat_transmission_factor.get_value( t ), @@ -367,8 +360,15 @@ def _get_optional_value( connection_length_divertor=_get_optional_value( self.connection_length_divertor, t ), - angle_of_incidence_target=self.angle_of_incidence_target.get_value(t), - toroidal_flux_expansion=self.toroidal_flux_expansion.get_value(t), + angle_of_incidence_target=_get_optional_value( + self.angle_of_incidence_target, t + ), + toroidal_flux_expansion=_get_optional_value( + self.toroidal_flux_expansion, t + ), + ratio_bpol_omp_to_bpol_avg=_get_optional_value( + self.ratio_bpol_omp_to_bpol_avg, t + ), seed_impurity_weights=seed_impurity_weights, fixed_impurity_concentrations=fixed_impurity_concentrations, enrichment_factor=enrichment_factor, diff --git a/torax/_src/edge/tests/edge_pydantic_model_test.py b/torax/_src/edge/tests/edge_pydantic_model_test.py index f2555ae60..981c4385e 100644 --- a/torax/_src/edge/tests/edge_pydantic_model_test.py +++ b/torax/_src/edge/tests/edge_pydantic_model_test.py @@ -379,6 +379,9 @@ def test_run_standalone_from_pydantic_config(self): 'fixed_impurity_concentrations': {'He': 0.01}, 'connection_length_target': 20.0, 'connection_length_divertor': 5.0, + 'angle_of_incidence_target': 3.0, + 'toroidal_flux_expansion': 1.0, + 'ratio_bpol_omp_to_bpol_avg': 4.0/3.0, 'computation_mode': 'inverse', 'solver_mode': 'fixed_point', 'use_enrichment_model': False, diff --git a/torax/_src/edge/tests/extended_lengyel_model_test.py b/torax/_src/edge/tests/extended_lengyel_model_test.py index 96a96d416..b6bbd0bd3 100644 --- a/torax/_src/edge/tests/extended_lengyel_model_test.py +++ b/torax/_src/edge/tests/extended_lengyel_model_test.py @@ -159,6 +159,9 @@ def test_call_inverse_mode(self): T_e_target=2.34, connection_length_target=20.0, connection_length_divertor=5.0, + angle_of_incidence_target=3.0, + ratio_bpol_omp_to_bpol_avg=4.0/3.0, + toroidal_flux_expansion=1.0, seed_impurity_weights={'N': 1.0, 'Ar': 0.05}, fixed_impurity_concentrations={'He': 0.01}, enrichment_factor={'N': 1.0, 'Ar': 1.0, 'He': 1.0}, @@ -301,6 +304,7 @@ def test_fixed_impurity_update_from_core(self, mock_run_standalone): connection_length_divertor=5.0, angle_of_incidence_target=1.0, toroidal_flux_expansion=1.0, + ratio_bpol_omp_to_bpol_avg=1.0, use_enrichment_model=False, diverted=True, ) @@ -417,6 +421,7 @@ def test_broadening_limited_logic( connection_length_divertor=5.0, angle_of_incidence_target=1.0, toroidal_flux_expansion=1.0, + ratio_bpol_omp_to_bpol_avg=1.0, use_enrichment_model=True, diverted=diverted, divertor_broadening_factor=_CONFIG_BROADENING, @@ -631,6 +636,7 @@ def test_edge_model_coupling_smoke(self): 'connection_length_divertor': 10.0, 'toroidal_flux_expansion': 4.0, 'angle_of_incidence_target': 3.0, + 'ratio_bpol_omp_to_bpol_avg': 4.0 / 3.0, 'use_enrichment_model': False, 'diverted': True, }, @@ -691,6 +697,7 @@ def test_temperature_boundary_condition_updates( 'connection_length_divertor': 10.0, 'toroidal_flux_expansion': 4.0, 'angle_of_incidence_target': 3.0, + 'ratio_bpol_omp_to_bpol_avg': 4.0 / 3.0, # Test parameters 'update_temperatures': update_temperatures, 'T_i_T_e_ratio_target': ion_to_electron_ratio, @@ -789,6 +796,7 @@ def test_inverse_mode_updates_impurities(self): 'connection_length_divertor': 15.0, 'toroidal_flux_expansion': 4.0, 'angle_of_incidence_target': 3.0, + 'ratio_bpol_omp_to_bpol_avg': 4.0 / 3.0, 'use_enrichment_model': False, 'diverted': True, }, @@ -881,6 +889,7 @@ def setUp(self): 'connection_length_divertor': 10.0, 'toroidal_flux_expansion': 4.0, 'angle_of_incidence_target': 3.0, + 'ratio_bpol_omp_to_bpol_avg': 4.0 / 3.0, 'diverted': True, # 'use_enrichment_model' and 'diverted' are overridden in tests. }, diff --git a/torax/_src/geometry/fbt.py b/torax/_src/geometry/fbt.py index 3aaf8b9c4..397d1be81 100644 --- a/torax/_src/geometry/fbt.py +++ b/torax/_src/geometry/fbt.py @@ -479,8 +479,8 @@ def _get_val(key): connection_length_target = _get_val('Lpar_target') # Lpar_div -> connection_length_divertor [m] connection_length_divertor = _get_val('Lpar_div') - # alpha_target -> angle_of_incidence_target [degrees] - angle_of_incidence_target = _get_val('alpha_target') + # alpha_target [radians] -> angle_of_incidence_target [degrees] + angle_of_incidence_target = np.rad2deg(_get_val('alpha_target')) # r_OMP -> R_OMP [m] R_OMP = _get_val('r_OMP') # r_target -> R_target [m] diff --git a/torax/_src/geometry/tests/fbt_test.py b/torax/_src/geometry/tests/fbt_test.py index 8d88df7d6..ec4220e5e 100644 --- a/torax/_src/geometry/tests/fbt_test.py +++ b/torax/_src/geometry/tests/fbt_test.py @@ -205,7 +205,7 @@ def test_fbt_edge_parameters( # Domain 1 (Upper): [15, 25] over time for Lpar_target. LY['Lpar_target'] = np.array([[10.0, 20.0], [15.0, 25.0]]) LY['Lpar_div'] = np.array([[1.0, 2.0], [1.5, 2.5]]) - LY['alpha_target'] = np.array([[3.0, 4.0], [3.5, 4.5]]) + LY['alpha_target'] = np.deg2rad(np.array([[3.0, 4.0], [3.5, 4.5]])) LY['r_OMP'] = np.array([[6.0, 7.0], [6.5, 7.5]]) LY['r_target'] = np.array([[5.0, 6.0], [5.5, 6.5]]) LY['Bp_OMP'] = np.array([[0.5, 0.7], [0.6, 0.8]]) diff --git a/torax/_src/output_tools/output.py b/torax/_src/output_tools/output.py index ebdc546ec..d7decda72 100644 --- a/torax/_src/output_tools/output.py +++ b/torax/_src/output_tools/output.py @@ -839,9 +839,20 @@ def _save_edge_outputs(self) -> xr.Dataset: xr_dict["solver_iterations"] = self._pack_into_data_array( "solver_iterations", numerics.iterations ) - xr_dict["solver_residual"] = self._pack_into_data_array( - "solver_residual", numerics.residual - ) + # Handle solver_residual explicitly because it is a vector + # (time, n_unknowns) which _pack_into_data_array doesn't support + # automatically. + if numerics.residual is not None and numerics.residual.ndim == 2: + xr_dict["solver_residual"] = xr.DataArray( + numerics.residual, + dims=[TIME, "solver_unknown_idx"], + name="solver_residual", + ) + else: + xr_dict["solver_residual"] = self._pack_into_data_array( + "solver_residual", numerics.residual + ) + xr_dict["solver_error"] = self._pack_into_data_array( "solver_error", numerics.error ) diff --git a/torax/_src/torax_pydantic/model_config.py b/torax/_src/torax_pydantic/model_config.py index 0a4141223..73c408412 100644 --- a/torax/_src/torax_pydantic/model_config.py +++ b/torax/_src/torax_pydantic/model_config.py @@ -274,7 +274,10 @@ def _validate_edge_core_impurity_consistency(self) -> typing_extensions.Self: raise ValueError( 'Mismatch between core plasma composition impurities and edge' f' impurities. Core: {core_species}, Edge: {edge_species}.' - f' Difference: {core_species.symmetric_difference(edge_species)}' + f' Difference: {core_species.symmetric_difference(edge_species)}.' + ' Likely reason: edge.fixed_impurity_concentrations and/or' + ' edge.seed_impurity_weights do not match plasma_composition.' + 'impurity.species.' ) return self diff --git a/torax/tests/sim_test.py b/torax/tests/sim_test.py index f647fdcdc..8a1be7368 100644 --- a/torax/tests/sim_test.py +++ b/torax/tests/sim_test.py @@ -226,6 +226,11 @@ class SimTest(sim_test_case.SimTestCase): 'test_iterhybrid_predictor_corrector_mavrin_n_e_ratios', 'test_iterhybrid_predictor_corrector_mavrin_n_e_ratios.py', ), + # Predictor-corrector w/ Mavrin, n_e_ratios, forward lengyel + ( + 'test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel', + 'test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.py', + ), # Predictor-corrector with Mavrin and n_e_ratios_Z_eff impurity mode. ( 'test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_z_eff', diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.nc new file mode 100644 index 000000000..32e702a86 Binary files /dev/null and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.py b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.py new file mode 100644 index 000000000..714ee47ec --- /dev/null +++ b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.py @@ -0,0 +1,46 @@ +# 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. + +"""test_iterhybrid_predictor_corrector with Mavrin radiation, n_e_ratios, x-lengyel.""" + +import copy + +from torax.tests.test_data import test_iterhybrid_predictor_corrector + +CONFIG = copy.deepcopy(test_iterhybrid_predictor_corrector.CONFIG) + +assert isinstance(CONFIG['plasma_composition'], dict) +# If present, Z_eff would be ignored in the 'n_e_ratios' impurity mode. +del CONFIG['plasma_composition']['Z_eff'] +CONFIG['plasma_composition']['impurity'] = { + 'impurity_mode': 'n_e_ratios', + 'species': {'Ne': 0.01, 'W': 1e-5}, +} +CONFIG['sources']['impurity_radiation'] = { + 'model_name': 'mavrin_fit', +} + +CONFIG['edge'] = { + 'model_name': 'extended_lengyel', + 'computation_mode': 'forward', + 'fixed_impurity_concentrations': {'Ne': 0.01, 'W': 1e-5}, + 'impurity_sot': 'core', + 'connection_length_target': 30.0, + 'connection_length_divertor': 10.0, + 'angle_of_incidence_target': 3.0, + 'ratio_bpol_omp_to_bpol_avg': 4.0 / 3.0, + 'toroidal_flux_expansion': 1.0, + 'use_enrichment_model': True, + 'diverted': True, +}