如何在numpy中以与Java相同的速度用函数值填充3D数组？

Question

我尝试用以下代码对3D圆柱体的体素进行建模：

import math
import numpy as np

R0 = 500
hz = 1

x = np.arange(-1000, 1000, 1)
y = np.arange(-1000, 1000, 1)
z = np.arange(-10, 10, 1)

xx, yy, zz = np.meshgrid(x, y, z)


def density_f(x, y, z):
    r_xy = math.sqrt(x ** 2 + y ** 2)
    if r_xy <= R0 and -hz <= z <= hz:
        return 1
    else:
        return 0


density = np.vectorize(density_f)(xx, yy, zz)

花了很多分钟来计算。在

相当于次优的Java代码运行10-15秒。在

如何让Python以相同的速度计算体素？在哪里优化？在

Answer 1

这是对迪帕克·赛尼的回答的长篇评论。在

主要的改变是不使用np.meshgrid生成的坐标，它包含不必要的重复。如果可以避免的话，这是不可取的（在内存使用和性能方面）

编码

import numba as nb
import numpy as np

@nb.jit(nopython=True,parallel=True)
def calc_density_2(x, y, z,R0,hz):
    threshold = R0 * R0

    density = 0
    for i in nb.prange(y.shape[0]):
        for j in range(x.shape[0]):
            r_xy = x[j] ** 2 + y[i] ** 2
            for k in range(z.shape[0]):
                if(r_xy <= threshold and -hz <= z[k] <= hz):
                    density+=1

    return density

计时

^{pr2}$

Answer 2

也可以使用函数的编译版本来计算密度。你可以用cython或numba。{a1}作为一个函数，可以很容易地将密度作为一个函数来编译。在

专业人士：

• 您可以编写注释中提到的if条件
• 比ans中提到的numpy版本稍快 @Willem Van Onsem，因为我们必须迭代布尔数组 计算density.astype(int).sum()中的和。在

缺点：

• 写一个丑陋的三级循环。解开了新的美丽。在

代码：

import numba as nb
@nb.jit(nopython=True, cache=True)
def calc_density(xx, yy, zz, R0, hz):
    threshold = R0 * R0
    dimensions = xx.shape

    density = 0
    for i in range(dimensions[0]):
        for j in range(dimensions[1]):
            for k in range(dimensions[2]):
                r_xy = xx[i][j][k] ** 2 + yy[i][j][k] ** 2

                if(r_xy <= threshold and -hz <= zz[i][j][k] <= hz):
                    density+=1
    return density

运行时间：

Willem Van Onsem solution, f2 variant : 1.28s without sum, 2.01 with sum.

Numba solution( calc_density, on second run, to discount the compile time) : 0.48s.

正如评论中所建议的，我们也不需要计算网格。我们可以直接将x, y, z传递给函数。因此：

^{pr2}$

现在，为了公平比较，我们在@Willem Van Onsem的ans中也包含了np.meshgrid的时间。 运行时间：

Willem Van Onsem solution, f2 variant(np.meshgrid time included) : 2.24s

Numba solution( calc_density2, on second run, to discount the compile time) : 0.079s.

Answer 3

请不要使用.vectorize(..)，因为它仍将在Python级别进行处理，因此效率不高。.vectorize()只能作为最后的手段使用，例如，如果函数不能以“bulk”的形式计算，因为它的“结构”太复杂了。在

但是您不需要使用.vectorize在这里，您可以实现您的函数来处理数组：

r_xy = np.sqrt(xx ** 2 + yy ** 2)
density = (r_xy <= R0) & (-hz <= zz) & (zz <= hz)

或者更快一点：

^{pr2}$

这将构造一个2000×2000×20布尔数组。我们可以使用：

intdens = density.astype(int)

构造ints的数组

在这里打印数组非常复杂，但它总共包含2'356'047个1：

>>> density.astype(int).sum()
2356047

基准测试：如果我在本地运行10次，我得到：

>>> timeit(f, number=10)
18.040479518999973
>>> timeit(f2, number=10)  # f2 is the optimized variant
13.287886952000008

所以平均来说，我们在1.3-1.8秒内计算出这个矩阵（包括将它转换成ints）。在