Papercut: enumify text based choices

bluemira/builders/pf_coil.py

@@ -9,6 +9,7 @@
 """
 
 from dataclasses import dataclass
+from enum import Enum, auto
 from typing import Dict, List, Union
 
 from bluemira.base.builder import Builder
@@ -24,6 +25,22 @@
 from bluemira.geometry.wire import BluemiraWire
 
 
+class CoilType(Enum):
+    """Enumification of text based choice of coil type"""
+
+    PF = auto()
+    CS = auto()
+
+    @classmethod
+    def _missing_(cls, value):
+        try:
+            return cls[value.upper()]
+        except KeyError:
+            raise ValueError(
+                f"{value} is not a valid CoilType. Choose from: PF or CS"
+            ) from None
+
+
 @dataclass
 class PFCoilBuilderParams(ParameterFrame):
     """
@@ -79,7 +96,8 @@
         c2 = make_circle(r_in)
 
         wp = PhysicalComponent(self.WINDING_PACK, BluemiraFace([c1, c2]))
-        idx = 0 if self.params.ctype.value == "CS" else 1
+
+        idx = CoilType(self.params.ctype.value.upper()).value - 1
         apply_component_display_options(wp, color=BLUE_PALETTE["PF"][idx])
 
         r_in -= self.params.tk_insulation.value
@@ -113,7 +131,7 @@
         Build the xz cross-section of the PF coil.
         """
         wp = PhysicalComponent(self.WINDING_PACK, BluemiraFace(shape))
-        idx = 0 if self.params.ctype.value == "CS" else 1
+        idx = CoilType(self.params.ctype.value.upper()).value - 1
         apply_component_display_options(wp, color=BLUE_PALETTE["PF"][idx])
 
         ins_shape = offset_wire(shape, self.params.tk_insulation.value)

bluemira/equilibria/equilibrium.py

@@ -9,7 +9,7 @@
 """
 
 from copy import deepcopy
-from enum import Enum
+from enum import Enum, auto
 from pathlib import Path
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
@@ -59,6 +59,26 @@
 from bluemira.utilities.tools import abs_rel_difference
 
 
+class VControlType(Enum):
+    """
+    Enumification of text based numerical stabilisation strategy
+    choices for vertical position controller
+    """
+
+    VIRTUAL = auto()
+    FEEDBACK = auto()
+
+    @classmethod
+    def _missing_(cls, value):
+        try:
+            return cls[value.upper()]
+        except KeyError:
+            raise ValueError(
+                "Please select a numerical stabilisation strategy"
+                ' from: 1) "virtual" \n 2) "feedback" 3) None.'
+            ) from None
+
+
 class MHDState:
     """
     Base class for magneto-hydrodynamic states
@@ -1048,17 +1068,19 @@
         vcontrol:
             Vertical control strategy
         """
-        if vcontrol == "virtual":
-            self.controller = VirtualController(self, gz=2.2)
-        elif vcontrol == "feedback":
-            raise NotImplementedError
-        elif vcontrol is None:
+        if vcontrol is None:
             self.controller = DummyController(self.plasma.psi())
         else:
-            raise ValueError(
-                "Please select a numerical stabilisation strategy"
-                ' from: 1) "virtual" \n 2) "feedback" 3) None.'
-            )
+            vcontrol_str = VControlType(vcontrol.upper())
+            if vcontrol_str is VControlType.VIRTUAL:
+                self.controller = VirtualController(self, gz=2.2)
+            elif vcontrol_str is VControlType.FEEDBACK:
+                raise NotImplementedError
+            else:
+                raise ValueError(
+                    "Please select a numerical stabilisation strategy"
+                    ' from: 1) "virtual" \n 2) "feedback" 3) None.'
+                )
 
     def solve(self, jtor: Optional[np.ndarray] = None, psi: Optional[np.ndarray] = None):
         """

bluemira/equilibria/plotting.py

@@ -11,6 +11,7 @@
 from __future__ import annotations
 
 import warnings
+from enum import Enum, auto
 from itertools import cycle
 from typing import TYPE_CHECKING
 
@@ -95,6 +96,22 @@
 }
 
 
+class CoilType(Enum):
+    """Enumification of text based choice of coil type"""
+
+    PF = auto()
+    CS = auto()
+
+    @classmethod
+    def _missing_(cls, value):
+        try:
+            return cls[value.upper()]
+        except KeyError:
+            raise ValueError(
+                f"{value} is not a valid CoilType. Choose from: PF or CS"
+            ) from None
+
+
 class Plotter:
     """
     Utility plotter abstract object
