Source code for torx.geometry.polygon_2d_m

"""Implementation of a 2D closed polygon class."""
import numpy as np
from skimage.measure import points_in_poly
from torx.autodoc_decorators_m import autodoc_class



@autodoc_class
class Polygon2D:
    """2D polygon class on a poloidal plane."""



    def __init__(self, x_points: np.ndarray, y_points: np.ndarray,
                 invert_polygon: bool=False):
        """
        Initialize the polygon with the vertices.

        The polygon is treated as closed, though the first and last points do
        not have to be the same.
        If invert_polygon is true, points outside the polygon are considered
        "inside," i.e. will return True from the routine points_inside.
        """
        self.x_points = np.array(x_points).flatten()
        self.y_points = np.array(y_points).flatten()
        self.invert_polygon = invert_polygon

        assert self.x_points.shape == self.y_points.shape
        assert len(self.x_points) == self.y_points.size




    def __getitem__(self, index):
        """
        Allow array indexing.

        Like:
            polygon[i] → (x, y) at index i
            polygon[i, 0] → x at index i
            polygon[i, 1] → y at index i
        """
        if isinstance(index, tuple):
            i, j = index
            if j == 0:
                return self.x_points[i]
            elif j == 1:
                return self.y_points[i]
            else:
                raise IndexError("Index must be 0 or 1 for second dimension")
        else:
            # Handle single index
            x = self.x_points[index]
            y = self.y_points[index]

            return np.column_stack([x, y])




    def flip_Z(self):
        """
        Flips the vertical direction.

        To interface with NUMERICAL equilibria which use this for field
        reversal. Also reverses the vertex arrays to
        ensure that the points remain defined in the same counter-/clockwise
        orientation as before
        """
        self.x_points = self.x_points[::-1]
        self.y_points = -self.y_points[::-1]


    @property
    def verts(self):
        """
        Return the point vertices in a (M, 2) array.

        Compatible with skimage.measure.points_in_poly
        """
        return np.column_stack((self.x_points, self.y_points))



    def points_inside(self, x_tests, y_tests):
        """Check whether points are inside the polygon."""
        assert x_tests.shape == y_tests.shape
        original_shape = x_tests.shape

        points = np.column_stack((x_tests.flatten(), y_tests.flatten()))

        mask = points_in_poly(points=points,
                              verts=self.verts).reshape(original_shape)

        if not (self.invert_polygon):
            return mask
        else:
            return np.logical_not(mask)




    def point_inside(self, x_test, y_test):
        """Return whether a test point is within the polygon."""
        return self.points_inside(np.atleast_1d(x_test),
                                  np.atleast_1d(y_test))[0]




    def signed_area(self):
        """
        Return the signed area of a non-intersecting polygon.

        If the sign is positive, points are counterclockwise. Otherwise, points
        are clockwise.

        See https://mathworld.wolfram.com/PolygonArea.html
        """
        area = np.sum( self.x_points * np.roll(self.y_points, -1) \
                     - self.y_points * np.roll(self.x_points, -1) \
                     ) / 2.0

        return area




    def repeat_first_point(self):
        """
        Conditionally repeat the first point at the end.

        If the polygon is not already closed, adds the final edge to link the
        first and last points. Useful for plotting.
        """
        if not np.isclose(self.x_points[0], self.x_points[-1]) or \
           not np.isclose(self.y_points[0], self.y_points[-1]):
            self.x_points = np.append(self.x_points, self.x_points[0])
            self.y_points = np.append(self.y_points, self.y_points[0])




    def to_ordered(self, start_idx: int=0, k: int=20,
                   n_interp: int=0, smooth: float=None):
        """
        Return an OrderedPolygon2D version of this polygon.

        Orders the points counterclockwise starting from the point with
        index start_idx using a nearest-neighbor algorithm with up to k
        nearest neighbors.
        """
        from torx.geometry.ordered_polygon_2d_m import OrderedPolygon2D

        return OrderedPolygon2D(self.x_points, self.y_points,
                                start_idx=start_idx, k=k,
                                n_interp=n_interp, smooth=smooth)




    @classmethod
    def from_2d_array(cls, data_slice: np.ndarray):
        """Create an instance from a single 2D array slice (npoints, xy)."""
        if data_slice.shape[-1] != 2:
            raise ValueError("The last dimension (dim_xy) must have length 2.")

        x_points = data_slice[:, 0]
        y_points = data_slice[:, 1]

        return cls(x_points, y_points)