diff --git a/bluemira/equilibria/coils/_grouping.py b/bluemira/equilibria/coils/_grouping.py index c4c46eed7c..0c6de258f3 100644 --- a/bluemira/equilibria/coils/_grouping.py +++ b/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,10 +334,11 @@ 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], @@ -333,7 +346,8 @@ def from_group_vecs(cls, eqdsk: EQDSKInterface) -> CoilGroup: 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) diff --git a/bluemira/equilibria/equilibrium.py b/bluemira/equilibria/equilibrium.py index 3fc1cc89cb..a51ddd9008 100644 --- a/bluemira/equilibria/equilibrium.py +++ b/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, diff --git a/bluemira/equilibria/file.py b/bluemira/equilibria/file.py index 10aa548983..9b9ba4a42c 100644 --- a/bluemira/equilibria/file.py +++ b/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" diff --git a/bluemira/utilities/tools.py b/bluemira/utilities/tools.py index 8ea06fb94a..6aec87523a 100644 --- a/bluemira/utilities/tools.py +++ b/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 diff --git a/tests/equilibria/test_equilibrium.py b/tests/equilibria/test_equilibrium.py index b1b84c182c..b52a301c00 100644 --- a/tests/equilibria/test_equilibrium.py +++ b/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)