🎨 Control coil enforcement, constraint stabilisation and general cleanup

bluemira/codes/process/api.py

@@ -11,7 +11,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from enum import Enum
+from enum import IntEnum, auto
 from importlib import resources
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal, TypeVar
@@ -21,7 +21,7 @@
 from bluemira.utilities.tools import flatten_iterable
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
+    from collections.abc import Sequence
 
 
 # Create dummy PROCESS objects. Required for docs to build properly and
@@ -113,25 +113,25 @@ def value(self, new_value):
         self._value = new_value
 
 
-class Impurities(Enum):
+class Impurities(IntEnum):
     """
     PROCESS impurities Enum
     """
 
-    H = 1
-    He = 2
-    Be = 3
-    C = 4
-    N = 5
-    O = 6  # noqa: E741
-    Ne = 7
-    Si = 8
-    Ar = 9
-    Fe = 10
-    Ni = 11
-    Kr = 12
-    Xe = 13
-    W = 14
+    H = auto()
+    He = auto()
+    Be = auto()
+    C = auto()
+    N = auto()
+    O = auto()  # noqa: E741
+    Ne = auto()
+    Si = auto()
+    Ar = auto()
+    Fe = auto()
+    Ni = auto()
+    Kr = auto()
+    Xe = auto()
+    W = auto()
 
     def files(self) -> dict[str, Path]:
         """
@@ -162,8 +162,8 @@ def id(self):
         return f"fimp({self.value:02})"
 
     def read_impurity_files(
-        self, filetype: Iterable[Literal["lz", "z2", "z"]]
-    ) -> tuple[list[ImpurityDataHeader]]:
+        self, filetype: Sequence[Literal["lz", "z2", "z"]]
+    ) -> tuple[list[ImpurityDataHeader], ...]:
         """Get contents of impurity data files"""
         files = self.files()
         return tuple(

bluemira/equilibria/coils/_coil.py

@@ -369,11 +369,7 @@ def position(self, values: np.ndarray):
     @ctype.setter
     def ctype(self, value: str | np.ndarray | CoilType):
         """Set coil type"""
-        self._ctype = (
-            value
-            if isinstance(value, CoilType)
-            else CoilType[value[0] if isinstance(value, np.ndarray) else value]
-        )
+        self._ctype = CoilType(value[0] if isinstance(value, np.ndarray) else value)
 
     @dx.setter
     def dx(self, value: float | None):

bluemira/equilibria/coils/_field.py

@@ -679,6 +679,27 @@ def Bz_response(
         """
         return self._sum(super().Bz_response(x, z), sum_coils=sum_coils, control=control)
 
+    def control_F(self, coil_grp: CoilGroup, *, control: bool = False) -> np.ndarray:
+        """
+        Returns the Green's matrix element for the coil mutual force.
+
+        \t:math:`Fz_{i,j}=-2\\pi X_i\\mathcal{G}(X_j,Z_j,X_i,Z_i)`
+
+        Parameters
+        ----------
+        coil_grp`:
+            the coil group to calculate against
+        control:
+            operations on control coils only
+        """
+        if control:
+            inds = self._control_ind
+            inds2 = coil_grp._control_ind
+        else:
+            inds = inds2 = slice(None)
+
+        return super().control_F(coil_grp)[inds][:, inds2]
+
     def _stored_greens(
         self, bgreen: np.ndarray, *, sum_coils: bool = True, control: bool = False
     ) -> np.ndarray:

bluemira/equilibria/optimisation/constraint_funcs.py

@@ -195,6 +195,10 @@ class FieldConstraintFunction(ConstraintFunction):
         Maximum fields inside the coils
     scale:
         Current scale with which to calculate the constraints
