From 0cc91ae71a65a9fb0df06534f27367f34b858d64 Mon Sep 17 00:00:00 2001 From: mn3981 Date: Tue, 5 May 2026 09:48:12 +0100 Subject: [PATCH] Add DivertorNumberModels enum and refactor i_single_null usage across models Co-authored-by: Copilot --- process/core/init.py | 11 ++++++----- process/models/build.py | 15 ++++++++++++--- process/models/divertor.py | 5 +++-- process/models/geometry/plasma.py | 3 ++- process/models/tfcoil/base.py | 3 ++- 5 files changed, 25 insertions(+), 12 deletions(-) diff --git a/process/core/init.py b/process/core/init.py index 989d04c19..2aa98cc29 100644 --- a/process/core/init.py +++ b/process/core/init.py @@ -32,6 +32,7 @@ init_superconducting_tf_coil_variables, ) from process.data_structure.tfcoil_variables import init_tfcoil_variables +from process.models.build import DivertorNumberModels from process.models.stellarator.initialization import st_init from process.models.superconductors import ( SuperconductorMaterial, @@ -569,14 +570,14 @@ def check_process(inputs, data): # noqa: ARG001 "REINKE IMPURITY MODEL: The Martin LH threshold scale is not being used and is recommended for the Reinke model", stacklevel=2, ) - - if data_structure.physics_variables.i_single_null == 0: + i_single_null = DivertorNumberModels(data_structure.physics_variables.i_single_null) + if i_single_null == DivertorNumberModels.DOUBLE_NULL: data.divertor.n_divertors = 2 data.build.dz_fw_plasma_gap = data.build.dz_xpoint_divertor data.build.dz_shld_upper = data.build.dz_shld_lower data.build.dz_vv_upper = data.build.dz_vv_lower warn("Double-null: Upper vertical build forced to match lower", stacklevel=2) - else: # i_single_null == 1 + else: # i_single_null == DivertorNumberModels.SINGLE_NULL data.divertor.n_divertors = 1 # Tight aspect ratio options (ST) @@ -681,7 +682,7 @@ def check_process(inputs, data): # noqa: ARG001 ) # Check if a single null divertor is used in double null machine - if data_structure.physics_variables.i_single_null == 0 and ( + if i_single_null == DivertorNumberModels.DOUBLE_NULL and ( data_structure.physics_variables.f_p_div_lower in {1.0, 0.0} ): warn("Operating with a single null in a double null machine", stacklevel=2) @@ -763,7 +764,7 @@ def check_process(inputs, data): # noqa: ARG001 raise ProcessValidationError( "More than 2 divertor coils (i_pf_location = 2) is not a valid configuration" ) - if data_structure.physics_variables.i_single_null == 1 and j < 2: + if i_single_null == DivertorNumberModels.SINGLE_NULL and j < 2: raise ProcessValidationError( "If i_single_null=1, use 2 individual divertor coils (i_pf_location = 2, 2; n_pf_coils_in_group = 1, 1)" ) diff --git a/process/models/build.py b/process/models/build.py index c9e4077a1..30941d905 100644 --- a/process/models/build.py +++ b/process/models/build.py @@ -22,6 +22,13 @@ logger = logging.getLogger(__name__) +class DivertorNumberModels(IntEnum): + """Enum for divertor number models. `i_single_null` is the index for this enum.""" + + DOUBLE_NULL = 0 + SINGLE_NULL = 1 + + class FwBlktVVShape(IntEnum): """Enum for first wall, blanket, and vacuum vessel shape options.""" @@ -172,7 +179,8 @@ def calculate_vertical_build(self, output: bool): physics_variables.i_single_null, ) - if physics_variables.i_single_null == 0: + i_single_null = DivertorNumberModels(physics_variables.i_single_null) + if i_single_null == DivertorNumberModels.DOUBLE_NULL: po.ocmmnt(self.outfile, "Double null case") # Start at the top and work down. @@ -804,7 +812,7 @@ def calculate_vertical_build(self, output: bool): ) # Vertical locations of divertor coils - if physics_variables.i_single_null == 0: + if i_single_null == DivertorNumberModels.DOUBLE_NULL: self.data.build.z_tf_top = ( self.data.build.z_tf_inside_half + self.data.build.dr_tf_inboard ) @@ -1645,7 +1653,8 @@ def calculate_radial_build(self, output: bool): self.data.build.dr_blkt_inboard + self.data.build.dr_blkt_outboard ) - if physics_variables.i_single_null == 1: + i_single_null = DivertorNumberModels(physics_variables.i_single_null) + if i_single_null == DivertorNumberModels.SINGLE_NULL: # Check if self.data.build.dz_fw_plasma_gap has been set too small self.data.build.dz_fw_plasma_gap = max( 0.5e0 diff --git a/process/models/divertor.py b/process/models/divertor.py index 0b8cffca7..2d3e85d60 100644 --- a/process/models/divertor.py +++ b/process/models/divertor.py @@ -8,6 +8,7 @@ from process.core.model import Model from process.data_structure import physics_variables as pv from process.data_structure import tfcoil_variables as tfv +from process.models.build import DivertorNumberModels class Divertor(Model): @@ -178,10 +179,10 @@ def divtart( # Total divertor area # Single null case - if i_single_null == 1: + if i_single_null == DivertorNumberModels.SINGLE_NULL: areadv = a1 + a2 + a3 # Double null case - elif i_single_null == 0: + elif i_single_null == DivertorNumberModels.DOUBLE_NULL: areadv = 2.0 * (a1 + a2 + a3) if self.data.divertor.i_div_heat_load == 1: diff --git a/process/models/geometry/plasma.py b/process/models/geometry/plasma.py index 5131897d1..999211db8 100644 --- a/process/models/geometry/plasma.py +++ b/process/models/geometry/plasma.py @@ -7,6 +7,7 @@ import numpy as np +from process.models.build import DivertorNumberModels from process.models.physics.plasma_geometry import PlasmaShapeModelType @@ -77,7 +78,7 @@ def plasma_geometry( theta2 = np.arcsin((kappa * rminor) / r2) inang = 1.0 / r1 outang = 1.5 / r2 - if i_single_null == 0: + if i_single_null == DivertorNumberModels.DOUBLE_NULL: angs1 = np.linspace( -(inang + theta1) + np.pi, (inang + theta1) + np.pi, 500, endpoint=True ) diff --git a/process/models/tfcoil/base.py b/process/models/tfcoil/base.py index e34f13975..8fb150f83 100644 --- a/process/models/tfcoil/base.py +++ b/process/models/tfcoil/base.py @@ -25,6 +25,7 @@ superconducting_tf_coil_variables, tfcoil_variables, ) +from process.models.build import DivertorNumberModels from process.models.superconductors import SuperconductorModel if TYPE_CHECKING: @@ -508,7 +509,7 @@ def tf_coil_shape_inner( r_tf_arc[3] = r_tf_arc[1] r_tf_arc[4] = r_tf_arc[0] - if i_single_null == 0: + if i_single_null == DivertorNumberModels.DOUBLE_NULL: z_tf_arc[0] = FSTRAIGHT * z_tf_inside_half z_tf_arc[1] = z_tf_inside_half z_tf_arc[2] = 0