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

请不要使用.vectorize(..)，因为它仍然会在Python级别进行处理，所以效率不高。 .vectorize()仅应在万不得已的情况下使用，例如，由于“结构”太复杂而无法在“批量”中计算该函数。

但是您无需在此处使用.vectorize，您可以通过以下方式实现函数来处理数组：

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

甚至更快：

r_xy = xx * xx + yy * yy
density = (r_xy <= R0 * R0) & (-hz <= zz) & (zz <= hz)

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

intdens = density.astype(int)

构造一个int个数组。

在此处打印数组非常麻烦，但是总共包含2'356'047个：

>>> 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秒内计算出该矩阵（包括将其转换为int）。

Answer 2

您还可以使用函数的编译版本来计算密度。您可以使用cython或numba。我使用numba来在ans中编译密度计算函数，因为它与放入装饰器一样容易。

专业人士：

您可以写评论中提到的if条件
比ans中提到的numpy版本快一点 @Willem Van Onsem，因为我们必须遍历布尔数组以计算density.astype(int).sum()中的总和。

缺点：

编写一个丑陋的三级循环。放松了单一衬板numpy解决方案的美丽。

代码：

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）解决方案，f2变体：不加和为1.28s，加和为2.01。

Numba解决方案（第二次运行calc_density，以减少编译时间）：0.48s。

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

@nb.jit(nopython=True, cache=True)
def calc_density2(x, y, z, R0, hz):
    threshold = R0 * R0
    dimensions = len(x), len(y), len(z)

    density = 0
    for i in range(dimensions[0]):
        for j in range(dimensions[1]):
            for k in range(dimensions[2]):

                r_xy = x[i] ** 2 + y[j] ** 2
                if(r_xy <= threshold and -hz <= z[k] <= hz):
                    density+=1
    return density

现在，为了公平起见，我们还将np.meshgrid的时间包含在@Willem Van Onsem的ans中。 运行时间：

Willem Van Onsem解决方案，f2变体（包括np.meshgrid时间）：2.24s

Numba解决方案（第二次运行calc_density2，以减少编译时间）：0.079s。

Answer 3

这是对Deepak Saini回答的冗长评论。

主要更改是不使用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

时间

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)

#after the first call (compilation overhead)
#calc_density_2          9.7 ms
#calc_density_2 parallel 3.9 ms
#@Deepak Saini           115 ms

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

3 个答案: