Add DivertorHeatLoadModel enum

process/data_structure/divertor_variables.py

@@ -1,4 +1,18 @@
 from dataclasses import dataclass
+from enum import IntEnum
+
+
+class DivertorHeatLoadModel(IntEnum):
+    """Divertor heat load model enumeration, controlled' by `i_div_heat_load`"""
+
+    USER_INPUT = 0
+    """User input for divertor heat load"""
+
+    PENG_CHAMBER = 1
+    """Divertor heat load model based on Peng chamber"""
+
+    WADE = 2
+    """Divertor heat load model based on Wade (Wade 2020)"""
 
 
 @dataclass(slots=True)
@@ -47,9 +61,9 @@ class DivertorData:
     i_div_heat_load: int = 2
     """switch for user input pflux_div_heat_load_mw:
 
-    - = 0: divtart model turned off and user inputs pflux_div_heat_load_mw
-    - = 1: divtart model calculates pflux_div_heat_load_mw
-    - = 2: divwade model calculates pflux_div_heat_load_mw"""
+    - = 0: User input
+    - = 1: Peng chamber model
+    - = 2: Wade model"""
 
     pflux_div_heat_load_max_mw: float = 5.0
     """heat load limit (MW/m2)"""

process/models/divertor.py

@@ -6,6 +6,7 @@
 from process.core import process_output as po
 from process.core.exceptions import ProcessValueError
 from process.core.model import Model
+from process.data_structure.divertor_variables import DivertorHeatLoadModel
 from process.data_structure.physics_variables import DivertorNumberModels
 
 
@@ -51,15 +52,22 @@ def run(self, output: bool = False):
             n_divertors=self.data.divertor.n_divertors,
         )
 
-        if self.data.divertor.i_div_heat_load == 0 and output:
+        if (
+            DivertorHeatLoadModel(self.data.divertor.i_div_heat_load)
+            == DivertorHeatLoadModel.USER_INPUT
+            and output
+        ):
             po.ovarre(
                 self.outfile,
                 "Divertor heat load (MW/m²)",
                 "(pflux_div_heat_load_mw)",
                 self.data.divertor.pflux_div_heat_load_mw,
             )
             return
-        if self.data.divertor.i_div_heat_load == 1:
+        if (
+            DivertorHeatLoadModel(self.data.divertor.i_div_heat_load)
+            == DivertorHeatLoadModel.PENG_CHAMBER
+        ):
             self.divtart(
                 self.data.physics.rmajor,
                 self.data.physics.rminor,
@@ -72,7 +80,10 @@ def run(self, output: bool = False):
                 dz_divertor=self.data.divertor.dz_divertor,
             )
             return
-        if self.data.divertor.i_div_heat_load == 2:
+        if (
+            DivertorHeatLoadModel(self.data.divertor.i_div_heat_load)
+            == DivertorHeatLoadModel.WADE
+        ):
             self.divwade(
                 self.data.physics.rmajor,
                 self.data.physics.rminor,
@@ -207,10 +218,17 @@ def divtart(
         elif i_single_null == DivertorNumberModels.DOUBLE_NULL:
             areadv = 2.0 * (a1 + a2 + a3)
 
-        if self.data.divertor.i_div_heat_load == 1:
+        if (
+            DivertorHeatLoadModel(self.data.divertor.i_div_heat_load)
+            == DivertorHeatLoadModel.PENG_CHAMBER
+        ):
             self.data.divertor.pflux_div_heat_load_mw = p_plasma_separatrix_mw / areadv
 
-        if output and self.data.divertor.i_div_heat_load == 1:
+        if (
+            output
+            and DivertorHeatLoadModel(self.data.divertor.i_div_heat_load)
+            == DivertorHeatLoadModel.PENG_CHAMBER
+        ):
             po.osubhd(self.outfile, "Divertor Heat Load")
             po.ocmmnt(self.outfile, "Assume an expanded divertor with a gaseous target")
             po.oblnkl(self.outfile)

tests/unit/models/test_divertor.py

@@ -2,6 +2,8 @@
 
 import pytest
 
+from process.data_structure.divertor_variables import DivertorHeatLoadModel
+
 
 @pytest.fixture
 def divertor(process_models):
@@ -38,7 +40,9 @@ def test_divtart(monkeypatch, divertor):
         p_plasma_separatrix_mw = 7.7197999809272062
         i_single_null = 0
         dz_divertor = 0.5
-        monkeypatch.setattr(divertor.data.divertor, "i_div_heat_load", 1)
+        monkeypatch.setattr(
+            divertor.data.divertor, "i_div_heat_load", DivertorHeatLoadModel.PENG_CHAMBER
+        )
 
         expected_pflux_div_heat_load_mw = 0.087770426974167357