From f04244f6a2150a5b041c031876c65729affd164d Mon Sep 17 00:00:00 2001 From: Tom Body <116750897+tbody-cfs@users.noreply.github.com> Date: Thu, 4 May 2023 14:27:09 -0400 Subject: [PATCH] Prevent out-of-bounds error in find_psisurface near axis --- freegs/critical.py | 27 ++++++++++++++++++--------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/freegs/critical.py b/freegs/critical.py index 5563f85..f79248c 100644 --- a/freegs/critical.py +++ b/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, )