使用 python 制作 3D 斑点的更快方法?

faster way to make a 3D blob with python?

是否有更好的方法来制作 3D 密度函数?

def make_spot_3d(bright, spread, x0,y0,z0):
    # Create x and y indices
    x = np.linspace(-50, 50, 200)
    y = np.linspace(-50, 50, 200)
    z = np.linspace(-50, 50, 200)

    X, Y, Z = np.meshgrid(x, y, z)
    Intensity = np.uint16(bright*np.exp(-((X-x0)/spread)**2
                                        -((Y-y0)/spread)**2
                                        -((Z-z0)/spread)**2))

    return Intensity

该函数可以生成一个可以用mayavi绘制的3D numpy数组

然而,当该函数用于生成一组斑点 (~100) 时,如下所示:

Spots = np.asarray([make_spot_3d(100,2, *loc) for loc in locations])
cluster = np.sum(Spots, axis=0)

产量例如:

执行时间大约为 1 分钟 (cpu i5);我打赌这会更快。

一个明显的改进是使用广播在 'sparse' 网格而不是完整的 meshgrid 上评估强度函数,例如:

X, Y, Z = np.meshgrid(x, y, z, sparse=True)

这在我的机器上将运行时间减少了大约 4 倍:

%timeit make_spot_3d(1., 1., 0, 0, 0)
1 loops, best of 3: 1.56 s per loop 

%timeit make_spot_3d_ogrid(1., 1., 0, 0, 0)
1 loops, best of 3: 359 ms per loop

您可以通过向量化位置、分布和亮度的计算来摆脱列表理解中涉及的开销,例如:

def make_spots(bright, spread, x0, y0, z0):

    # Create x and y indices
    x = np.linspace(-50, 50, 200)
    y = np.linspace(-50, 50, 200)
    z = np.linspace(-50, 50, 200)

    # this will broadcast out to an (nblobs, ny, nx, nz) array
    dx = x[None, None, :, None] - x0[:, None, None, None]
    dy = y[None, :, None, None] - y0[:, None, None, None]
    dz = z[None, None, None, :] - z0[:, None, None, None]
    spread = spread[:, None, None, None]
    bright = bright[:, None, None, None]

    # we can save time by performing the exponentiation over 2D arrays
    # before broadcasting out to 4D, since exp(a + b) == exp(a) * exp(b)
    s2 = spread * spread
    a = np.exp(-(dx * dx) / s2)
    b = np.exp(-(dy * dy) / s2)
    c = np.exp(-(dz * dz) / s2)

    intensity = bright * a * b * c

    return intensity.astype(np.uint16)

其中 brightspreadx0y0z0 是一维向量。这将生成一个 (nblobs, ny, nx, nz) 数组,然后您可以对第一个轴求和。根据生成的 blob 数量以及评估它们的网格有多大,创建此中间数组在内存方面可能会变得非常昂贵。

另一种选择是初始化单个 (ny, nx, nz) 输出数组并就地计算总和:

def sum_spots_inplace(bright, spread, x0, y0, z0):

    # Create x and y indices
    x = np.linspace(-50, 50, 200)
    y = np.linspace(-50, 50, 200)
    z = np.linspace(-50, 50, 200)

    dx = x[None, None, :, None] - x0[:, None, None, None]
    dy = y[None, :, None, None] - y0[:, None, None, None]
    dz = z[None, None, None, :] - z0[:, None, None, None]
    spread = spread[:, None, None, None]
    bright = bright[:, None, None, None]

    s2 = spread * spread
    a = np.exp(-(dx * dx) / s2)
    b = np.exp(-(dy * dy) / s2)
    c = np.exp(-(dz * dz) / s2)

    out = np.zeros((200, 200, 200), dtype=np.uint16)

    for ii in xrange(bright.shape[0]):
        out += bright[ii] * a[ii] * b[ii] * c[ii]

    return out

这将需要更少的内存,但潜在的缺点是它需要在 Python.

中循环

让您了解相对性能:

def sum_spots_listcomp(bright, spread, x0, y0, z0):
    return np.sum([make_spot_3d(bright[ii], spread[ii], x0[ii], y0[ii], z0[ii])
                   for ii in xrange(len(bright))], axis=0)

def sum_spots_vec(bright, spread, x0, y0, z0):
    return make_spots(bright, spread, x0, y0, z0).sum(0)

# some fake data
bright = np.random.rand(10) * 100
spread = np.random.rand(10) * 100
x0 = (np.random.rand(10) - 0.5) * 50
y0 = (np.random.rand(10) - 0.5) * 50
z0 = (np.random.rand(10) - 0.5) * 50

%timeit sum_spots_listcomp(bright, spread, x0, y0, z0)
# 1 loops, best of 3: 16.6 s per loop

%timeit sum_spots_vec(bright, spread, x0, y0, z0)
# 1 loops, best of 3: 1.03 s per loop

%timeit sum_spots_inplace(bright, spread, x0, y0, z0)
# 1 loops, best of 3: 330 ms per loop

所以,你在任期内每项操作都做了800万(=200*200*200)次;首先,您可以通过仅计算 if 的八分之一并对其进行镜像,将其减少到 100 万(如果球体恰好位于网格的中心)。镜像不是免费的,但仍然比 exp.

便宜得多

此外,很有可能你应该在强度值下降到 0 后才停止计算。使用一点对数魔法,你可以得出一个可能比 200* 小得多的感兴趣区域200*200网格。

既然你有一个i5处理器,而且点之间是相互独立的,那么实现多线程就好了。您不一定需要多个 进程 ,因为许多 Numpy 操作都会释放 GIL。附加代码可以很简单:

from multiprocessing.dummy import Pool

if __name__ == '__main__':
    wrap = lambda pos: make_spot_3d(100, 2, *pos)
    cluster = sum(Pool().imap_unordered(wrap, positions))

更新

在工作时对我的 PC 进行了一些测试后,我必须承认上面的代码太幼稚且效率低下。在 8 核上,相对于单核性能,加速仅为 ~1.5 倍。

我仍然认为多线程是个好主意,但成功很大程度上取决于实现。