如何加速numpy.all与numpy.nonzero()？海量包围盒点检测优化求助

咱们先拆解下你这段代码慢的核心原因——原代码里的(points > cuboid_min) & (points < cuboid_max)会生成一个**(4000000, 3, 64)**的布尔数组，算下来光这个中间数组就占了近800MB内存！内存带宽被占满，自然计算和后续的np.all、np.nonzero都慢得离谱。

下面给你几个实用的优化方案，按实现复杂度和提升效果排序：

1. 换计算顺序：按点循环，大幅降低内存压力 #

因为你只有64个点，数量极少，完全可以逐个点处理，每次只处理单个点和所有包围盒的判断，这样每个循环的中间数组只有**(4000000, 3)**，内存占用直接降到12MB左右，内存压力骤减，速度会明显提升。

修改后的代码：

import numpy as np

n_points = 64
n_cuboid = 4000000
points = np.random.rand(1, 3, n_points)*512
cuboid_min = np.random.rand(n_cuboid, 3, 1)*512
cuboid_max = cuboid_min + np.random.rand(n_cuboid, 3, 1)*8

# 先把包围盒数组改回更高效的形状（去掉多余的最后一维）
cuboid_min = cuboid_min.reshape(n_cuboid, 3)
cuboid_max = cuboid_max.reshape(n_cuboid, 3)
points = points.reshape(3, n_points)  # 形状(3,64)

inside_list = []
for i in range(n_points):
    point = points[:, i]  # 当前点，形状(3,)
    # 对单个点判断是否在每个包围盒内
    lower_ok = point > cuboid_min
    upper_ok = point < cuboid_max
    inside = np.all(lower_ok & upper_ok, axis=1)
    inside_list.append(inside)

# 合并所有点的结果
inside_cuboid = np.column_stack(inside_list)
indices = np.nonzero(inside_cuboid)

这个改动几乎不需要额外依赖，只是调整了计算顺序，就能让np.all的耗时大幅下降。

2. 用Numba JIT编译，把循环变成机器码 #

如果想进一步榨干CPU性能，试试Numba——它能把Python循环直接编译成机器码，还支持多核并行，完全规避numpy的中间数组开销。

代码示例：

import numpy as np
import numba as nb

n_points = 64
n_cuboid = 4000000
points = np.random.rand(1, 3, n_points)*512
cuboid_min = np.random.rand(n_cuboid, 3, 1)*512
cuboid_max = cuboid_min + np.random.rand(n_cuboid, 3, 1)*8

# 调整数组形状为连续内存布局
cuboid_min = cuboid_min.reshape(n_cuboid, 3)
cuboid_max = cuboid_max.reshape(n_cuboid, 3)
points = points.reshape(3, n_points)

@nb.njit(parallel=True)
def check_inside(points, cuboid_min, cuboid_max):
    n_p = points.shape[1]
    n_c = cuboid_min.shape[0]
    # 预分配结果数组
    result = np.zeros((n_c, n_p), dtype=nb.boolean)
    # 并行处理每个点（每个点的计算独立，适合多核）
    for i in nb.prange(n_p):
        px, py, pz = points[0, i], points[1, i], points[2, i]
        for j in range(n_c):
            min_x, min_y, min_z = cuboid_min[j]
            max_x, max_y, max_z = cuboid_max[j]
            # 直接判断，无中间数组
            result[j, i] = (px > min_x) and (px < max_x) and \
                           (py > min_y) and (py < max_y) and \
                           (pz > min_z) and (pz < max_z)
    return result

inside_cuboid = check_inside(points, cuboid_min, cuboid_max)
indices = np.nonzero(inside_cuboid)

这里用了parallel=True让每个点的计算跑在不同CPU核心上，对于400万包围盒+64个点的场景，速度能比纯numpy快3-5倍甚至更多。

3. 构建空间索引（适合多次查询的场景） #

如果你的需求是多次查询不同的点，而不是只查这64个，那提前给400万个包围盒构建R-Tree空间索引会是最优解——它能把每个点需要检查的包围盒数量从400万降到几十个甚至几个，彻底改变计算量级。

需要先安装rtree库（pip install rtree），代码示例：

import numpy as np
from rtree import index

n_points = 64
n_cuboid = 4000000
points = np.random.rand(1, 3, n_points)*512
cuboid_min = np.random.rand(n_cuboid, 3, 1)*512
cuboid_max = cuboid_min + np.random.rand(n_cuboid, 3, 1)*8

# 调整数组形状
cuboid_min = cuboid_min.reshape(n_cuboid, 3)
cuboid_max = cuboid_max.reshape(n_cuboid, 3)
points = points.reshape(3, n_points).T  # 形状(64,3)

# 构建R-Tree索引
idx = index.Index()
for j in range(n_cuboid):
    # R-Tree需要的边界格式：(min_x, min_y, min_z, max_x, max_y, max_z)
    bounds = (*cuboid_min[j], *cuboid_max[j])
    idx.insert(j, bounds)

# 逐个查询点
inside_cuboid = np.zeros((n_cuboid, n_points), dtype=bool)
for i, point in enumerate(points):
    # 查询包含该点的所有包围盒ID
    match_ids = list(idx.intersection((*point, *point)))
    # 标记这些包围盒为包含当前点
    inside_cuboid[match_ids, i] = True

indices = np.nonzero(inside_cuboid)

构建索引的过程可能需要几秒，但后续每个点的查询都能在毫秒级完成，适合需要反复查询的场景。

内容的提问来源于stack exchange，提问作者f. c.