Import coilset details from eqdsk if available (#3356)

* ✨ Add coilset additions to equilibrium

* ✅ Tests

* Update bluemira/equilibria/coils/_grouping.py

Co-authored-by: Oliver Funk <oliverfunk@users.noreply.github.com>

* Update bluemira/equilibria/coils/_grouping.py

Co-authored-by: Oliver Funk <oliverfunk@users.noreply.github.com>

---------

Co-authored-by: Oliver Funk <oliverfunk@users.noreply.github.com>

bluemira/equilibria/coils/_grouping.py

@@ -17,11 +17,6 @@
 from operator import attrgetter
 from typing import TYPE_CHECKING
 
-if TYPE_CHECKING:
-    from matplotlib.axes import Axes
-
-    from bluemira.equilibria.file import EQDSKInterface
-
 import numpy as np
 
 from bluemira.base.constants import CoilType
@@ -39,6 +34,12 @@
 from bluemira.equilibria.plotting import CoilGroupPlotter
 from bluemira.utilities.tools import flatten_iterable, yintercept
 
+if TYPE_CHECKING:
+    import numpy.typing as npt
+    from matplotlib.axes import Axes
+
+    from bluemira.equilibria.file import EQDSKInterface
+
 
 def symmetrise_coilset(coilset: CoilSet) -> CoilSet:
     """
@@ -296,7 +297,7 @@ def from_group_vecs(cls, eqdsk: EQDSKInterface) -> CoilGroup:
                     current=0,
                     dx=0,
                     dz=0,
-                    ctype="DUM",
+                    ctype=CoilType.DUM,
                     j_max=0,
                     b_max=0,
                 )
@@ -308,10 +309,21 @@ def from_group_vecs(cls, eqdsk: EQDSKInterface) -> CoilGroup:
                 "Please replace with an appropriate coilset."
             )
             return cls(*coils)
+
+        def _get_val(lst: npt.ArrayLike | None, idx: int, default=None):
+            if lst is None:
+                return None
+            try:
+                return lst[idx]
+            except IndexError:
+                return default
+
         for i in range(eqdsk.ncoil):
             dx = eqdsk.dxc[i]
             dz = eqdsk.dzc[i]
-            if abs(eqdsk.Ic[i]) < I_MIN:
+            cn = _get_val(eqdsk.coil_names, i)
+            ct = None if (v := _get_val(eqdsk.coil_types, i)) is None else CoilType(v)
+            if ct is CoilType.NONE or (abs(eqdsk.Ic[i]) < I_MIN and ct is None):
                 # Some eqdsk formats (e.g., CREATE) contain 'quasi-coils'
                 # with currents very close to 0.
                 # Catch these cases and make sure current is set to zero.
@@ -322,18 +334,20 @@ def from_group_vecs(cls, eqdsk: EQDSKInterface) -> CoilGroup:
                         current=0,
                         dx=dx,
                         dz=dz,
-                        ctype="NONE",
+                        ctype=ct or CoilType.NONE,
+                        name=cn,
                     )
                 )
-            elif dx != dz:  # Rough and ready
+            elif ct is CoilType.CS or (dx != dz and ct is None):  # Rough and ready
                 cscoils.append(
                     Coil(
                         eqdsk.xc[i],
                         eqdsk.zc[i],
                         current=eqdsk.Ic[i],
                         dx=dx,
                         dz=dz,
-                        ctype="CS",
+                        ctype=CoilType.CS,
+                        name=cn,
                     )
                 )
             else:
@@ -343,7 +357,8 @@ def from_group_vecs(cls, eqdsk: EQDSKInterface) -> CoilGroup:
                     current=eqdsk.Ic[i],
                     dx=dx,
                     dz=dz,
-                    ctype="PF",
+                    ctype=ct or CoilType.PF,
+                    name=cn,
                 )
                 coil.fix_size()  # Oh ja
                 pfcoils.append(coil)

bluemira/equilibria/equilibrium.py

@@ -187,6 +187,10 @@ def to_eqdsk(
             filename = Path(directory, filename)
 
         self.filename = filename  # Convenient
+        if data.get("coil_types") is not None:
+            data["coil_types"] = [
+                ct if isinstance(ct, str) else ct.name for ct in data["coil_types"]
+            ]
         eqdsk = EQDSKInterface(**data)
         eqdsk.write(filename.as_posix(), file_format=filetype, **kwargs)
 
@@ -594,6 +598,8 @@ def to_dict(self) -> dict[str, Any]:
             "Bz": self.Bz(),
             "Bp": self.Bp(),
             "ncoil": self.coilset.n_coils(),
+            "coil_names": self.coilset.name,
+            "coil_types": self.coilset.ctype,
             "xc": xc,
             "zc": zc,
             "dxc": dxc,
@@ -974,6 +980,8 @@ def to_dict(self, qpsi_calcmode: int = 0) -> dict[str, Any]:
             "xbdry": lcfs.x,
             "zbdry": lcfs.z,
             "ncoil": self.coilset.n_coils(),
+            "coil_names": self.coilset.name,
+            "coil_types": self.coilset.ctype,
             "xc": x_c,
             "zc": z_c,
             "dxc": dxc,

bluemira/equilibria/file.py

@@ -230,7 +230,7 @@ def _read_json(file) -> dict[str, Any]:
     data_has_pnorm = False
     data_has_psinorm = False
     for k, value in data.items():
-        if isinstance(value, list):
+        if isinstance(value, list) and k not in {"coil_type", "coil_names"}:
             data[k] = np.asarray(value)
         data_has_pnorm |= k == "pnorm"
         data_has_psinorm |= k == "psinorm"

bluemira/utilities/tools.py

@@ -13,6 +13,7 @@
 import operator
 import string
 import warnings
+from collections import Counter
 from collections.abc import Callable, Iterable
 from functools import wraps
 from importlib import import_module as imp
@@ -30,7 +31,12 @@
 
 from bluemira.base.constants import E_I, E_IJ, E_IJK
 from bluemira.base.file import force_file_extension
-from bluemira.base.look_and_feel import bluemira_debug, bluemira_print, bluemira_warn
+from bluemira.base.look_and_feel import (
+    bluemira_debug,
+    bluemira_error,
+    bluemira_print,
+    bluemira_warn,
+)
 
 if TYPE_CHECKING:
     from types import ModuleType
@@ -558,6 +564,9 @@ def num_almost_eq(val1, val2):
     def array_is_eq(val1, val2):
         return (np.asarray(val1) == np.asarray(val2)).all()
 
+    def list_eq(val1, val2):
+        return Counter(val1) == Counter(val2)
+
     if almost_equal:
         array_eq = array_almost_eq
         num_eq = num_almost_eq
@@ -569,10 +578,13 @@ def array_is_eq(val1, val2):
     comp_map = {
         key: (
             array_eq
-            if isinstance(val, np.ndarray | list)
+            if isinstance(val, np.ndarray)
+            or (isinstance(val, list) and is_num(next(flatten_iterable(val))))
             else (
                 dict_eq
                 if isinstance(val, dict)
+                else list_eq
+                if isinstance(val, list)
                 else num_eq
                 if is_num(val)
                 else operator.eq
@@ -614,7 +626,7 @@ def array_is_eq(val1, val2):
     else:
         result += "==========================================================="
         if verbose:
-            print(result)  # noqa: T201
+            bluemira_error(result)
     return the_same
 
 

tests/equilibria/test_equilibrium.py

@@ -365,6 +365,19 @@ def test_woops_no_coils(self, grouping):
         if grouping is CoilSet:
             assert len(coil.control) == 0
 
+    def test_eq_coilnames(self):
+        testfile = Path(
+            get_bluemira_path("equilibria/test_data", subfolder="tests"),
+            "DN-DEMO_eqref_withCoilNames.json",
+        )
+        e = Equilibrium.from_eqdsk(testfile)
+        assert e.coilset.name == [
+            *("PF_1", "PF_2", "PF_3", "PF_4", "PF_5", "PF_6"),
+            *("CS_1", "CS_2", "CS_3", "CS_4", "CS_5"),
+        ]
+        assert e.coilset.n_coils(ctype="PF") == 6
+        assert e.coilset.n_coils(ctype="CS") == 5
+
 
 class TestEqReadWrite:
     @pytest.mark.parametrize("qpsi_calcmode", [0, 1])
@@ -387,6 +400,9 @@ def test_read_write(self, qpsi_calcmode, file_format):
         eq2 = Equilibrium.from_eqdsk(new_file_path)
         d2 = eq2.to_dict(qpsi_calcmode=qpsi_calcmode)
         new_file_path.unlink()
+        if file_format == "eqdsk":
+            d1.pop("coil_names")
+            d2.pop("coil_names")
         assert compare_dicts(d1, d2, almost_equal=True)