diff --git a/torax/config/tests/runtime_params_slice.py b/torax/config/tests/runtime_params_slice.py index 4bd6a6ab..a71682fa 100644 --- a/torax/config/tests/runtime_params_slice.py +++ b/torax/config/tests/runtime_params_slice.py @@ -168,16 +168,20 @@ def test_source_formula_config_has_time_dependent_params(self): dcs = runtime_params_slice_lib.DynamicRuntimeParamsSliceProvider( runtime_params=runtime_params, sources={ - 'gas_puff_source': electron_density_sources.GasPuffRuntimeParams( - puff_decay_length={0.0: 0.0, 1.0: 4.0}, - S_puff_tot={0.0: 0.0, 1.0: 5.0}, + electron_density_sources.GAS_PUFF_SOURCE_NAME: ( + electron_density_sources.GasPuffRuntimeParams( + puff_decay_length={0.0: 0.0, 1.0: 4.0}, + S_puff_tot={0.0: 0.0, 1.0: 5.0}, + ) ), - 'pellet_source': electron_density_sources.PelletRuntimeParams( - pellet_width={0.0: 0.0, 1.0: 1.0}, - pellet_deposition_location={0.0: 0.0, 1.0: 2.0}, - S_pellet_tot={0.0: 0.0, 1.0: 3.0}, + electron_density_sources.PELLET_SOURCE_NAME: ( + electron_density_sources.PelletRuntimeParams( + pellet_width={0.0: 0.0, 1.0: 1.0}, + pellet_deposition_location={0.0: 0.0, 1.0: 2.0}, + S_pellet_tot={0.0: 0.0, 1.0: 3.0}, + ) ), - 'nbi_particle_source': ( + electron_density_sources.GENERIC_PARTICLE_SOURCE_NAME: ( electron_density_sources.NBIParticleRuntimeParams( nbi_particle_width={0.0: 0.0, 1.0: 6.0}, nbi_deposition_location={0.0: 0.0, 1.0: 7.0}, @@ -189,9 +193,13 @@ def test_source_formula_config_has_time_dependent_params(self): )( t=0.5, ) - pellet_source = dcs.sources['pellet_source'] - gas_puff_source = dcs.sources['gas_puff_source'] - nbi_particle_source = dcs.sources['nbi_particle_source'] + pellet_source = dcs.sources[electron_density_sources.PELLET_SOURCE_NAME] + gas_puff_source = dcs.sources[ + electron_density_sources.GAS_PUFF_SOURCE_NAME + ] + nbi_particle_source = dcs.sources[ + electron_density_sources.GENERIC_PARTICLE_SOURCE_NAME + ] assert isinstance( pellet_source, electron_density_sources.DynamicPelletRuntimeParams, @@ -222,11 +230,13 @@ def test_source_formula_config_has_time_dependent_params(self): dcs = runtime_params_slice_lib.DynamicRuntimeParamsSliceProvider( runtime_params=runtime_params, sources={ - 'gas_puff_source': sources_params_lib.RuntimeParams( - formula=formula_config.Exponential( - total={0.0: 0.0, 1.0: 1.0}, - c1={0.0: 0.0, 1.0: 2.0}, - c2={0.0: 0.0, 1.0: 3.0}, + electron_density_sources.GAS_PUFF_SOURCE_NAME: ( + sources_params_lib.RuntimeParams( + formula=formula_config.Exponential( + total={0.0: 0.0, 1.0: 1.0}, + c1={0.0: 0.0, 1.0: 2.0}, + c2={0.0: 0.0, 1.0: 3.0}, + ) ) ), }, @@ -234,29 +244,29 @@ def test_source_formula_config_has_time_dependent_params(self): )( t=0.25, ) - gas_puff_source = dcs.sources['gas_puff_source'] + gas_puff_source = dcs.sources[ + electron_density_sources.GAS_PUFF_SOURCE_NAME + ] assert isinstance( gas_puff_source.formula, formula_config.DynamicExponential, ) - np.testing.assert_allclose( - gas_puff_source.formula.total, 0.25 - ) + np.testing.assert_allclose(gas_puff_source.formula.total, 0.25) np.testing.assert_allclose(gas_puff_source.formula.c1, 0.5) - np.testing.assert_allclose( - gas_puff_source.formula.c2, 0.75 - ) + np.testing.assert_allclose(gas_puff_source.formula.c2, 0.75) with self.subTest('gaussian_formula'): runtime_params = general_runtime_params.GeneralRuntimeParams() dcs = runtime_params_slice_lib.DynamicRuntimeParamsSliceProvider( runtime_params=runtime_params, sources={ - 'gas_puff_source': sources_params_lib.RuntimeParams( - formula=formula_config.Gaussian( - total={0.0: 0.0, 1.0: 1.0}, - c1={0.0: 0.0, 1.0: 2.0}, - c2={0.0: 0.0, 1.0: 3.0}, + electron_density_sources.GAS_PUFF_SOURCE_NAME: ( + sources_params_lib.RuntimeParams( + formula=formula_config.Gaussian( + total={0.0: 0.0, 1.0: 1.0}, + c1={0.0: 0.0, 1.0: 2.0}, + c2={0.0: 0.0, 1.0: 3.0}, + ) ) ), }, @@ -264,17 +274,13 @@ def test_source_formula_config_has_time_dependent_params(self): )( t=0.25, ) - gas_puff_source = dcs.sources['gas_puff_source'] - assert isinstance( - gas_puff_source.formula, formula_config.DynamicGaussian - ) - np.testing.assert_allclose( - gas_puff_source.formula.total, 0.25 - ) + gas_puff_source = dcs.sources[ + electron_density_sources.GAS_PUFF_SOURCE_NAME + ] + assert isinstance(gas_puff_source.formula, formula_config.DynamicGaussian) + np.testing.assert_allclose(gas_puff_source.formula.total, 0.25) np.testing.assert_allclose(gas_puff_source.formula.c1, 0.5) - np.testing.assert_allclose( - gas_puff_source.formula.c2, 0.75 - ) + np.testing.assert_allclose(gas_puff_source.formula.c2, 0.75) def test_wext_in_dynamic_runtime_params_cannot_be_negative(self): """Tests that wext cannot be negative.""" @@ -283,30 +289,34 @@ def test_wext_in_dynamic_runtime_params_cannot_be_negative(self): runtime_params=runtime_params, transport=transport_params_lib.RuntimeParams(), sources={ - 'jext': external_current_source.RuntimeParams( - wext={0.0: 1.0, 1.0: -1.0} + external_current_source.SOURCE_NAME: ( + external_current_source.RuntimeParams( + wext={0.0: 1.0, 1.0: -1.0} + ) ), }, stepper=stepper_params_lib.RuntimeParams(), torax_mesh=self._geo.torax_mesh, ) # While wext is positive, this should be fine. - dcs = dcs_provider(t=0.0,) - jext = dcs.sources['jext'] - assert isinstance( - jext, external_current_source.DynamicRuntimeParams + dcs = dcs_provider( + t=0.0, ) + jext = dcs.sources[external_current_source.SOURCE_NAME] + assert isinstance(jext, external_current_source.DynamicRuntimeParams) np.testing.assert_allclose(jext.wext, 1.0) # Even 0 should be fine. - dcs = dcs_provider(t=0.5,) - jext = dcs.sources['jext'] - assert isinstance( - jext, external_current_source.DynamicRuntimeParams + dcs = dcs_provider( + t=0.5, ) + jext = dcs.sources[external_current_source.SOURCE_NAME] + assert isinstance(jext, external_current_source.DynamicRuntimeParams) np.testing.assert_allclose(jext.wext, 0.0) # But negative values will cause an error. with self.assertRaises(RuntimeError): - dcs_provider(t=1.0,) + dcs_provider( + t=1.0, + ) @parameterized.parameters( ( @@ -471,7 +481,9 @@ def test_update_dynamic_slice_provider_updates_sources( """Tests that the dynamic slice provider can be updated.""" runtime_params = general_runtime_params.GeneralRuntimeParams() source_models_builder = default_sources.get_default_sources_builder() - source_models_builder.runtime_params['jext'].Iext = 1.0 + source_models_builder.runtime_params[ + external_current_source.SOURCE_NAME + ].Iext = 1.0 geo = geometry.build_circular_geometry(n_rho=4) provider = runtime_params_slice_lib.DynamicRuntimeParamsSliceProvider( runtime_params=runtime_params, @@ -485,7 +497,9 @@ def test_update_dynamic_slice_provider_updates_sources( self.assertIn(key, dcs.sources) # Update an interpolated variable. - source_models_builder.runtime_params['jext'].Iext = 2.0 + source_models_builder.runtime_params[ + external_current_source.SOURCE_NAME + ].Iext = 2.0 # Check pre-update that nothing has changed. dcs = provider( @@ -493,7 +507,7 @@ def test_update_dynamic_slice_provider_updates_sources( ) for key in source_models_builder.runtime_params.keys(): self.assertIn(key, dcs.sources) - jext_source = dcs.sources['jext'] + jext_source = dcs.sources[external_current_source.SOURCE_NAME] assert isinstance(jext_source, external_current_source.DynamicRuntimeParams) self.assertEqual(jext_source.Iext, 1.0) @@ -508,7 +522,7 @@ def test_update_dynamic_slice_provider_updates_sources( ) for key in source_models_builder.runtime_params.keys(): self.assertIn(key, dcs.sources) - jext_source = dcs.sources['jext'] + jext_source = dcs.sources[external_current_source.SOURCE_NAME] assert isinstance(jext_source, external_current_source.DynamicRuntimeParams) self.assertEqual(jext_source.Iext, 2.0) diff --git a/torax/sources/bootstrap_current_source.py b/torax/sources/bootstrap_current_source.py index 13f4c9d1..8a00d9dd 100644 --- a/torax/sources/bootstrap_current_source.py +++ b/torax/sources/bootstrap_current_source.py @@ -34,6 +34,9 @@ from torax.sources import source_profiles +SOURCE_NAME = 'j_bootstrap' + + @dataclasses.dataclass(kw_only=True) class RuntimeParams(runtime_params_lib.RuntimeParams): """Configuration parameters for the bootstrap current source.""" @@ -184,7 +187,7 @@ def get_source_profile_for_affected_core_profile( ) -> jax.Array: return jnp.where( affected_core_profile in self.affected_core_profiles_ints, - profile['j_bootstrap'], + profile[SOURCE_NAME], jnp.zeros_like(geo.rho), ) diff --git a/torax/sources/bremsstrahlung_heat_sink.py b/torax/sources/bremsstrahlung_heat_sink.py index 470cbb21..9e9928cd 100644 --- a/torax/sources/bremsstrahlung_heat_sink.py +++ b/torax/sources/bremsstrahlung_heat_sink.py @@ -31,6 +31,9 @@ from torax.sources import source_models +SOURCE_NAME = 'bremsstrahlung_heat_sink' + + @dataclasses.dataclass(kw_only=True) class RuntimeParams(runtime_params_lib.RuntimeParams): use_relativistic_correction: bool = False diff --git a/torax/sources/electron_density_sources.py b/torax/sources/electron_density_sources.py index 37fc8cd3..6922ed3a 100644 --- a/torax/sources/electron_density_sources.py +++ b/torax/sources/electron_density_sources.py @@ -94,6 +94,9 @@ def _calc_puff_source( ) +GAS_PUFF_SOURCE_NAME = 'gas_puff_source' + + @dataclasses.dataclass(kw_only=True, frozen=True, eq=True) class GasPuffSource(source.Source): """Gas puff source for the ne equation.""" @@ -183,6 +186,9 @@ def _calc_nbi_source( ) +GENERIC_PARTICLE_SOURCE_NAME = 'nbi_particle_source' + + @dataclasses.dataclass(kw_only=True, frozen=True, eq=True) class NBIParticleSource(source.Source): """Neutral-beam injection source for the ne equation.""" @@ -257,6 +263,9 @@ def _calc_pellet_source( ) +PELLET_SOURCE_NAME = 'pellet_source' + + @dataclasses.dataclass(kw_only=True, frozen=True, eq=True) class PelletSource(source.Source): """Pellet source for the ne equation.""" diff --git a/torax/sources/external_current_source.py b/torax/sources/external_current_source.py index 539b1278..0efdaac3 100644 --- a/torax/sources/external_current_source.py +++ b/torax/sources/external_current_source.py @@ -34,6 +34,7 @@ from torax.sources import source +SOURCE_NAME = 'jext' # pylint: disable=invalid-name diff --git a/torax/sources/fusion_heat_source.py b/torax/sources/fusion_heat_source.py index faf2b679..d44cd699 100644 --- a/torax/sources/fusion_heat_source.py +++ b/torax/sources/fusion_heat_source.py @@ -30,6 +30,9 @@ from torax.sources import source +SOURCE_NAME = 'fusion_heat_source' + + def calc_fusion( geo: geometry.Geometry, core_profiles: state.CoreProfiles, diff --git a/torax/sources/generic_ion_el_heat_source.py b/torax/sources/generic_ion_el_heat_source.py index 41fdd41b..4d5859fd 100644 --- a/torax/sources/generic_ion_el_heat_source.py +++ b/torax/sources/generic_ion_el_heat_source.py @@ -31,6 +31,7 @@ from torax.sources import source +SOURCE_NAME = 'generic_ion_el_heat_source' # Many variables throughout this function are capitalized based on physics # notational conventions rather than on Google Python style # pylint: disable=invalid-name diff --git a/torax/sources/ohmic_heat_source.py b/torax/sources/ohmic_heat_source.py index a11a859b..c122a3b4 100644 --- a/torax/sources/ohmic_heat_source.py +++ b/torax/sources/ohmic_heat_source.py @@ -31,6 +31,9 @@ from torax.sources import source_models as source_models_lib +SOURCE_NAME = 'ohmic_heat_source' + + @functools.partial( jax_utils.jit, static_argnames=[ diff --git a/torax/sources/qei_source.py b/torax/sources/qei_source.py index b5c041d9..09b395e3 100644 --- a/torax/sources/qei_source.py +++ b/torax/sources/qei_source.py @@ -30,6 +30,7 @@ from torax.sources import source_profiles +SOURCE_NAME = 'qei_source' # pylint: disable=invalid-name diff --git a/torax/sources/register_source.py b/torax/sources/register_source.py index 3b33a56c..a7ea8862 100644 --- a/torax/sources/register_source.py +++ b/torax/sources/register_source.py @@ -109,53 +109,53 @@ def get_registered_source(source_name: str) -> RegisteredSource: def register_torax_sources(): """Register a set of sources commonly used in TORAX.""" register_new_source( - 'j_bootstrap', + bootstrap_current_source.SOURCE_NAME, source_class=bootstrap_current_source.BootstrapCurrentSource, default_runtime_params_class=bootstrap_current_source.RuntimeParams, ) register_new_source( - 'jext', + external_current_source.SOURCE_NAME, external_current_source.ExternalCurrentSource, default_runtime_params_class=external_current_source.RuntimeParams, ) register_new_source( - 'nbi_particle_source', + electron_density_sources.GENERIC_PARTICLE_SOURCE_NAME, electron_density_sources.NBIParticleSource, default_runtime_params_class=electron_density_sources.NBIParticleRuntimeParams, ) register_new_source( - 'gas_puff_source', + electron_density_sources.GAS_PUFF_SOURCE_NAME, electron_density_sources.GasPuffSource, default_runtime_params_class=electron_density_sources.GasPuffRuntimeParams, ) register_new_source( - 'pellet_source', + electron_density_sources.PELLET_SOURCE_NAME, electron_density_sources.PelletSource, default_runtime_params_class=electron_density_sources.PelletRuntimeParams, ) register_new_source( - 'generic_ion_el_heat_source', + ion_el_heat.SOURCE_NAME, ion_el_heat.GenericIonElectronHeatSource, default_runtime_params_class=ion_el_heat.RuntimeParams, ) register_new_source( - 'fusion_heat_source', + fusion_heat_source.SOURCE_NAME, fusion_heat_source.FusionHeatSource, default_runtime_params_class=fusion_heat_source.FusionHeatSourceRuntimeParams ) register_new_source( - 'qei_source', + qei_source.SOURCE_NAME, qei_source.QeiSource, default_runtime_params_class=qei_source.RuntimeParams, ) register_new_source( - 'ohmic_heat_source', + ohmic_heat_source.SOURCE_NAME, ohmic_heat_source.OhmicHeatSource, default_runtime_params_class=ohmic_heat_source.OhmicRuntimeParams, links_back=True, ) register_new_source( - 'bremsstrahlung_heat_sink', + bremsstrahlung_heat_sink.SOURCE_NAME, bremsstrahlung_heat_sink.BremsstrahlungHeatSink, default_runtime_params_class=bremsstrahlung_heat_sink.RuntimeParams, ) diff --git a/torax/sources/source_models.py b/torax/sources/source_models.py index cbf87e00..8dd158f3 100644 --- a/torax/sources/source_models.py +++ b/torax/sources/source_models.py @@ -577,11 +577,11 @@ def __init__( # Some sources are accessed for specific use cases, so we extract those # ones and expose them directly. self._j_bootstrap = None - self._j_bootstrap_name = 'j_bootstrap' # default, can be overridden below. + self._j_bootstrap_name = bootstrap_current_source.SOURCE_NAME self._jext = None - self._jext_name = 'jext' # default, can be overridden below. + self._jext_name = external_current_source.SOURCE_NAME self._qei_source = None - self._qei_source_name = 'qei_source' # default, can be overridden below. + self._qei_source_name = qei_source_lib.SOURCE_NAME # The rest of the sources are "standard". self._standard_sources = {} @@ -811,9 +811,19 @@ def __init__( source_builders = source_builders or {} # Validate that these sources are found - bootstrap_found = False if 'j_bootstrap' not in source_builders else True - qei_found = False if 'qei_source' not in source_builders else True - jext_found = False if 'jext' not in source_builders else True + bootstrap_found = ( + False + if bootstrap_current_source.SOURCE_NAME not in source_builders + else True + ) + qei_found = ( + False if qei_source_lib.SOURCE_NAME not in source_builders else True + ) + jext_found = ( + False + if external_current_source.SOURCE_NAME not in source_builders + else True + ) # These are special sources that must be present for every TORAX run. # If these sources are missing, we need to include builders for them. @@ -821,27 +831,35 @@ def __init__( # The SourceModels would also build them, but then there'd be no # user-editable runtime params for them. if not bootstrap_found: - source_builders['j_bootstrap'] = source_lib.make_source_builder( - bootstrap_current_source.BootstrapCurrentSource, - runtime_params_type=bootstrap_current_source.RuntimeParams, - )() - source_builders['j_bootstrap'].runtime_params.mode = ( - runtime_params_lib.Mode.ZERO + source_builders[bootstrap_current_source.SOURCE_NAME] = ( + source_lib.make_source_builder( + bootstrap_current_source.BootstrapCurrentSource, + runtime_params_type=bootstrap_current_source.RuntimeParams, + )() ) + source_builders[ + bootstrap_current_source.SOURCE_NAME + ].runtime_params.mode = runtime_params_lib.Mode.ZERO if not qei_found: - source_builders['qei_source'] = source_lib.make_source_builder( - qei_source_lib.QeiSource, - runtime_params_type=qei_source_lib.RuntimeParams, - )() - source_builders['qei_source'].runtime_params.mode = ( + source_builders[qei_source_lib.SOURCE_NAME] = ( + source_lib.make_source_builder( + qei_source_lib.QeiSource, + runtime_params_type=qei_source_lib.RuntimeParams, + )() + ) + source_builders[qei_source_lib.SOURCE_NAME].runtime_params.mode = ( runtime_params_lib.Mode.ZERO ) if not jext_found: - source_builders['jext'] = source_lib.make_source_builder( - external_current_source.ExternalCurrentSource, - runtime_params_type=external_current_source.RuntimeParams, - )() - source_builders['jext'].runtime_params.mode = runtime_params_lib.Mode.ZERO + source_builders[external_current_source.SOURCE_NAME] = ( + source_lib.make_source_builder( + external_current_source.ExternalCurrentSource, + runtime_params_type=external_current_source.RuntimeParams, + )() + ) + source_builders[ + external_current_source.SOURCE_NAME + ].runtime_params.mode = runtime_params_lib.Mode.ZERO self.source_builders = source_builders diff --git a/torax/sources/tests/bootstrap_current_source.py b/torax/sources/tests/bootstrap_current_source.py index 74d141bb..198cf921 100644 --- a/torax/sources/tests/bootstrap_current_source.py +++ b/torax/sources/tests/bootstrap_current_source.py @@ -48,10 +48,10 @@ def test_source_value(self): runtime_params = general_runtime_params.GeneralRuntimeParams() geo = geometry.build_circular_geometry() source_models_builder = source_models_lib.SourceModelsBuilder( - {'j_bootstrap': source_builder} + {bootstrap_current_source.SOURCE_NAME: source_builder} ) source_models = source_models_builder() - source = source_models.sources['j_bootstrap'] + source = source_models.sources[bootstrap_current_source.SOURCE_NAME] dynamic_runtime_params_slice = ( runtime_params_slice.DynamicRuntimeParamsSliceProvider( runtime_params, diff --git a/torax/sources/tests/external_current_source.py b/torax/sources/tests/external_current_source.py index 8c636efd..e0f1c2de 100644 --- a/torax/sources/tests/external_current_source.py +++ b/torax/sources/tests/external_current_source.py @@ -48,7 +48,7 @@ def test_source_value(self): dynamic_slice = runtime_params_slice.DynamicRuntimeParamsSliceProvider( runtime_params, sources={ - 'jext': source_builder.runtime_params, + external_current_source.SOURCE_NAME: source_builder.runtime_params, }, torax_mesh=geo.torax_mesh, )( @@ -59,14 +59,18 @@ def test_source_value(self): self.assertIsNotNone( source.get_value( dynamic_runtime_params_slice=dynamic_slice, - dynamic_source_runtime_params=dynamic_slice.sources['jext'], + dynamic_source_runtime_params=dynamic_slice.sources[ + external_current_source.SOURCE_NAME + ], geo=geo, ) ) self.assertIsNotNone( source.jext_hires( dynamic_runtime_params_slice=dynamic_slice, - dynamic_source_runtime_params=dynamic_slice.sources['jext'], + dynamic_source_runtime_params=dynamic_slice.sources[ + external_current_source.SOURCE_NAME + ], geo=geo, ) ) @@ -83,7 +87,9 @@ def test_invalid_source_types_raise_errors(self): dynamic_slice = runtime_params_slice.DynamicRuntimeParamsSliceProvider( runtime_params, sources={ - 'jext': source_builder.runtime_params, + external_current_source.SOURCE_NAME: ( + source_builder.runtime_params + ), }, torax_mesh=geo.torax_mesh, )( @@ -91,7 +97,9 @@ def test_invalid_source_types_raise_errors(self): ) source.get_value( dynamic_runtime_params_slice=dynamic_slice, - dynamic_source_runtime_params=dynamic_slice.sources['jext'], + dynamic_source_runtime_params=dynamic_slice.sources[ + external_current_source.SOURCE_NAME + ], geo=geo, ) @@ -108,7 +116,7 @@ def test_profile_is_on_face_grid(self): dynamic_runtime_params_slice = runtime_params_slice.DynamicRuntimeParamsSliceProvider( runtime_params, sources={ - 'jext': source_builder.runtime_params, + external_current_source.SOURCE_NAME: source_builder.runtime_params, }, torax_mesh=geo.torax_mesh, )( @@ -117,7 +125,9 @@ def test_profile_is_on_face_grid(self): self.assertEqual( source.get_value( dynamic_runtime_params_slice, - dynamic_runtime_params_slice.sources['jext'], + dynamic_runtime_params_slice.sources[ + external_current_source.SOURCE_NAME + ], geo, core_profiles=None, ).shape, diff --git a/torax/sources/tests/formulas.py b/torax/sources/tests/formulas.py index 74650489..76620a0c 100644 --- a/torax/sources/tests/formulas.py +++ b/torax/sources/tests/formulas.py @@ -23,8 +23,12 @@ from torax.config import numerics as numerics_lib from torax.config import profile_conditions as profile_conditions_lib from torax.config import runtime_params as general_runtime_params +from torax.sources import bremsstrahlung_heat_sink +from torax.sources import electron_density_sources from torax.sources import formula_config from torax.sources import formulas +from torax.sources import fusion_heat_source +from torax.sources import ohmic_heat_source from torax.sources import runtime_params as runtime_params_lib from torax.sources.tests import test_lib from torax.stepper import linear_theta_method @@ -68,27 +72,31 @@ def test_custom_exponential_source_can_replace_puff_source(self): ) # Set the sources to match test_particle_sources_constant as well. source_models_builder = default_sources.get_default_sources_builder() - source_models_builder.runtime_params['pellet_source'].S_pellet_tot = 2.0e22 + source_models_builder.runtime_params[ + electron_density_sources.PELLET_SOURCE_NAME + ].S_pellet_tot = 2.0e22 S_puff_tot = 1.0e22 # pylint: disable=invalid-name puff_decay_length = 0.05 - source_models_builder.runtime_params['gas_puff_source'].S_puff_tot = ( - S_puff_tot - ) source_models_builder.runtime_params[ - 'gas_puff_source' + electron_density_sources.GAS_PUFF_SOURCE_NAME + ].S_puff_tot = S_puff_tot + source_models_builder.runtime_params[ + electron_density_sources.GAS_PUFF_SOURCE_NAME ].puff_decay_length = puff_decay_length - source_models_builder.runtime_params['nbi_particle_source'].S_nbi_tot = 0.0 + source_models_builder.runtime_params[ + electron_density_sources.GENERIC_PARTICLE_SOURCE_NAME + ].S_nbi_tot = 0.0 # We need to turn off some other sources for test_particle_sources_constant # that are unrelated to our test for the ne custom source. - source_models_builder.runtime_params['fusion_heat_source'].mode = ( - runtime_params_lib.Mode.ZERO - ) - source_models_builder.runtime_params['ohmic_heat_source'].mode = ( - runtime_params_lib.Mode.ZERO - ) - source_models_builder.runtime_params['bremsstrahlung_heat_sink'].mode = ( + source_models_builder.runtime_params[ + fusion_heat_source.SOURCE_NAME + ].mode = runtime_params_lib.Mode.ZERO + source_models_builder.runtime_params[ohmic_heat_source.SOURCE_NAME].mode = ( runtime_params_lib.Mode.ZERO ) + source_models_builder.runtime_params[ + bremsstrahlung_heat_sink.SOURCE_NAME + ].mode = runtime_params_lib.Mode.ZERO # Add the custom source to the source_models, but keep it turned off for the # first run. @@ -172,9 +180,9 @@ def test_custom_exponential_source_can_replace_puff_source(self): ) ) # And turn off the gas puff source it is replacing. - source_models_builder.runtime_params['gas_puff_source'].mode = ( - runtime_params_lib.Mode.ZERO - ) + source_models_builder.runtime_params[ + electron_density_sources.GAS_PUFF_SOURCE_NAME + ].mode = runtime_params_lib.Mode.ZERO sim = simulation_app.update_sim( sim, test_particle_sources_constant_runtime_params,