Removes _index from cocos params, improves & fixes docs, raises when cannot det. single cocos or no qpsi present

eqdsk/cocos.py

@@ -9,54 +9,17 @@
 from copy import deepcopy
 from dataclasses import dataclass
 from enum import Enum, unique
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 import numpy as np
 
 from eqdsk.log import eqdsk_warn
+from eqdsk.models import Sign, ZeroOne
 
 if TYPE_CHECKING:
     from eqdsk.file import EQDSKInterface
 
 
-class ZeroOne(Enum):
-    """An enum representing the values 0 and 1 for the 2pi exponent of Bp."""
-
-    ZERO = 0
-    ONE = 1
-
-    def __sub__(self, other: Any) -> ZeroOne:
-        """Return the difference between the value and the other value.
-
-        - If it is another ZeroOne, return the difference of the values.
-        - Raise a `TypeError` otherwise.
-        """
-        if type(other) is ZeroOne:
-            return ZeroOne(self.value - other.value)
-        raise TypeError(
-            f"Cannot subtract {type(other)} from {type(self)}.",
-        )
-
-
-class Sign(Enum):
-    """An enum representing the positive or negative sign of
-    a COCOS parameter.
-    """
-
-    POSITIVE = 1
-    NEGATIVE = -1
-
-    def __mul__(self, other: Any):
-        """Return the product of the sign with the other value.
-
-        - If it is another Sign, return the product of the values.
-        - If it is a number, return the product of the value and the number.
-        """
-        if type(other) is Sign:
-            return Sign(self.value * other.value)
-        return self.value * other
-
-
 @dataclass(frozen=True)
 class COCOSParams:
     """The parameters for a single COCOS definition."""
@@ -68,11 +31,11 @@ class COCOSParams:
     sign_Bp: Sign
     """The sign of Bp, depends on the sign of Ip and the gradient of psi."""
     sign_R_phi_Z: Sign
-    """The sign of (R, phi, Z), positive if theta and phi have
-    opposite directions, negative if the same."""
+    """The sign of (R, phi, Z), positive if phi (toroidal)
+    is CCW from the top, negative if CW."""
     sign_rho_theta_phi: Sign
-    """The sign of (rho, theta, phi), positive if phi (toroidal)
-    is CW from the top."""
+    """The sign of (rho, theta, phi), positive if theta and phi have
+    opposite directions, negative if the same."""
 
 
 @unique
