使用Numba计算快速凸包

Question

我发现了一个使用Numpy实现计算二维点凸包的好方法。我想将这个函数@njit化，以便在我的其他Numba jitted代码中使用它。但是，由于它使用递归和不支持的Numba功能，我无法修改它使其运行。有人能帮我重写一下吗？

import numpy as np
from numba import njit

def process(S, P, a, b):
    signed_dist = np.cross(S[P] - S[a], S[b] - S[a])
    K = [i for s, i in zip(signed_dist, P) if s > 0 and i != a and i != b]

    if len(K) == 0:
        return (a, b)

    c = max(zip(signed_dist, P))[1]
    return process(S, K, a, c)[:-1] + process(S, K, c, b)

def quickhull_2d(S: np.ndarray) -> np.ndarray:
    a, b = np.argmin(S[:,0]), np.argmax(S[:,0])
    max_index = np.argmax(S[:,0])
    max_element = S[max_index]
    return process(S, np.arange(S.shape[0]), a, max_index)[:-1] + process(S, np.arange(S.shape[0]), max_index, a)[:-1]

示例数据输入和输出

points = np.array([[0, 0], [1, 1], [0.5, 0.5], [0, 1], [1, 0]])
ch = quickhull_2d(points)

print(ch)
[0, 4, 1, 3]

print(points[ch])
[[0. 0.]
 [1. 0.]
 [1. 1.]
 [0. 1.]]

- Hokyjack

---

当我使用输入np.full((3, 3), 2)运行quickhull_2d时，我在process函数内部收到一个错误：The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()。您确定提供的代码在其未修改的形式下是正确的吗？如果是，请提供一个示例输入数组。谢谢。 - Rafnus

当然，抱歉，我忘记添加检查编辑后的帖子。输入数据应该是一个numpy 2D数组，包含X、Y点（形状=（N，2））。 - Hokyjack

我看到的问题是，quickhull函数将不同类型的元组（S - np.ndarray、a - np.int64、max_index - np.int64等）传递给process函数，如果在递归中调用process，则它要么再次输出该元组，要么最终只输出两个元素的列表。而numba函数每次都必须返回相同类型的元组。 - Hokyjack

Answer 1

这段代码存在许多问题，无法在Numba中使用。

首先，返回变量大小的元组在Numba中不可能，因为元组的类型隐含地包括其大小。元组基本上是一种结构化类型而不是列表。有关此问题的更多信息，请参见this post和this one。解决方案基本上是返回列表（慢）或数组（快）。

此外，参数的类型从一个函数到另一个函数会发生变化。实际上，在quickhull_2d中调用process时，P被定义为一个Numpy数组，并且在process本身中调用时，P被定义为一个列表。列表和数组完全不同。在Numba中最好尽可能使用数组，除非您使用列表添加未知数量的项目（既不小也不有界）。

此外，Numba不支持max(zip(signed_dist, P))[1]，而且这种方法效率也不高（对于NumPy代码来说也不太常用）。应该使用P[np.argmax(signed_dist)]。

另外，通用情况下似乎也不支持np.cross，因此需要使用cross2d（来自numba.np.extensions）。

最后，当使用这样的递归函数时，最好指定参数的输入类型，以避免奇怪的错误。这可以通过签名字符串来实现。

生成的代码如下：

import numpy as np
from numba import njit
from numba.np.extensions import cross2d

@njit('(float64[:,:], int64[:], int64, int64)')
def process(S, P, a, b):
    signed_dist = cross2d(S[P] - S[a], S[b] - S[a])
    K = np.array([i for s, i in zip(signed_dist, P) if s > 0 and i != a and i != b], dtype=np.int64)

    if len(K) == 0:
        return [a, b]

    c = P[np.argmax(signed_dist)]
    return process(S, K, a, c)[:-1] + process(S, K, c, b)

@njit('(float64[:,:],)')
def quickhull_2d(S: np.ndarray) -> np.ndarray:
    a, b = np.argmin(S[:,0]), np.argmax(S[:,0])
    max_index = np.argmax(S[:,0])
    max_element = S[max_index]
    return process(S, np.arange(S.shape[0]), a, max_index)[:-1] + process(S, np.arange(S.shape[0]), max_index, a)[:-1]

points = np.array([[0, 0], [1, 1], [0.5, 0.5], [0, 1], [1, 0]])
ch = quickhull_2d(points)

print(ch) # print [0, 4, 1, 3]

请注意，编译时间较慢，执行时间不应过长。这是由于列表（以及运行时性能的临时数组）造成的。下一步是简单地使用数组。坏消息是 Numba 不支持 concatenate（因为一般情况下难以实现，尽管特定情况很容易）。您可以创建一个新数组并复制每个部分（或者更好的方法是：在递归调用期间预分配数组并对其进行切片）。

此外，请注意，任何递归函数都可以使用手动堆栈转换为非递归函数。也就是说，这可能会更慢并使代码更加冗长。然而，这种方法有一些好处：它避免了递归深度过大导致的堆栈溢出，并且如果重写函数以避免其中一个函数调用堆叠，它可能会更快，这要归功于 尾调用优化。

- Jérôme Richard

---

非常感谢您的帮助和付出，以及详细的解释！我对300个随机的2D点进行了一些计时：

python quickhull_2d - 923 µs

scipy ConvexHull - 394 µs

numba quickhull_2d - 43.8 µs

这比Python和Scipy都有很大的改进！ - Hokyjack