Merge pull request #98 from tbody-cfs/patch-1

Prevent out-of-bounds error in find_psisurface near axis

freegs/critical.py

@@ -358,13 +358,22 @@ def find_psisurface(eq, psifunc, r0, z0, r1, z1, psival=1.0, n=100, axis=None):
         # Only one value
         ind = argmax(pnorm > psival)
 
-        # Edited by Bhavin 31/07/18
-        # Changed 1.0 to psival in f
-        # make f gradient to psival surface
-        f = (pnorm[ind] - psival) / (pnorm[ind] - pnorm[ind - 1])
-
-        r = (1.0 - f) * r[ind] + f * r[ind - 1]
-        z = (1.0 - f) * z[ind] + f * z[ind - 1]
+        if ind == 0:
+            # If the point is very close to the magnetic axis, don't
+            # try to do extrapolation.
+            r = r[ind]
+            z = z[ind]
+        else:
+            # Edited by Bhavin 31/07/18
+            # Changed 1.0 to psival in f
+            # make f gradient to psival surface
+            f = (pnorm[ind] - psival) / (pnorm[ind] - pnorm[ind - 1])
+
+            # Interpolate between points
+            r = (1.0 - f) * r[ind] + f * r[ind - 1]
+            z = (1.0 - f) * z[ind] + f * z[ind - 1]
+
+            if f > 1.0: warn(f"find_psisurface has encountered an extrapolation. This will probably result in a point where you don't expect it.")
 
     if axis is not None:
         axis.plot(r, z, "bo")
@@ -505,8 +514,8 @@ def find_safety(
                 psifunc,
                 r0,
                 z0,
-                r0 + 8.0 * sin(theta),
-                z0 + 8.0 * cos(theta),
+                r0 + np.ptp(eq.R) * sin(theta),
+                z0 + np.ptp(eq.Z) * cos(theta),
                 psival=psin,
                 axis=axis,
             )