如何用纯Numpy无循环计算网格内3D点关联值的均值？

问题描述 #

我有一批3D点数据，每个点的前两个值对应图像坐标，第三个值是该点的关联数值。需要在图像上划分2D网格，生成一个数组，其中每个元素对应网格单元内所有点的关联数值的均值。

已用Numpy实现该功能，但代码包含两个for循环，希望移除循环实现纯Numpy版本，尝试过np.mgrid但没找到可行方法，寻求帮助。

初始实现代码：

import numpy as np

points_range = np.array([2, 5, 1])
points = np.random.random((100, 3))  # points[:, i] 对应 x, y, z（i=0,1,2）
points *= points_range

x_steps = 10  # 将x轴空间划分为10列
y_steps = 15  # 将y轴空间划分为15行
x_step_size = points_range[0] / x_steps
y_step_size = points_range[1] / y_steps

# 尝试使用np.mgrid，但未找到合适用法
X, Y = np.mgrid[0:points_range[0]:x_step_size, 0:points_range[1]:y_step_size]

# 当前可行的带循环实现
means = [[0 for _ in range(y_steps)] for _ in range(x_steps)]
for x_step_id in range(x_steps):
    for y_step_id in range(y_steps):
        # 筛选当前网格单元内的点
        p = points[(points[:, 0] >= x_step_size * x_step_id) &
                   (points[:, 0] < x_step_size * (x_step_id + 1)) &
                   (points[:, 1] >= y_step_size * y_step_id) &
                   (points[:, 1] < y_step_size * (y_step_id + 1))]
        # 计算z值的均值
        mean_z = np.mean(p[:, 2])
        means[x_step_id][y_step_id] = 0 if np.isnan(mean_z) else mean_z

解决方案 #

以下是四种实现方式的性能对比（测试环境：1000万点数据，15×15网格，单位：秒）：

Method for_loop_1 took 37.779709815979004
Method for_loop_2 took 14.891589879989624
Method for_plus_numpy took 16.47166681289673
Method full_numpy took 1.1438612937927246

四种实现代码 #

import numpy as np
import time

points_range = np.array([2, 5, 1])
points = np.random.random((10000000, 3))  # points[:, i] 对应 x, y, z（i=0,1,2）
points *= points_range

x_steps = 15  # 将x轴空间划分为15列
y_steps = 15  # 将y轴空间划分为15行
x_step_size = points_range[0] / x_steps
y_step_size = points_range[1] / y_steps


methods = []
def for_loop_1():
    start = time.time()
    means = [[0 for _ in range(y_steps)] for _ in range(x_steps)]
    n = np.zeros_like(means)
    for p in points:
        tile_x = int(p[0] // x_step_size)
        tile_y = int(p[1] // y_step_size)
        means[tile_x][tile_y] = (p[2] + n[tile_x][tile_y] * means[tile_x][tile_y]) / (n[tile_x][tile_y] + 1)
        n[tile_x][tile_y] += 1
    stop = time.time()
    return means, stop - start
methods.append(for_loop_1)

def for_loop_2():
    start = time.time()
    sums = np.zeros((x_steps, y_steps))
    counts = np.zeros((x_steps, y_steps))
    for point in points:
        tile_index_x = int(point[0] // x_step_size)
        tile_index_y = int(point[1] // y_step_size)
        sums[tile_index_x, tile_index_y] += point[2]
        counts[tile_index_x, tile_index_y] += 1
    means = np.where(counts > 0, sums / counts, 0)
    stop = time.time()
    return means, stop - start
methods.append(for_loop_2)

# 原带循环的Numpy实现
def for_plus_numpy():
    start = time.time()
    means = [[0 for _ in range(y_steps)] for _ in range(x_steps)]
    for x_step_id in range(x_steps):
        for y_step_id in range(y_steps):
            # 筛选当前网格单元内的点
            p = points[(points[:, 0] >= x_step_size * x_step_id) &
                       (points[:, 0] < x_step_size * (x_step_id + 1)) &
                       (points[:, 1] >= y_step_size * y_step_id) &
                       (points[:, 1] < y_step_size * (y_step_id + 1))]
            # 计算z值的均值
            mean_z = np.mean(p[:, 2])
            means[x_step_id][y_step_id] = 0 if np.isnan(mean_z) else mean_z
    stop = time.time()
    return means, stop - start
methods.append(for_plus_numpy)


# 纯Numpy无循环实现
def full_numpy():
    start = time.time()
    edges = np.linspace(0, points_range[0], x_steps+1), np.linspace(0, points_range[1], y_steps+1)
    sums, _, _ = np.histogram2d(points[:, 0], points[:, 1], edges, weights=points[:, 2])
    counts, _, _ = np.histogram2d(points[:, 0], points[:, 1], edges)
    means = np.where(counts > 0, sums / counts, 0)
    stop = time.time()
    return means, stop - start
methods.append(full_numpy)


for method in methods:
    means, duration = method()
    print("Method ", method.__name__ + " took ", duration, sep="")

内容的提问来源于stack exchange，提问作者Hedwin Bonnavaud