@@ -235,9 +252,10 @@
         for i, (x, z, dx, x_b, z_b, ct, n, cur, ctrl) in enumerate(zip(*arrays)):
             if ctrl:
                 if self.colors is not None:
-                    if ct.name == "PF":
+                    ctype = CoilType(ct.name.upper())
+                    if ctype is CoilType.PF:
                         kwargs["facecolor"] = self.colors[0]
-                    elif ct.name == "CS":
+                    elif ctype is CoilType.CS:
                         kwargs["facecolor"] = self.colors[1]
 
                 self._plot_coil(
@@ -307,7 +325,8 @@
         Single coil annotation utility function
         """
         off = max(0.2, dx + 0.02)
-        if ctype.name == "CS":
+        ctype_name = CoilType(ctype.name.upper())
+        if ctype_name is CoilType.CS:
             drs = -1.5 * off
             ha = "right"
         else:
@@ -318,7 +337,7 @@
             text = "\n".join([text, f"{raw_uc(force[1], 'N', 'MN'):.2f} MN"])
         x = float(x) + drs
         z = float(z)
-        if centre is not None and ctype.name == "PF":
+        if centre is not None and ctype_name is CoilType.PF:
             v = np.array([x - centre[0], z - centre[1]])
             v /= np.sqrt(sum(v**2))
             d = 1 + np.sqrt(2) * dx

bluemira/equilibria/positioner.py

@@ -11,7 +11,7 @@
 import re
 from copy import deepcopy
 from typing import Dict, List, Tuple, Union
-
+from enum import Enum,auto
 import numpy as np
 from scipy.interpolate import InterpolatedUnivariateSpline, interp1d
 from scipy.optimize import minimize_scalar
@@ -40,6 +40,30 @@
 from bluemira.geometry.tools import boolean_cut, boolean_fuse, make_polygon, offset_wire
 from bluemira.utilities import tools
 
+class ReactorType(Enum):
+    """Enumification of text based choices of the type of reactor"""
+
+    NORMAL = auto()
+    ST = auto()
+
+    @classmethod
+    def _missing_(cls, value):
+        try:
+            return cls[value.upper()]
+        except KeyError:
+            raise ValueError(f"{value} is a wrong ReactorType.  Choose from: ST or Normal") from None
+class CoilsetLayoutType(Enum):
+    """Enumification of text based choices for the layout of the CS modules"""
+
+    ITER = auto()
+    DEMO = auto()
+
+    @classmethod
+    def _missing_(cls, value):
+        try:
+            return cls[value.upper()]
+        except KeyError:
+            raise ValueError(f"{value} is not a valid coilset Layout. Choose from: 'ITER' and 'DEMO'") from None
 
 class CoilPositioner:
     """
@@ -99,8 +123,8 @@
         self.n_PF = n_PF
         self.n_CS = n_CS
         self.csgap = csgap
-        self.rtype = rtype
-        self.cslayout = cslayout
+        self.rtype = ReactorType(rtype.upper())
+        self.cslayout = CoilsetLayoutType(cslayout.upper())
 
     def equispace_PF(self, track: Coordinates, n_PF: int) -> List[Coil]:
         """
@@ -109,10 +133,11 @@
         based on plasma shape considerations (mirror about X-points)
         """
         a = np.rad2deg(np.arctan(abs(self.delta) / self.kappa))
-        if self.rtype == "Normal":
+
+        if self.rtype is ReactorType.NORMAL:
             angle_upper = 90 + a * 1.6
             angle_lower = -90 - a * 1.6
-        elif self.rtype == "ST":
+        elif self.rtype is ReactorType.ST:
             angle_upper = 90 + a * 1.2
             angle_lower = -90 - a * 1.0
 
@@ -224,18 +249,14 @@
         z_max = max(self.track.z)
         z_min = -z_max
         if self.n_CS != 0:
-            if self.cslayout == "ITER":
+            if self.cslayout  is CoilsetLayoutType.ITER:
                 coils.append(
                     self.equispace_CS(self.x_cs, self.tk_cs, z_min, z_max, self.n_CS)
                 )
-            elif self.cslayout == "DEMO":
+            elif self.cslayout is CoilsetLayoutType.DEMO:
                 coils.extend(
                     self.demospace_CS(self.x_cs, self.tk_cs, z_min, z_max, self.n_CS)
                 )
-            else:
-                raise ValueError(
-                    f"Valid options are 'ITER' and 'DEMO', not '{self.cslayout}'"
-                )
         cset = CoilSet(*coils)
         cset.discretisation = d_coil
         return cset

bluemira/fuel_cycle/analysis.py

@@ -14,6 +14,7 @@
 
 if TYPE_CHECKING:
     from bluemira.fuel_cycle.cycle import EUDEMOFuelCycleModel
+from enum import Enum, auto
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -24,6 +25,25 @@
 __all__ = ["FuelCycleAnalysis"]
 
 
+class QueryType(Enum):
+    """Enumification of text based choices for the type of statistical value to return"""
+
+    MIN = auto()
+    MAX = auto()
+    MEAN = auto()
+    MEDIAN = auto()
+    P95TH = auto()  # identifier should start with a letter
+
+    @classmethod
+    def _missing_(cls, value):
+        try:
+            return cls[value.upper()]
+        except KeyError:
+            raise ValueError(
+                f"Invalid query: {value}. Choose from: min, max, mean, median, 95th"
+            ) from None
+
+
 class FuelCycleAnalysis:
     """
     Analysis class for compiling and analysing fuel cycle statistics.
@@ -117,17 +137,21 @@
         return self._query("t_d", s=query)
 
     def _query(self, p: str, s: str) -> float:
-        if s == "min":
+        inp_s = QueryType.P95TH if s == "95th" else QueryType(s.upper())
+
+        if inp_s is QueryType.MIN:
             return min(self.__dict__[p])
-        if s == "max":
+        if inp_s is QueryType.MAX:
             return max(self.__dict__[p])
-        if s == "mean":
+        if inp_s is QueryType.MEAN:
             return np.mean(self.__dict__[p])
-        if s == "median":
+        if inp_s is QueryType.MEDIAN:
             return np.median(self.__dict__[p])
-        if s == "95th":
+        if inp_s is QueryType.P95TH:
             return np.percentile(self.__dict__[p], 95)
-        raise ValueError(f"Unknown query: '{s}'")
+        raise ValueError(
+            f"Unknown query: '{s}'. Choose from: min, max, mean, median, 95th"
+        )
 
     def plot(self, figsize=(12, 6), bins=20, **kwargs):
         """

bluemira/fuel_cycle/lifecycle.py

@@ -11,6 +11,7 @@
 
 import json
 from dataclasses import dataclass, field
+from enum import Enum, auto
 from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
 import matplotlib.pyplot as plt
@@ -31,6 +32,22 @@
 __all__ = ["LifeCycle"]
 
 
+class PlotType(Enum):
+    """Enumification of text based choices for the type of plot"""
+
+    PIE = auto()
+    BAR = auto()
+
+    @classmethod
+    def _missing_(cls, value):
+        try:
+            return cls[value.upper()]
+        except KeyError:
+            raise ValueError(
+                f"{value} is not a valid plot type. Choose from: pie or bar"
+            ) from None
+
+
 class LifeCycle:
     """
     A life cycle object for a fusion reactor.
@@ -412,7 +429,8 @@
             self.total_planned_maintenance,
             self.t_unplanned_m,
         ]
-        if typ == "pie":
+        plt_typ = PlotType(typ.upper())
+        if plt_typ is PlotType.PIE:
             plt.pie(
                 sizes,
                 labels=labels,
@@ -422,7 +440,7 @@
                 counterclock=False,
             )
             plt.axis("equal")
-        elif typ == "bar":
+        elif plt_typ is PlotType.BAR:
             bottom = 0
             for i, s in enumerate(sizes):
                 ax.bar(1, s / sum(sizes) * 100, bottom=bottom, label=labels[i])