在python中不使用循环填充3dmarray

Question

我用一个依赖于其他1D数组值的函数填充一个3D数组，如下面的代码所示。涉及我的真实数据的代码需要花费很长时间，因为我的1d数组（因此我的3D数组）的长度是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

您可以使用@njit中的@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))

---

^{pr2}$

njit与矩阵法的比较

Answer 2

只要有可能，就使用基本的矩阵运算来代替矩阵运算。在

import numpy as np
import 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)))

def func_matrix_operation(kx, ky, kz, _):
  kx_ = np.expand_dims(kx ** 2, 1) # shape: (100, 1)
  ky_ = np.expand_dims(ky ** 2, 0) # shape: (1, 100)
  # Make use of broadcasting such that kxy[i, j] = kx[i] ** 2 + ky[j] ** 2     
  kxy = kx_ + ky_ # shape: (100, 100)
  kxy_ = np.expand_dims(kxy, 2) # shape: (100, 100, 1)
  kz_ = np.reshape(kz ** 2, (1, 1, len(kz))) # shape: (1, 1, 100)
  kxyz = kxy_ + kz_ # kxyz[i, j, k] = kx[i] ** 2 + ky[j] ** 2 + kz[k] ** 2

  kh = np.sqrt(kxyz)
  kh[kxy == 0] = 1
  return kh

start_time = time.time()
Kh1 = func_matrix_operation(kx, ky, kz, Kh)
print('Matrix operation 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()
Kh2 = func_normal(kx, ky, kz, Kh)
print('Normal function Finished in %s seconds' % (time.time() - start_time))

assert np.array_equal(Kh1, Kh2)

输出为：

^{pr2}$