python - 在 python 中不使用循环填充 3D 数组

Question

我正在用一个函数填充 3 维数组，该函数取决于其他 1 维数组的值，如下面的代码所示。涉及我的真实数据的代码需要很长时间，因为我的 1 维数组(以及我的 3 维数组)的长度为 100 万。有什么方法可以更快地做到这一点，例如不在 python 中使用循环？

这个想法可能看起来很愚蠢，但我仍然想知道在我的程序中用 C++ 填充这个对象导入代码是否会更快......我是 C++ 的新手，所以我没有尝试过.

import numpy as np
import time

start_time = time.time()
kx = np.linspace(0,400,100)
ky = np.linspace(0,400,100)
kz = np.linspace(0,400,100)

Kh = np.empty((len(kx),len(ky),len(kz)))

for i in range(len(kx)):
    for j in range(len(ky)):
        for k in range(len(kz)):
            if np.sqrt(kx[i]**2+ky[j]**2) != 0:
                Kh[i][j][k] = np.sqrt(kx[i]**2+ky[j]**2+kz[k]**2)
            else:
                Kh[i][j][k] = 1


print('Finished in %s seconds' % (time.time() - start_time))

Answer 1

您可以使用 numba 中的 @njit 装饰器，这是一个高性能的 JIT 编译器。它将时间减少了一个数量级以上。下面是对比和代码。它就像导入 njit 然后使用 @njit 作为函数的装饰器一样简单。 This是官方网站。

我还使用 njit 计算了 1000*1000*1000 数据点的时间，仅用了 17.856173038482666 秒。使用并行版本作为 @njit(parallel=True) 进一步将时间减少到 9.36257791519165 秒。使用正常功能执行相同操作需要几分钟时间。

我还对 njit 和@Bily 在他的 answer 中建议的矩阵运算做了一些时间比较以下。虽然对于高达 700 的点数而言，时间是可比的，但 njit 方法明显胜过大于 700 的点数，如下图所示。

import numpy as np
import time
from numba import njit

kx = np.linspace(0,400,100)
ky = np.linspace(0,400,100)
kz = np.linspace(0,400,100)

Kh = np.empty((len(kx),len(ky),len(kz)))

@njit  # <----- Decorating your function here
def func_njit(kx, ky, kz, Kh):
    for i in range(len(kx)):
        for j in range(len(ky)):
            for k in range(len(kz)):
                if np.sqrt(kx[i]**2+ky[j]**2) != 0:
                    Kh[i][j][k] = np.sqrt(kx[i]**2+ky[j]**2+kz[k]**2)
                else:
                    Kh[i][j][k] = 1
    return Kh                

start_time = time.time()
Kh = func_njit(kx, ky, kz, Kh)
print('NJIT Finished in %s seconds' % (time.time() - start_time))

def func_normal(kx, ky, kz, Kh):
    for i in range(len(kx)):
        for j in range(len(ky)):
            for k in range(len(kz)):
                if np.sqrt(kx[i]**2+ky[j]**2) != 0:
                    Kh[i][j][k] = np.sqrt(kx[i]**2+ky[j]**2+kz[k]**2)
                else:
                    Kh[i][j][k] = 1
    return Kh 

start_time = time.time()
Kh = func_normal(kx, ky, kz, Kh)
print('Normal function Finished in %s seconds' % (time.time() - start_time))

---

NJIT Finished in 0.36797094345092773 seconds
Normal function Finished in 5.540749788284302 seconds

njit与矩阵法的比较