@@ -275,7 +238,7 @@ def identify_eqdsk(
         A list of the identified COCOS definitions.
     """
     if eqdsk.qpsi is None:
-        eqdsk_warn("WARNING: qpsi is not defined in the eqdsk file. Setting to 1")
+        eqdsk_warn("qpsi is not defined in the eqdsk file. Setting to 1")
         eqdsk.qpsi = np.array([1])
 
     cw_phi_l = [True, False] if clockwise_phi is None else [clockwise_phi]
@@ -333,12 +296,12 @@ def identify_cocos(
 
     Returns
     -------
-        The identified COCOS convention.
+    The identified COCOS convention.
 
     Raises
     ------
-        ValueError: If the sign of qpsi is not consistent across the flux
-            surfaces.
+    ValueError: If the sign of qpsi is not consistent across the flux
+        surfaces.
     """
     sign_R_phi_Z = Sign.NEGATIVE if phi_clockwise_from_top else Sign.POSITIVE
     exp_Bp = ZeroOne.ZERO if volt_seconds_per_radian else ZeroOne.ONE

eqdsk/errors.py

@@ -0,0 +1,34 @@
+# SPDX-FileCopyrightText: 2023-present The Bluemira Developers <https://github.com/Fusion-Power-Plant-Framework/bluemira>
+#
+# SPDX-License-Identifier: LGPL-2.1-or-later
+
+"""Custom errors"""
+
+from eqdsk.cocos import COCOS
+
+
+class NoSingleConventionError(Exception):
+    """Raised when no single COCOS can be determined from the EQDSK file."""
+
+    def __init__(self, conventions: list[COCOS], message_extra: str):
+        self.conventions = conventions
+        self.message_extra = message_extra
+        super().__init__(
+            f"A single COCOS could not be determined, "
+            f"found conventions "
+            f"({', '.join([str(c.index) for c in self.conventions])}) "
+            f"for the EQDSK file.\n{self.message_extra}"
+        )
+
+
+class MissingQpsiDataError(Exception):
+    """Raised when attempting to identify the COCOS for an EQDSK file,
+    that does not have an qpsi data in it.
+    """
+
+    def __init__(self, message_extra: str):
+        self.message_extra = message_extra
+        super().__init__(
+            f"In order to properly identify the COCOS of this EQDSK file, "
+            f"qpsi data must be present in the file.\n{self.message_extra}"
+        )

eqdsk/file.py

@@ -16,12 +16,18 @@
 import numpy as np
 
 from eqdsk.cocos import COCOS, convert_eqdsk, identify_eqdsk
+from eqdsk.errors import (
+    MissingQpsiDataError,
+    NoSingleConventionError,
+)
 from eqdsk.log import eqdsk_print, eqdsk_warn
 from eqdsk.tools import is_num, json_writer
 
 if TYPE_CHECKING:
     from io import TextIOWrapper
 
+    from eqdsk.models import Sign
+
 EQDSK_EXTENSIONS = [".eqdsk", ".eqdsk_out", ".geqdsk"]
 
 
@@ -47,7 +53,7 @@ class EQDSKInterface:
     Plasma current direction is not enforced here!
     """
 
-    DEFAULT_COCOS_INDEX = 11
+    DEFAULT_COCOS = 11
 
     bcentre: float
     """Vacuum toroidal Magnetic field at the reference radius [T]."""
@@ -138,15 +144,16 @@ def __post_init__(self):
         self._cocos = None
 
     @classmethod
-    def from_file(
+    def from_file(  # noqa: PLR0913
         cls,
         file_path: str | Path,
-        from_cocos_index: int | None = None,
-        to_cocos_index: int | None = DEFAULT_COCOS_INDEX,
+        from_cocos: int | None = None,
+        to_cocos: int | None = DEFAULT_COCOS,
         *,
         clockwise_phi: bool | None = None,
         volt_seconds_per_radian: bool | None = None,
         no_cocos: bool = False,
+        qpsi_sign: Sign | None = None,
     ) -> EQDSKInterface:
         """Create an EQDSKInterface object from a file.
 
@@ -158,19 +165,23 @@ def from_file(
                 - eqdsk
                 - eqdsk_out
                 - geqdsk
+        from_cocos:
+            The user set COCOS of the EQDSK file.
+            This sets what COCCOS the file will be id's as
+            and will raise it's not one of the determined COCOS.
+        to_cocos:
+            The COCOS to convert the EQDSK file to. If None, the file
+            will not be converted.
         clockwise_phi:
             Whether the EQDSK file's phi is clockwise or not.
         volt_seconds_per_radian:
             Whether the EQDSK file's psi is in volt seconds per radian.
-        from_cocos_index:
-            The COCOS index of the EQDSK file. Used when the determined
-            COCOS is ambiguous. Will raise if given and not one of
-            the determined COCOS indices.
-        to_cocos_index:
-            The COCOS index to convert the EQDSK file to.
         no_cocos:
-            Whether to return the EQDSK data without identifying
-            and converting to `to_cocos_index` COCOS index.
+            Whether to return the EQDSK data without performing
+            any identifying COCOS identification or conversion.
+        qpsi_sign:
+            The sign of the qpsi, required for identification
+            when qpsi is not present in the file.
 
         Returns
         -------
@@ -189,13 +200,22 @@ def from_file(
         if no_cocos:
             return inst
 
-        inst.identify(
-            as_cocos_index=from_cocos_index,
-            clockwise_phi=clockwise_phi,
-            volt_seconds_per_radian=volt_seconds_per_radian,
-        )
-        if to_cocos_index is not None:
-            inst = inst.as_cocos(to_cocos_index)
+        try:
+            inst.identify(
+                as_cocos=from_cocos,
+                clockwise_phi=clockwise_phi,
+                volt_seconds_per_radian=volt_seconds_per_radian,
+                qpsi_sign=qpsi_sign,
+            )
+        except NoSingleConventionError as e:
+            raise NoSingleConventionError(
+                e.conventions,
+                message_extra="You need to specify `from_cocos` or "
+                "`clockwise_phi` and `volt_seconds_per_radian`.",
+            ) from None
+
+        if to_cocos is not None:
+            inst = inst.to_cocos(to_cocos)
 
         return inst
 
@@ -212,64 +232,100 @@ def cocos(self) -> COCOS:
 
     def identify(
         self,
-        as_cocos_index: int | None = None,
+        as_cocos: int | None = None,
         *,
         clockwise_phi: bool | None = None,
         volt_seconds_per_radian: bool | None = None,
+        qpsi_sign: Sign | None = None,
     ):
-        """Identify the COCOS of this eqdsk and set the COCOS attribute.
+        """Identifies the COCOS of this eqdsk.
+
+        Note
+        ----
+        This sets the internal _cocos attribute and does not return
+        anything.
 
         Parameters
         ----------
+        as_cocos:
+            The COCOS index to convert the EQDSK file to.
+            If given, the file will be id's as the given COCOS,
+            only if it one of the possible identified COCOS.
         clockwise_phi:
             Whether the EQDSK file's phi is clockwise or not.
         volt_seconds_per_radian:
             Whether the EQDSK file's psi is in volt seconds per radian.
-        as_cocos_index:
-            The COCOS index to convert the EQDSK file to. If given,
-            the COCOS will be converted to this COCOS index.
+        qpsi_sign:
+            The sign of the qpsi, required for identification
+            when qpsi is not present in the file.
 
         Raises
         ------
         ValueError:
-            If as_cocos_index is given but does not match any
-            identified COCOS index.
+            If as_cocos is given but does not match any identified COCOS.
         ValueError:
             If no COCOS can be identified.
 
         """
+        qpsi_is_not_set = self.qpsi is None or np.allclose(self.qpsi, 0)
+        if qpsi_is_not_set:
+            if qpsi_sign:
+                eqdsk_warn(
+                    "No qpsi data found but `qpsi_sign` provided. "
+                    f"Setting qpsi as array of {qpsi_sign.value}'s."
+                )
+                self.qpsi = np.ones(self.nx) * qpsi_sign.value
+            else:
+                raise MissingQpsiDataError(
+                    message_extra="Setting the `qpsi_sign` parameter will resolve this. "
+                    "This is the sign of the qpsi across the flux surfaces."
+                    "\nRefer to the COCOS spec or the implementation of it "
+                    "in this package, if you know what the desired direction (CW/CCW) "
+                    "is for theta and phi for this EQDSK."
+                    "This can help you determine what sign qpsi should be."
+                )
+
         conventions = identify_eqdsk(
             self,
             clockwise_phi=clockwise_phi,
             volt_seconds_per_radian=volt_seconds_per_radian,
         )
 
-        if as_cocos_index is not None:
-            matching_conv = [c for c in conventions if c.index == as_cocos_index]
-            if not matching_conv:
-                raise ValueError(
-                    f"No convention found that matches "
-                    f"the given COCOS index {as_cocos_index}, "
-                    f"from the possible ({', '.join([str(c.index) for c in conventions])}).",  # noqa: E501
+        def _id():
+            if as_cocos:
+                matching_conv = next(
+                    (c for c in conventions if c.index == as_cocos), None
+                )
+                if not matching_conv:
+                    raise ValueError(
+                        f"No convention found that matches "
+                        f"the given COCOS index {as_cocos}, "
+                        f"from the possible ({', '.join([str(c.index) for c in conventions])}).",  # noqa: E501
+                    )
+                return matching_conv
+            if len(conventions) != 1:
+                raise NoSingleConventionError(
+                    conventions,
+                    message_extra="You need to specify `as_cocos` or "
+                    "`clockwise_phi` and `volt_seconds_per_radian`.",
                 )
-            conventions = matching_conv
+            return conventions[0]
 
-        conv = conventions[0]
-        if len(conventions) != 1:
-            eqdsk_warn(
-                f"A single COCOS could not be determined, "
-                f"found conventions ({', '.join([str(c.index) for c in conventions])}) "
-                f"for the EQDSK file. Choosing COCOS {conv.index}.",
-            )
-        eqdsk_print(f"EQDSK identified as COCOS {conv.index}.")
-        self._cocos = conv
+        c = _id()
+        eqdsk_print(f"EQDSK identified as COCOS {c.index}.")
+        self._cocos = c
+
+    def to_cocos(self, to_cocos: int) -> EQDSKInterface:
+        """Returns a copy of this eqdsk converted to the given COCOS.
 
-    def as_cocos(self, cocos_index: int) -> EQDSKInterface:
-        """Return a copy of this eqdsk converted to the given COCOS."""
-        if self.cocos.index == cocos_index:
+        Note
+        ----
+            This returns a new instance of the EQDSKInterface class.
+        """
+        if self.cocos.index == to_cocos:
             return self
-        eqdsk_print(f"Converting EQDSK to COCOS {cocos_index}.")
-        return convert_eqdsk(self, cocos_index)
+        eqdsk_print(f"Converting EQDSK to COCOS {to_cocos}.")
+        return convert_eqdsk(self, to_cocos)
 
     def to_dict(self) -> dict:
         """Return a dictionary of the EQDSK data."""
@@ -612,9 +668,9 @@ def write_array(fortran_format: ff.FortranRecordWriter, array: np.ndarray):
 
         # Define dummy data for qpsi if it has not been previously defined.
         qpsi = (
-            np.ones(data["nx"]) if data["qpsi"] is None else np.atleast_1d(data["qpsi"])
+            np.zeros(data["nx"]) if data["qpsi"] is None else np.atleast_1d(data["qpsi"])
         )
-
+        eqdsk_warn("No qpsi data found. Setting as array of 0's.")
         if len(qpsi) == 1:
             qpsi = np.full(data["nx"], qpsi)
         elif len(qpsi) != data["nx"]: