From 60b4d738ffad24dd2af83602fc4c96fa687d6fdc Mon Sep 17 00:00:00 2001
From: Tom Body <tbody@cfs.energy>
Date: Fri, 16 Dec 2022 11:30:25 -0500
Subject: [PATCH] Pull changes from master

Fix invalid escape sequences

Try/catch for potential UnboundLocalError

Prevent an UnboundLocalError

Add __getitem__ for Circuit

Tidy up optimiser

Fix error where find_psisurface will return random point if approaching magnetic axis

Fix on-axis q calculation
---
 freegs/coil.py          |  1 +
 freegs/critical.py      |  1 +
 freegs/equilibrium.py   | 14 +++++++++++---
 freegs/gradshafranov.py |  2 +-
 freegs/machine.py       | 26 --------------------------
 freegs/optimise.py      |  4 ++--
 freegs/optimiser.py     |  2 +-
 freegs/picard.py        |  2 +-
 8 files changed, 18 insertions(+), 34 deletions(-)

diff --git a/freegs/coil.py b/freegs/coil.py
index 963dcb1..26f4589 100644
--- a/freegs/coil.py
+++ b/freegs/coil.py
@@ -27,6 +27,7 @@
 from .gradshafranov import Greens, GreensBr, GreensBz, mu0
 import numpy as np
 import numbers
+import warnings
 from shapely.geometry import Point
 import warnings
 
diff --git a/freegs/critical.py b/freegs/critical.py
index f79248c..c869d8b 100644
--- a/freegs/critical.py
+++ b/freegs/critical.py
@@ -524,6 +524,7 @@ def find_safety(
     # Get variables for loop integral around flux surface
     r = psisurf[:, :, 0]
     z = psisurf[:, :, 1]
+
     fpol = eq.fpol(psirange[:]).reshape(npsi, 1)
     Br = eq.Br(r, z)
     Bz = eq.Bz(r, z)
diff --git a/freegs/equilibrium.py b/freegs/equilibrium.py
index 79f3508..4eacab8 100644
--- a/freegs/equilibrium.py
+++ b/freegs/equilibrium.py
@@ -346,10 +346,18 @@ def q(self, psinorm=None, npsi=100):
             psinorm = linspace(1.0 / (npsi + 1), 1.0, npsi, endpoint=False)
             return psinorm, critical.find_safety(self, psinorm=psinorm)
 
-        result = critical.find_safety(self, psinorm=psinorm)
+        elif np.any((psinorm < 0.01) | (psinorm > 0.99)):
+            psinorm_inner = np.linspace(0.01, 0.99, num=npsi)
+            q_inner = critical.find_safety(self, psinorm=psinorm_inner)
+
+            interp = interpolate.interp1d(psinorm_inner, q_inner, kind = "quadratic", fill_value = "extrapolate")
+            result = interp(psinorm)
+        else:
+            result = critical.find_safety(self, psinorm=psinorm)
+        
         # Convert to a scalar if only one result
-        if len(result) == 1:
-            return result[0]
+        if np.size(result) == 1:
+            return float(result)
         return result
 
     def tor_flux(self, psi=None):
diff --git a/freegs/gradshafranov.py b/freegs/gradshafranov.py
index 2b63a56..f8c8c04 100644
--- a/freegs/gradshafranov.py
+++ b/freegs/gradshafranov.py
@@ -38,7 +38,7 @@ class GSElliptic:
     Represents the Grad-Shafranov elliptic operator
 
     .. math::
-        \Delta^* = R^2 \nabla\cdot\frac{1}{R^2}\nabla
+        \\Delta^* = R^2 \\nabla\\cdot\\frac{1}{R^2}\\nabla
 
     which is
 
diff --git a/freegs/machine.py b/freegs/machine.py
index 3062634..7b97fa2 100644
--- a/freegs/machine.py
+++ b/freegs/machine.py
@@ -98,32 +98,6 @@ def __getitem__(self, name):
             if label == name:
                 return coil
         raise KeyError("Circuit does not contain coil with label '{0}'".format(name))
-    
-    @property
-    def current(self) -> float:
-        return self._current
-
-    @current.setter
-    def current(self, value: float):
-        """
-        When updating the circuit current, also update the current in the coils in the circuit (with multipliers)
-        """
-        self._current = value
-        for _, coil, multiplier in self.coils:
-            coil.current = multiplier * value
-
-    @property
-    def control(self) -> bool:
-        return self._control
-
-    @control.setter
-    def control(self, value: bool):
-        """
-        When updating the circuit control switch, also update the control switch in the coils in the circuit
-        """
-        self._control = value
-        for _, coil, _ in self.coils:
-            coil.control = value
 
     def psi(self, R, Z):
         """
diff --git a/freegs/optimise.py b/freegs/optimise.py
index f2e842b..cbe1942 100644
--- a/freegs/optimise.py
+++ b/freegs/optimise.py
@@ -27,7 +27,7 @@
 import matplotlib.pyplot as plt
 from freegs.plotting import plotEquilibrium
 
-from math import sqrt
+from numpy import sqrt, inf
 
 # Measures which operate on Equilibrium objects
 
@@ -199,7 +199,7 @@ def solve_and_measure(eq):
             return measure(eq)
         except:
             # Solve failed.
-            return float("inf")
+            return float(inf)
 
     # Call the generic optimiser,
     return optimiser.optimise(
diff --git a/freegs/optimiser.py b/freegs/optimiser.py
index 2c6ea49..1701559 100644
--- a/freegs/optimiser.py
+++ b/freegs/optimiser.py
@@ -24,7 +24,7 @@
 along with FreeGS.  If not, see <http://www.gnu.org/licenses/>.
 """
 
-import random
+from numpy import random
 import copy
 import bisect
 import sys
diff --git a/freegs/picard.py b/freegs/picard.py
index 81474d1..94d40ee 100644
--- a/freegs/picard.py
+++ b/freegs/picard.py
@@ -93,7 +93,7 @@ def solve(
     bndry = 0.0
 
     # Set an initial value for bndry_change (set high to prevent immediate convergence)
-    bndry_change = 10.0
+    bndry_change = np.inf
 
     # Plasma assumed to not be limited at first
     has_been_limited = False