+    name:
+        name of constraint (used in error reporting)
+    round_dp:
+        round the output of the constraint (to maintain numerical stability)
     """
 
     def __init__(
@@ -206,6 +210,8 @@ def __init__(
         B_max: npt.NDArray[np.float64],
         scale: float,
         name: str | None = None,
+        *,
+        round_dp: int = 16,
     ):
         self.ax_mat = ax_mat
         self.az_mat = az_mat
@@ -214,6 +220,7 @@ def __init__(
         self.B_max = B_max
         self.scale = scale
         self.name = name
+        self._round_dp = round_dp
 
     def f_constraint(self, vector: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
         """Constraint function"""
@@ -223,7 +230,7 @@ def f_constraint(self, vector: npt.NDArray[np.float64]) -> npt.NDArray[np.float6
         Bz_a = self.az_mat @ currents
 
         B = np.hypot(Bx_a + self.bxp_vec, Bz_a + self.bzp_vec)
-        return B - self.B_max
+        return np.round(B - self.B_max, self._round_dp)
 
     def df_constraint(self, vector: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
         """Constraint derivative"""
@@ -235,7 +242,7 @@ def df_constraint(self, vector: npt.NDArray[np.float64]) -> npt.NDArray[np.float
 
         Bx = Bx_a * (Bx_a * currents + self.bxp_vec)
         Bz = Bz_a * (Bz_a * currents + self.bzp_vec)
-        return (Bx + Bz) / (B * self.scale**2)
+        return np.round((Bx + Bz) / (B * self.scale**2), self._round_dp)
 
 
 class CurrentMidplanceConstraint(ConstraintFunction):
@@ -394,17 +401,17 @@ def cs_z_sep_constraint(self, f_matx, max_value):
         """Constraint Function: CS separation constraints."""
         scaled_max_value = max_value / self.scale
         cs_fz = self.cs_fz(f_matx)
-        for i in range(self.n_CS - 1):  # evaluate each gap in CS stack
-            f_sep = np.sum(cs_fz[: i + 1]) - np.sum(cs_fz[i + 1 :])
-            self.constraint[self.n_PF + 1 + i] = f_sep - scaled_max_value
+        for i in range(1, self.n_CS):  # evaluate each gap in CS stack
+            f_sep = np.sum(cs_fz[:i]) - np.sum(cs_fz[i:])
+            self.constraint[self.n_PF + i] = f_sep - scaled_max_value
 
     def cs_z_sep_grad(self, df_matx):
         """Constraint Derivative: CS separation constraints."""
-        for i in range(self.n_CS - 1):  # evaluate each gap in CS stack
+        for i in range(1, self.n_CS):  # evaluate each gap in CS stack
             # CS separation constraint Jacobians
-            f_up = np.sum(df_matx[self.n_PF : self.n_PF + i + 1, :, 1], axis=0)
-            f_down = np.sum(df_matx[self.n_PF + i + 1 :, :, 1], axis=0)
-            self.grad[self.n_PF + 1 + i] = f_up - f_down
+            f_up = np.sum(df_matx[self.n_PF : self.n_PF + i, :, 1], axis=0)
+            f_down = np.sum(df_matx[self.n_PF + i :, :, 1], axis=0)
+            self.grad[self.n_PF + i] = f_up - f_down
 
 
 class CoilForceConstraint(ConstraintFunction, CoilForceConstraintFunctions):
@@ -432,6 +439,10 @@ class CoilForceConstraint(ConstraintFunction, CoilForceConstraintFunctions):
         The individual maximum force in the z direction for the CS coils
     scale:
         Current scale with which to calculate the constraints
+    name:
+        name of constraint (used in error reporting)
+    round_dp:
+        round the output of the constraint (to maintain numerical stability)
     """
 
     def __init__(
@@ -445,12 +456,15 @@ def __init__(
         CS_Fz_sep_max: float,
         scale: float,
         name: str | None = None,
+        *,
+        round_dp: int = 16,
     ):
         super().__init__(a_mat, b_vec, n_PF, n_CS, scale)
         self.PF_Fz_max = PF_Fz_max
         self.CS_Fz_sum_max = CS_Fz_sum_max
         self.CS_Fz_sep_max = CS_Fz_sep_max
         self.name = name
+        self._round_dp = round_dp
 
     def f_constraint(self, vector):
         """Constraint function"""
@@ -460,7 +474,7 @@ def f_constraint(self, vector):
         if self.n_CS != 0:
             self.cs_z_constraint(f_matx, self.CS_Fz_sum_max)
             self.cs_z_sep_constraint(f_matx, self.CS_Fz_sep_max)
-        return self.constraint
+        return np.round(self.constraint, self._round_dp)
 
     def df_constraint(self, vector):
         """Constraint derivative"""
@@ -470,4 +484,4 @@ def df_constraint(self, vector):
         if self.n_CS != 0:
             self.cs_z_grad(df_matx)
             self.cs_z_sep_grad(df_matx)
-        return self.grad
+        return np.round(self.grad, self._round_dp)

bluemira/equilibria/optimisation/constraints.py

@@ -238,7 +238,8 @@ def __init__(
         B_max: float | np.ndarray,
         tolerance: float | np.ndarray | None = None,
     ):
-        n_coils = coilset.n_coils()
+        cc = coilset.get_control_coils()
+        n_coils = cc.n_coils()
         if is_num(B_max):
             B_max *= np.ones(n_coils)
         if len(B_max) != n_coils:
@@ -252,6 +253,7 @@ def __init__(
 
     @staticmethod
     def _get_constraint_points(coilset):
+        coilset = coilset.get_control_coils()
         return coilset.x - coilset.dx, coilset.z
 
     def prepare(
@@ -365,20 +367,19 @@ def control_response(coilset: CoilSet) -> np.ndarray:
         """
         Calculate control response of a CoilSet to the constraint.
         """
-        return coilset.control_F(coilset)
+        return coilset.control_F(coilset, control=True)
 
     @staticmethod
     def evaluate(equilibrium: Equilibrium) -> np.ndarray:
         """
         Calculate the value of the constraint in an Equilibrium.
         """
-        fp = np.zeros((equilibrium.coilset.n_coils(), 2))
-        current = equilibrium.coilset.current
+        cc = equilibrium.coilset.get_control_coils()
+        fp = np.zeros((cc.n_coils(), 2))
+        current = cc.current
         non_zero = np.nonzero(current)[0]
         if non_zero.size:
-            fp[non_zero] = (
-                equilibrium.coilset.F(equilibrium)[non_zero] / current[non_zero][:, None]
-            )
+            fp[non_zero] = cc.F(equilibrium)[non_zero] / current[non_zero][:, None]
         return fp
 
     def f_constraint(self) -> CoilForceConstraintFunction:
@@ -855,7 +856,9 @@ def build_control_matrix(self):
         i = 0
         for constraint in self.constraints:
             n = len(constraint)
-            self.A[i : i + n, :] = constraint.control_response(self.coilset)
+            self.A[i : i + n, :] = constraint.control_response(
+                self.coilset.get_control_coils()
+            )
             i += n
 
     def build_target(self):