Equilibria analysis

bluemira/equilibria/analysis.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: 2021-present M. Coleman, J. Cook, F. Franza
+# SPDX-FileCopyrightText: 2021-present I.A. Maione, S. McIntosh
+# SPDX-FileCopyrightText: 2021-present J. Morris, D. Short
+#
+# SPDX-License-Identifier: LGPL-2.1-or-later
+"""Equilibria and Equilibria optimisation analysis tools"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from bluemira.equilibria.equilibrium import MHDState
+    from bluemira.equilibria.optimisation.problem.base import CoilsetOptimisationProblem
+
+
+class EqAnalysis:
+    """Equilibria analysis toolbox"""
+
+    def __init__(self, eq: MHDState):
+        self._eq = eq
+
+    def plot(self):
+        """Plot equilibria"""
+        return self.eq.plot()
+
+
+class COPAnalysis(EqAnalysis):
+    """Coilset Optimisation Problem analysis toolbox"""
+
+    def __init__(self, cop: CoilsetOptimisationProblem):
+        super().__init__(cop.eq)
+        self._cop = cop

bluemira/equilibria/equilibrium.py

@@ -914,7 +914,9 @@ def from_eqdsk(
             filename=filename,
         )
 
-    def to_dict(self, qpsi_calcmode: int = 0) -> dict[str, Any]:
+    def to_dict(
+        self, qpsi_calcmode: int | QpsiCalcMode = QpsiCalcMode.NO_CALC
+    ) -> dict[str, Any]:
         """
         Creates dictionary for equilibrium object, in preparation for saving
         to a file format
@@ -999,14 +1001,15 @@ def to_eqdsk(
         header: str = "BP_equilibria",
         directory: str | None = None,
         filetype: str = "json",
-        qpsi_calcmode: int = 0,
+        qpsi_calcmode: int | QpsiCalcMode = QpsiCalcMode.NO_CALC,
         **kwargs,
     ):
         """
         Writes the Equilibrium Object to an eqdsk file
         """
-        if "eqdsk" in filetype and qpsi_calcmode == 0:
-            qpsi_calcmode = 2
+        qpsi_calcmode = QpsiCalcMode(qpsi_calcmode)
+        if "eqdsk" in filetype and qpsi_calcmode is QpsiCalcMode.NO_CALC:
+            qpsi_calcmode = QpsiCalcMode.ZEROS
 
         super().to_eqdsk(
             self.to_dict(qpsi_calcmode),