逻辑索引python

Question

我正在努力提高 Python 中逻辑索引的速度。所以，目前我必须绘制一些热图，为此我将输入数据划分为指定数量的 x 和 y 箱，然后通过函数 return_val，我使用逻辑索引来计算平均值给定 bin 中的值

当我的 bin 大小很小时，这很有效，但是当我尝试增加 bin 大小时，比如 100x100，那么程序会变慢很多

我知道可以通过使用 Python 中的 stats.binned_statistic_2d 函数来提高速度。但是，我想了解如何优化我当前的代码以使平均过程更快

import numpy as np

arr_len = 932826
x = np.random.uniform(low=0, high=4496, size=arr_len)
y = np.random.uniform(low=-74, high=492, size=arr_len)
z = np.random.uniform(low=-30, high=97, size=arr_len)

# Check points
bin_x = 10
bin_y = 10
x1 = np.linspace(x.min(), x.max(), bin_x)
y1 = np.linspace(y.min(), y.max(), bin_y)


def return_val(x, y, z, x1, y1, i, j):
    idx = np.logical_and(np.logical_and(x > x1[i - 1], x < x1[i]), np.logical_and(y > y1[j - 1], y < y1[j]))
    if np.count_nonzero(idx) == 0:
        return np.nan
    else:
        return np.mean(z[idx])


z1 = np.zeros((len(x1), len(y1)))
for i in range(1, len(x1)):
    for j in range(1, len(y1)):
        z1[i - 1, j - 1] = return_val(x, y, z, x1, y1, i, j)

z1 = z1.transpose()

Answer 1

代码的一半时间花在隐式分配临时数组（由于logical_and 和比较运算符），另一半时间花在慢速嵌套循环调用使用慢速 CPython 解释器 的函数。解决这些问题的一种方法是简单地使用 Numba 的 JIT，使用不带临时数组的无分支操作并使用 并行性。这是一个例子：

import numpy as np
import numba as nb

arr_len = 932826
x = np.random.uniform(low=0, high=4496, size=arr_len)
y = np.random.uniform(low=-74, high=492, size=arr_len)
z = np.random.uniform(low=-30, high=97, size=arr_len)

# Check points
bin_x = 10
bin_y = 10
x1 = np.linspace(x.min(), x.max(), bin_x)
y1 = np.linspace(y.min(), y.max(), bin_y)

@nb.njit('float64(float64[::1], float64[::1], float64[::1], float64[::1], float64[::1], int32, int32)')
def return_val(x, y, z, x1, y1, i, j):
    count = 0
    s = 0.0
    # Branchless mean
    for k in range(len(x)):
        valid = (x[k] > x1[i - 1]) & (x[k] < x1[i]) & (y[k] > y1[j - 1]) & (y[k] < y1[j])
        s += z[k] * valid
        count += valid
    if count == 0:
        return np.nan
    else:
        return s / count

@nb.njit('float64[:,:](float64[::1], float64[::1], float64[::1], float64[::1], float64[::1])', parallel=True)
def compute(x, y, z, x1, y1):
    z1 = np.zeros((len(x1), len(y1)))
    for i in nb.prange(1, len(x1)):
        for j in range(1, len(y1)):
            z1[i - 1, j - 1] = return_val(x, y, z, x1, y1, i, j)
    return z1

z1 = compute(x, y, z, x1, y1)

上面的代码在我的机器上快了 11 倍。它可以通过处理循环进一步改进，以便计算可以更加缓存友好。