The Big Programming Thread - Page 984

from typing import Any, Dict, List, Optional, Set, Tuple, Union  # for mypy type checking

import numpy as np
from scipy.spatial.distance import cdist

from position import Point2


class Points(np.ndarray):
    def __new__(cls, units_or_points):
        obj = np.asarray(units_or_points).view(cls)
        return obj

    def closest_to(self, point: Point2) -> Point2:
        """Returns the point of self that is closest to another point."""
        if point in self:
            return Point2(tuple(point))
        deltas = self - point
        distances = np.einsum("ij,ij->i", deltas, deltas)
        result = self[np.argmin(distances)]
        return Point2(tuple(result))

    def closer_than(self, point: Point2, distance: Union[int, float]) -> "Points":
        """Returns a new Points object with all points of self that
        are closer than distance to point."""
        position = np.array([point])
        M = cdist(self, position)
        find = np.where(np.all(M < distance, axis=1))
        selection = np.array([self[i] for i in find[0]])
        return Points(selection)

    def in_distance_between(
        self, point: Point2, distance1: Union[int, float], distance2: Union[int, float]
    ) -> "Points":
        """Returns a new Points object with all points of self
        that are between distance1 and distance2 away from point."""
        p = np.array([point])
        M = cdist(self, p)
        find = np.where(np.any(np.logical_and(distance1 < M, M < distance2), axis=1))
        selection = np.array([self[i] for i in find[0]])
        return Points(selection)

    def sort_by_distance_to(self, point: Point2, reverse: bool = False) -> "Points":
        """Returns a new Points object with all points of self sorted by distance to point.
        Ordered from smallest to biggest distance. Reverse order with keyword reverse=True."""
        deltas = self - point
        distances = (1 if reverse else -1) * np.einsum("ij,ij->i", deltas, deltas)
        result = self[distances.argsort()[::-1]]
        return Points(result)

    def closest_n_points(self, point: Point2, n: int) -> "Points":
        """Returns a new Points object with the n points of self that are closest to point."""
        deltas = self - point
        distances = np.einsum("ij,ij->i", deltas, deltas)
        result = (self[distances.argsort()[::-1]])[-n:]
        return Points(result)

    def in_distance_of_points(self, points: "Points", distance: Union[int, float]) -> "Points":
        """Returns a new Points object with every point of self that
        is in distance of any point in points."""
        M = cdist(self, points)
        find = np.where(np.any(M < distance, axis=1))
        selection = np.array([self[i] for i in find[0]])
        return Points(selection)

    def n_closest_to_distance(self, point: Point2, distance: Union[int, float], n: int) -> "Points":
        """Returns a new Points object with the n points of self
        which calculated distance to point is closest to distance."""
        deltas = self - point
        distances = np.absolute(distance - np.einsum("ij,ij->i", deltas, deltas))
        result = (self[distances.argsort()[::-1]])[-n:]
        return Points(result)