快速替换有条件设置数组元素的方法 [英] Fast alternative to conditionally set an array elements

问题描述

我有两个给定的3d数组 x_dist 和 y_dist ，每个数组的形状为(36,50,50). x_dist 和 y_dist 中的元素的类型为 np.float32 ，可以为正，负或零.我需要创建一个新的数组 res_array ，在这里我将所有索引处的值都设置为(1-y_dist)*(x_dist)，条件((x_dist< = 0)|((x_dist> 0)&(y_dist>(1 + x_dist))))是 True .我当前的实现如下.

I have two given 3d arrays x_dist and y_dist each of shape (36,50,50). Elements in x_dist and y_dist are of type np.float32 and can be positive, negative or zero. I need to create a new array res_array, where I set its value to (1-y_dist)*(x_dist) at all indexes except for where the condition ((x_dist <= 0) | ((x_dist > 0) & (y_dist > (1 + x_dist)))) is True. My current implementation is as follows.

res_array  = (1-y_dist)*(x_dist)
res_array[((x_dist <= 0) | ((x_dist > 0) & (y_dist > (1 + x_dist))))] = 0.0

但是，我需要运行包含此代码的代码，这要花数千次，并且我确信有一种更聪明，更快捷的方法来执行此操作.您能在这里帮我获得性能上更好的代码或单行代码吗?

However, I need to run the code that contains this code snipet thousands of time and I am sure there is a smarter and more faster way to do the same. Can you please help me here to get a performance wise better code or one-liner?

感谢您的帮助！

推荐答案

Numba JIT 可以有效地做到这一点.这是一个实现:

Numba JIT can be used to do that efficiently. Here is an implementation:

@njit
def fastImpl(x_dist, y_dist):
    res_array = np.empty(x_dist.shape)
    for z in range(res_array.shape[0]):
        for y in range(res_array.shape[1]):
            for x in range(res_array.shape[2]):
                xDist = x_dist[z,y,x]
                yDist = y_dist[z,y,x]
                if xDist > 0.0 and yDist <= (1.0 + xDist):
                    res_array[z,y,x] = (1.0 - yDist) * xDist
    return res_array

以下是随机输入矩阵的性能结果:

Here are performance results on random input matrices:

Original implementation: 494 µs ± 6.23 µs per loop (mean ± std. dev. of 7 runs, 500 loops each)
New implementation: 37.8 µs ± 236 ns per loop (mean ± std. dev. of 7 runs, 500 loops each)

新的实现大约快13倍(不考虑编译/预热时间).

The new implementation is about 13 times faster (without taking into account the compilation/warm-up time).

这篇关于快速替换有条件设置数组元素的方法的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！