From 5186d75e1e3bb7d96243720b38c9a625b3e1cc1a Mon Sep 17 00:00:00 2001 From: kpentland Date: Fri, 10 Oct 2025 10:55:12 +0100 Subject: [PATCH] patching small issue with dr_sep for some equilibria --- freegs4e/equilibrium.py | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/freegs4e/equilibrium.py b/freegs4e/equilibrium.py index 3a91ec2..4a21718 100644 --- a/freegs4e/equilibrium.py +++ b/freegs4e/equilibrium.py @@ -2439,13 +2439,24 @@ def dr_sep( psi_bndry = self.psi_bndry xpts = self.xpt - # compute absolute difference between each point's ψ and psi_bndry and then - # get indices of the two smallest differences - closest_indices = np.argsort(np.abs(xpts[:, 2] - psi_bndry))[:2] + # build polygon of the tokamak wall + polygon = sh.Polygon( + np.array([self.tokamak.wall.R, self.tokamak.wall.Z]).T + ) + + # find x-points inside the wall + mask = np.array( + [polygon.contains(sh.Point(x, y)) for x, y in xpts[:, 0:2]] + ) + + # select these x-points + xpts_inside_wall = xpts[mask, :] - # extract the corresponding rows (two X-points closest to psi_bndry, then sort by lowest z coord z-point) - closest_xpts = xpts[closest_indices] - closest_xpts_sorted = closest_xpts[np.argsort(closest_xpts[:, 1])] + # get indices of the two cloest to magnetic axis + closest_xpts_idxs = np.argsort( + np.linalg.norm(xpts_inside_wall[:, 0:2] - self.opt[0, 0:2], axis=1) + ) + closest_xpts_sorted = xpts_inside_wall[closest_xpts_idxs[0:2], :] # find the flux contours for the values of psi_boundary at each X-point contour0 = plt.contour(