Patching small issue with dr_sep for some equilibria

freegs4e/equilibrium.py

@@ -2405,7 +2405,8 @@ def dr_sep(
     ):
         """
         Computes the radial separation (on the inboard and outboard sides) at vertical position `Z_level`
-        between flux surfaces passing through the lower and upper X-points.
+        between flux surfaces passing through the lower and upper X-points. This requires two X-points to
+        be located within the vessel region.
 
         The inboard dr_sep is defined as the difference in R (radial position) between
         the two flux surfaces intersecting `Z_level` on the inboard side (smaller R values),
@@ -2417,7 +2418,7 @@ def dr_sep(
         Parameters
         ----------
         Z_level : float
-            Vertical height at which to calculate dr_sep [m].
+            Vertical height at which to calculate dr_sep [m]. Defaults to eq.Zgeoemtric().
         sign_flip : bool
             By convention, the function assumes the first X-point corresponds to the lower
             X-point, and the second to the upper X-point. So for an upper single null plasma,
@@ -2439,13 +2440,34 @@ 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, :]
+
+        # if less than two X-points inside wall, throw error
+        if xpts_inside_wall.shape[0] < 2:
+            raise ValueError(
+                f"Expected at least two X-points inside wall, got {xpts_inside_wall.shape[0]}."
+            )
+
+        # get the two closest to magnetic axis
+        mag_axis = self.magneticAxis()
+        distances = np.linalg.norm(
+            xpts_inside_wall[:, 0:2] - mag_axis[0:2], axis=1
+        )
+        closest_two = xpts_inside_wall[np.argsort(distances)[:2]]
 
-        # 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])]
+        # sort those two by Z (puts lowest X-point first)
+        closest_xpts_sorted = closest_two[np.argsort(closest_two[:, 1])]
 
         # find the flux contours for the values of psi_boundary at each X-point
         contour0 = plt.contour(
@@ -2510,12 +2532,11 @@ def dr_sep(
         R_in_out1_stacked = np.vstack(R_in_out1)
 
         # clip rows so that we only obtain the two rows with R values closest to the magnetic axis
-        R_mag = self.Rmagnetic()
         R_in_out0_filtered = R_in_out0_stacked[
-            np.argsort(np.abs(R_in_out0_stacked[:, 0] - R_mag))[:2]
+            np.argsort(np.abs(R_in_out0_stacked[:, 0] - mag_axis[0]))[:2]
         ]
         R_in_out1_filtered = R_in_out1_stacked[
-            np.argsort(np.abs(R_in_out1_stacked[:, 0] - R_mag))[:2]
+            np.argsort(np.abs(R_in_out1_stacked[:, 0] - mag_axis[0]))[:2]
         ]
 
         # sort so that the lower X-point is first