Gfrktyl

I am coding a bot for Starcraft 2, in which many distances have to be calculated every frame.

Here is the part of the library that is being used and I want to improve: https://github.com/Dentosal/python-sc2/blob/develop/sc2/position.py

I built a new class Points that inherits from np.ndarray.
It is not yet connected to the rest of the library, but the functions are done. I removed the functions furthest_to, further_than and so on because the closer-versions are basically the same execpt a -1 or <.

Are these functions implemented in most efficient way? Is there a way to improve the parts that look like this:

find = np.where(np.any(M < distance, axis=1))
selection = np.array([self[i] for i in find[0]])

Any other comments or suggestions are also welcome :)

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)

You should probably have another look at the scipy.spatial module. It provides (hopefully) faster methods for most of these checks using a k-d-tree.

from scipy.spatial import cKDTree



class Points(np.ndarray):

    def __new__(cls, units_or_points):

        obj = np.asarray(units_or_points).view(cls)

        obj.kd_tree = cKDTree(obj)

        return obj



    def closest_to(self, point: Point2) -> Point2:

        """Returns the point of self that is closest to another point."""

        _, i = self.kd_tree.query([[point.x, point.y]])

        return Point2(self[i][0])



    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."""

        selection = self.kd_tree.query_ball_point([point.x, point.y], distance)

        return self[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."""

        selection_close = self.kd_tree.query_ball_point([point.x, point.y], distance1)

        selection_far = self.kd_tree.query_ball_point([point.x, point.y], distance2)

        selection = list(set(selection_far) - set(selection_close))

        return self[selection]



    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."""

        _, indices = self.kd_tree.query([[point.x, point.y]], k=n)

        return self[indices]



    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."""

        pairs = self.kd_tree.query_ball_tree(points.kd_tree, distance)

        return points[[i for closest in pairs for i in closest]]

These are all the ones I could quickly find a way for using the tree. Not included are sort_by_distance_to, n_closest_to_distance and n_closest_to_distance.

In order to test if this is really faster, here are some tests, with the following setup:

np.random.seed(42)

points = np.random.rand(300, 2)

points_graipher = Points(points)

points_op = PointsOP(points)

point = Point2(np.random.rand(2))

points2 = np.random.rand(10, 2)

points2_graipher = Points(points2)

Here PointsOP is you class and Points is the class defined in this answer.

%timeit points_op.closest_to(point)

# 38.3 µs ± 1.35 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit points_graipher.closest_to(point)

# 43.7 µs ± 249 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)



%timeit points_op.closer_than(point, 0.1)

# 39.5 µs ± 238 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)    

%timeit points_graipher.closer_than(point, 0.1)

# 11 µs ± 26 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)



%timeit points_op.in_distance_between(point, 0.1, 0.2)

# 52.9 µs ± 275 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit points_graipher.in_distance_between(point, 0.1, 0.2)

# 21.9 µs ± 180 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)



%timeit points_op.closest_n_points(point, 10)

# 29.5 µs ± 359 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit points_graipher.closest_n_points(point, 10)

# 41.7 µs ± 287 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)



%timeit points_op.in_distance_of_points(points2, 0.1)

# 116 µs ± 727 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit points_graipher.in_distance_of_points(points2_graipher, 0.1)

# 89.2 µs ± 500 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

As you can see for $N = 300$ points there are some methods which are faster with the KDTree (up to four times), some that are basically the same and some that are slower (by up to two times).

To get a feeling how the different approaches, scale, here are some plots. The only thing changing is the number of points. The steps are 30, 300, 3000, 30000.

To summarize, you should check this for some actual cases you have. Depending on the size of points your implementation or this implementation is faster.